Alawn/chatlog
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

parse multi-partition wxgf format (#195)

Browse Source
This commit is contained in:
Sarv
2025-08-18 17:25:18 +08:00
committed by GitHub
parent 3ec6a9a81b
commit a16b689f4e
1 changed files with 301 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/util/dat2img/wxgf.go
+301 -15
View File
@@ -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
}