diff --git a/pkg/util/dat2img/wxgf.go b/pkg/util/dat2img/wxgf.go index 694f9b7..8d98d9a 100644 --- a/pkg/util/dat2img/wxgf.go +++ b/pkg/util/dat2img/wxgf.go @@ -5,16 +5,19 @@ import ( "fmt" "os" "os/exec" + "path/filepath" "github.com/Eyevinn/mp4ff/avc" "github.com/Eyevinn/mp4ff/bits" "github.com/Eyevinn/mp4ff/hevc" "github.com/Eyevinn/mp4ff/mp4" + "github.com/google/uuid" ) const ( ENV_FFMPEG_PATH = "FFMPEG_PATH" MinRatio = 0.6 + FixSliceHeaders = true ) var ( @@ -39,11 +42,39 @@ func Wxam2pic(data []byte) ([]byte, string, error) { return nil, "", fmt.Errorf("invalid wxgf") } - offset, size, err := findDataPartition(data) + partitions, err := findDataPartition(data) if err != nil { return nil, "", err } + if partitions.LikeAnime() { + // FIXME mask frame not work + animeFrames := make([][]byte, 0) + maskFrames := make([][]byte, 0) + for i, partition := range partitions.Partitions { + if i%2 == 0 { + maskFrames = append(maskFrames, data[partition.Offset:partition.Offset+partition.Size]) + } else { + animeFrames = append(animeFrames, data[partition.Offset:partition.Offset+partition.Size]) + } + } + if FFmpegMode { + mp4Data, err := ConvertAnime2GIF(animeFrames, maskFrames) + if err != nil { + return nil, "", err + } + return mp4Data, "gif", nil + } + mp4Data, err := TransmuxAnime2MP4(animeFrames, maskFrames) + if err != nil { + return nil, "", err + } + return mp4Data, "mp4", nil + } + + offset := partitions.Partitions[partitions.MaxIndex].Offset + size := partitions.Partitions[partitions.MaxIndex].Size + if FFmpegMode { jpgData, err := Convert2JPG(data[offset : offset+size]) if err != nil { @@ -52,18 +83,34 @@ func Wxam2pic(data []byte) ([]byte, string, error) { return jpgData, JPG.Ext, nil } - mp4Data, err := Convert2MP4(data[offset : offset+size]) + mp4Data, err := Transmux2MP4(data[offset : offset+size]) if err != nil { return nil, "", err } return mp4Data, "mp4", nil } -func findDataPartition(data []byte) (offset int, size int, err error) { +type Partitions struct { + Partitions []Partition + MaxRatio float64 + MaxIndex int +} + +func (p *Partitions) LikeAnime() bool { + return len(p.Partitions) > 1 && p.MaxRatio < MinRatio +} + +type Partition struct { + Offset int + Size int + Ratio float64 +} + +func findDataPartition(data []byte) (*Partitions, error) { headerLen := int(data[4]) if headerLen >= len(data) { - return 0, 0, fmt.Errorf("invalid wxgf") + return nil, fmt.Errorf("invalid wxgf") } patterns := [][]byte{ @@ -72,17 +119,24 @@ func findDataPartition(data []byte) (offset int, size int, err error) { } for _, pattern := range patterns { + ret := &Partitions{ + Partitions: make([]Partition, 0), + } offset := 0 for { + if headerLen+offset > len(data) { + break + } + index := bytes.Index(data[headerLen+offset:], pattern) if index == -1 { break } absIndex := headerLen + offset + index - offset += index + 1 if absIndex < 4 { + offset += index + 1 continue } @@ -90,20 +144,29 @@ func findDataPartition(data []byte) (offset int, size int, err error) { int(data[absIndex-2])<<8 | int(data[absIndex-1]) if length <= 0 || absIndex+length > len(data) { + offset += index + 1 continue } - ratio := float64(length) / float64(len(data)) - if ratio < MinRatio { - continue + partition := Partition{ + Offset: absIndex, + Size: length, + Ratio: float64(length) / float64(len(data)), } - - return absIndex, length, nil + ret.Partitions = append(ret.Partitions, partition) + if partition.Ratio > ret.MaxRatio { + ret.MaxRatio = partition.Ratio + ret.MaxIndex = len(ret.Partitions) - 1 + } + offset += index + length } + if len(ret.Partitions) > 0 { + return ret, nil + } } - return 0, 0, fmt.Errorf("no partition found") + return nil, fmt.Errorf("no partition found") } func Convert2JPG(data []byte) ([]byte, error) { @@ -132,7 +195,66 @@ func Convert2JPG(data []byte) ([]byte, error) { return jpegData, nil } -func Convert2MP4(data []byte) ([]byte, error) { +func writeTempFile(data [][]byte) (string, error) { + path := filepath.Join(os.TempDir(), fmt.Sprintf("anime-%s", uuid.New().String())) + file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + return "", fmt.Errorf("failed to open anime temp file: %w", err) + } + defer file.Close() + for _, frame := range data { + _, err := file.Write(frame) + if err != nil { + return "", fmt.Errorf("failed to write anime frame to temp file: %w", err) + } + } + return path, nil +} + +// ConvertAnime2GIF convert anime frames and mask frames to mp4 +// FIXME No longer need to write to temporary files +func ConvertAnime2GIF(animeFrames [][]byte, maskFrames [][]byte) ([]byte, error) { + animeFilePath, err := writeTempFile(animeFrames) + if err != nil { + return nil, fmt.Errorf("failed to write anime temp file: %w", err) + } + defer os.Remove(animeFilePath) + + maskFilePath, err := writeTempFile(maskFrames) + if err != nil { + return nil, fmt.Errorf("failed to write mask temp file: %w", err) + } + defer os.Remove(maskFilePath) + + cmd := exec.Command(FFMpegPath, + "-i", animeFilePath, + "-i", maskFilePath, + "-filter_complex", "[0:v][1:v]alphamerge,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse", + "-f", "gif", + "-") + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + return nil, fmt.Errorf("ffmpeg failed: %w", err) + } + + gifData := stdout.Bytes() + if len(gifData) == 0 { + return nil, fmt.Errorf("ffmpeg output is empty") + } + + return gifData, nil +} + +func isFFmpegAvailable() bool { + cmd := exec.Command(FFMpegPath, "-version") + return cmd.Run() == nil +} + +func Transmux2MP4(data []byte) ([]byte, error) { vpsNALUs, spsNALUs, ppsNALUs := hevc.GetParameterSetsFromByteStream(data) @@ -177,7 +299,171 @@ func Convert2MP4(data []byte) ([]byte, error) { return sw.Bytes(), nil } -func isFFmpegAvailable() bool { - cmd := exec.Command(FFMpegPath, "-version") - return cmd.Run() == nil +func TransmuxAnime2MP4(animeFrames [][]byte, maskFrames [][]byte) ([]byte, error) { + + if len(maskFrames) != len(animeFrames) { + return nil, fmt.Errorf("mask frame num (%d) not equal to anime frame num (%d)", len(maskFrames), len(animeFrames)) + } + + init := mp4.CreateEmptyInit() + seg := mp4.NewMediaSegment() + trackIDs := []uint32{1, 2} + frag, err := mp4.CreateMultiTrackFragment(1, trackIDs) + if err != nil { + return nil, err + } + seg.AddFragment(frag) + + err = Add2Trak(init, frag, 0, animeFrames) + if err != nil { + return nil, fmt.Errorf("add full sample to track failed: %w", err) + } + + err = Add2Trak(init, frag, 1, maskFrames) + if err != nil { + return nil, fmt.Errorf("add full sample to track failed: %w", err) + } + + totalSize := init.Size() + seg.Size() + sw := bits.NewFixedSliceWriter(int(totalSize)) + + init.EncodeSW(sw) + seg.EncodeSW(sw) + + return sw.Bytes(), nil +} + +func Add2Trak(init *mp4.InitSegment, frag *mp4.Fragment, index int, data [][]byte) error { + videoTimescale := uint32(90000) + init.AddEmptyTrack(videoTimescale, "video", "und") + trak := init.Moov.Traks[index] + + vps, sps, pps := hevc.GetParameterSetsFromByteStream(data[0]) + + // FIXME Two slices reporting being the first in the same frame. + if FixSliceHeaders { + spsMap, ppsMap, err := createSPSPPSMaps(vps, sps, pps) + if err != nil { + return fmt.Errorf("create sps pps map failed: %w", err) + } + for i := range data { + fixedFrame, err := fixSliceHeadersInFrame(data[i], spsMap, ppsMap) + if err != nil { + return fmt.Errorf("fix slice header failed: %w", err) + } + data[i] = fixedFrame + } + } + + err := trak.SetHEVCDescriptor("hev1", vps, sps, pps, nil, true) + if err != nil { + return fmt.Errorf("set trak failed: %w", err) + } + + var decodeTime uint64 = 0 + frameDuration := uint32(3000) + for i := 0; i < len(data); i++ { + sampleData := avc.ConvertByteStreamToNaluSample(removeParameterSets(data[i])) + simple := mp4.FullSample{ + Sample: mp4.Sample{ + Flags: getSampleFlags(sampleData, i == 0), + Dur: frameDuration, + Size: uint32(len(sampleData)), + }, + DecodeTime: decodeTime, + Data: sampleData, + } + + err = frag.AddFullSampleToTrack(simple, uint32(index+1)) + if err != nil { + return err + } + decodeTime += uint64(frameDuration) + } + + return nil +} + +func getSampleFlags(frameData []byte, isFirstFrame bool) uint32 { + if isFirstFrame || hevc.IsRAPSample(frameData) { + return 0x02000000 + } + return 0x01010000 +} + +func createSPSPPSMaps(vpsNalus, spsNalus, ppsNalus [][]byte) (map[uint32]*hevc.SPS, map[uint32]*hevc.PPS, error) { + spsMap := make(map[uint32]*hevc.SPS) + for _, spsNalu := range spsNalus { + sps, err := hevc.ParseSPSNALUnit(spsNalu) + if err != nil { + return nil, nil, fmt.Errorf("parse sps failed: %w", err) + } + spsMap[uint32(sps.SpsID)] = sps + } + + ppsMap := make(map[uint32]*hevc.PPS) + for _, ppsNalu := range ppsNalus { + pps, err := hevc.ParsePPSNALUnit(ppsNalu, spsMap) + if err != nil { + return nil, nil, fmt.Errorf("parse pps failed: %w", err) + } + ppsMap[pps.PicParameterSetID] = pps + } + + return spsMap, ppsMap, nil +} + +func fixSliceHeadersInFrame(frameData []byte, spsMap map[uint32]*hevc.SPS, ppsMap map[uint32]*hevc.PPS) ([]byte, error) { + nalus := avc.ExtractNalusFromByteStream(frameData) + var fixedNalus [][]byte + var firstSliceFound bool + + for _, nalu := range nalus { + naluType := hevc.GetNaluType(nalu[0]) + + if naluType < hevc.NALU_TRAIL_N || naluType > hevc.NALU_CRA { + fixedNalus = append(fixedNalus, nalu) + continue + } + + if !firstSliceFound { + fixedNalus = append(fixedNalus, nalu) + firstSliceFound = true + } else { + sliceHeader, err := hevc.ParseSliceHeader(nalu, spsMap, ppsMap) + if err != nil { + return nil, fmt.Errorf("parse slice header failed: %w", err) + } + if !sliceHeader.FirstSliceSegmentInPicFlag { + fixedNalus = append(fixedNalus, nalu) + } + } + } + + return reconstructAnnexBStream(fixedNalus), nil +} + +func removeParameterSets(annexBData []byte) []byte { + nalus := avc.ExtractNalusFromByteStream(annexBData) + var videoNalus [][]byte + + for _, nalu := range nalus { + naluType := hevc.GetNaluType(nalu[0]) + if naluType <= 31 { + videoNalus = append(videoNalus, nalu) + } + } + + return reconstructAnnexBStream(videoNalus) +} + +func reconstructAnnexBStream(nalus [][]byte) []byte { + var result []byte + startCode := []byte{0x00, 0x00, 0x01} + + for _, nalu := range nalus { + result = append(result, startCode...) + result = append(result, nalu...) + } + return result }