Ignite/decoder/y4m/y4m.go

package y4m

import (
	"errors"
	"fmt"
	"git.gammaspectra.live/S.O.N.G/Ignite/color"
	"git.gammaspectra.live/S.O.N.G/Ignite/frame"
	"git.gammaspectra.live/S.O.N.G/Ignite/utilities"
	"github.com/ulikunitz/xz"
	"io"
	"strconv"
	"strings"
)

type Decoder struct {
	r              io.Reader
	frameSeekTable []int64
	parameters     map[Parameter][]string

	properties frame.StreamProperties

	frameSize int

	frameStartOffset int64

	frameCounter int
	timecodes    utilities.Timecodes

	pool *frame.Pool
}

type Parameter byte

const (
	ParameterFrameWidth       Parameter = 'W'
	ParameterFrameHeight      Parameter = 'H'
	ParameterFrameRate        Parameter = 'F'
	ParameterInterlacing      Parameter = 'I'
	ParameterPixelAspectRatio Parameter = 'A'
	ParameterColorFormat      Parameter = 'C'
	ParameterExtension        Parameter = 'X'
)

const fileMagic = "YUV4MPEG2 "
const frameMagic = "FRAME"

func NewXZCompressedDecoder(reader io.Reader, settings map[string]any) (*Decoder, error) {
	r, err := xz.NewReader(reader)
	if err != nil {
		return nil, err
	}
	return NewDecoder(r, settings)
}

func NewDecoder(reader io.Reader, settings map[string]any) (*Decoder, error) {
	s := &Decoder{
		r:          reader,
		parameters: make(map[Parameter][]string),
	}
	var err error

	if err = s.readHeader(); err != nil {
		return nil, err
	}

	if err = s.parseParameters(); err != nil {
		return nil, err
	}

	if t, ok := settings["seek_table"]; ok {
		if table, ok := t.([]int64); ok {
			s.frameSeekTable = table
		}
	}

	if t, ok := settings["timecodes"]; ok {
		if tc, ok := t.(utilities.Timecodes); ok {
			s.timecodes = tc
		}
	}

	// Flagged VFR
	if s.properties.FrameRate.Numerator == 0 && s.properties.FrameRate.Denominator == 0 && len(s.timecodes) == 0 {
		return nil, errors.New("Y4M indicates VFR but timecodes is not set")
	}

	if s.properties.PixelAspectRatio.Numerator == 0 && s.properties.PixelAspectRatio.Denominator == 0 {
		// Assume 1:1
		s.properties.PixelAspectRatio = utilities.NewRatio(1, 1)
	}

	s.properties.VFR = s.IsVFR()

	s.pool, err = frame.NewPool(s.properties.FrameProperties())
	if err != nil {
		return nil, err
	}

	return s, nil
}

func (s *Decoder) IsVFR() bool {
	return len(s.timecodes) != 0
}

// FramePTS Returns -1 if not found
func (s *Decoder) FramePTS(n int) int64 {
	if s.IsVFR() {
		if n >= len(s.timecodes) {
			return -1
		}
		return s.timecodes[n]
	}
	return int64(n)
}

func (s *Decoder) Properties() frame.StreamProperties {
	return s.properties
}

func (s *Decoder) SeekToFrame(frameNumber int) (err error) {
	if s.frameCounter == frameNumber {
		return nil
	}

	if seeker, ok := s.r.(io.Seeker); ok {
		if frameNumber >= 0 && len(s.frameSeekTable) > frameNumber && s.frameSeekTable[frameNumber] != 0 {
			if _, err = seeker.Seek(s.frameSeekTable[frameNumber], io.SeekStart); err != nil {
				return err
			}

			s.frameCounter = frameNumber

			return nil
		} else if frameNumber >= 0 && len(s.frameSeekTable) > 0 {
			//attempt blind seek from last decoded frame
			framesToSeekFromLast := frameNumber + 1 - len(s.frameSeekTable)

			if _, err = seeker.Seek(s.frameSeekTable[len(s.frameSeekTable)-1]+int64(5+1+s.frameSize)*int64(framesToSeekFromLast), io.SeekStart); err != nil {
				return err
			}

			s.frameCounter = frameNumber

			return nil
		} else if frameNumber >= 0 && s.frameStartOffset != 0 {
			//attempt full blind seek from start
			if _, err = seeker.Seek(s.frameStartOffset+int64(5+1+s.frameSize)*int64(frameNumber), io.SeekStart); err != nil {
				return err
			}

			s.frameCounter = frameNumber

			return nil

		} else {
			return errors.New("frameNumber out of range")
		}
	} else {
		return errors.New("reader is not io.Seeker")
	}
}

func (s *Decoder) Decode() (frame.Frame, error) {
	_, f, err := s.GetFrame()
	return f, err
}
func (s *Decoder) DecodeStream() *frame.Stream {
	stream, channel := frame.NewStream(s.properties)
	go func() {
		defer close(channel)
		for {
			if f, err := s.Decode(); err != nil {
				return
			} else {
				channel <- f
			}
		}
	}()
	return stream
}

func (s *Decoder) GetFrameSeekTable() []int64 {
	return s.frameSeekTable
}

func (s *Decoder) GetFrame() (parameters map[Parameter][]string, frameObject frame.Frame, err error) {
	var index int64

	if seeker, ok := s.r.(io.Seeker); ok {
		if index, err = seeker.Seek(0, io.SeekCurrent); err != nil {
			return nil, nil, err
		}
	}

	if parameters, err = s.readFrameHeader(); err != nil {
		return nil, nil, err
	}

	if index > 0 {
		for len(s.frameSeekTable) <= s.frameCounter {
			s.frameSeekTable = append(s.frameSeekTable, 0)
		}
		s.frameSeekTable[s.frameCounter] = index
	}

	pts := s.FramePTS(s.frameCounter)
	if pts == -1 {
		return nil, nil, fmt.Errorf("frame %d PTS could not be calculated", s.frameCounter)
	}

	if frameObject, err = s.readFrameData(pts); err != nil {
		return nil, nil, err
	}

	s.frameCounter++

	return
}

func (s *Decoder) readHeader() (err error) {
	var header [10]byte
	if _, err = io.ReadFull(s.r, header[:]); err != nil {
		return err
	}

	if string(header[:]) != fileMagic {
		return fmt.Errorf("invalid file signature, %s != %s", string(header[:]), fileMagic)
	}

	data := make([]byte, 0, 1024)
	buf := make([]byte, 1)
	for {
		if _, err = s.r.Read(buf); err != nil {
			return err
		}

		if buf[0] == '\n' {
			break
		}

		data = append(data, buf[0])
	}

	for _, v := range strings.Split(string(data), " ") {
		if len(v) > 1 {
			if slice, ok := s.parameters[Parameter(v[0])]; ok {
				s.parameters[Parameter(v[0])] = append(slice, v[1:])
			} else {
				s.parameters[Parameter(v[0])] = []string{v[1:]}
			}
		}
	}

	return nil
}

func (s *Decoder) readFrameHeader() (parameters map[Parameter][]string, err error) {
	var header [5]byte
	if _, err = io.ReadFull(s.r, header[:]); err != nil {
		return nil, err
	}

	if string(header[:]) != frameMagic {
		return nil, fmt.Errorf("invalid frame signature, %s != %s", string(header[:]), frameMagic)
	}

	var data []byte
	var buf [1]byte
	for {
		if _, err = s.r.Read(buf[:]); err != nil {
			return nil, err
		}

		if buf[0] == '\n' {
			break
		}

		data = append(data, buf[0])
	}

	for _, v := range strings.Split(string(data), " ") {
		if len(v) > 1 {
			if parameters == nil {
				parameters = make(map[Parameter][]string)
			}
			if slice, ok := parameters[Parameter(v[0])]; ok {
				parameters[Parameter(v[0])] = append(slice, v[1:])
			} else {
				parameters[Parameter(v[0])] = []string{v[1:]}
			}
		}
	}

	return parameters, nil
}

func (s *Decoder) readFrameData(pts int64) (f frame.Frame, err error) {
	f = s.pool.Get(pts)
	_, err = io.ReadFull(s.r, f.GetJoint())
	return f, err
}

func (s *Decoder) parseParameters() (err error) {
	for k, values := range s.parameters {
		switch k {
		case ParameterFrameWidth:
			if s.properties.Width, err = strconv.Atoi(values[0]); err != nil {
				return err
			}
		case ParameterFrameHeight:
			if s.properties.Height, err = strconv.Atoi(values[0]); err != nil {
				return err
			}
		case ParameterFrameRate:
			v := strings.Split(values[0], ":")
			if len(v) != 2 {
				return fmt.Errorf("wrong frame rate %s", values[0])
			}
			if s.properties.FrameRate.Numerator, err = strconv.Atoi(v[0]); err != nil {
				return err
			}
			if s.properties.FrameRate.Denominator, err = strconv.Atoi(v[1]); err != nil {
				return err
			}
		case ParameterInterlacing:
			if values[0] != "p" {
				return fmt.Errorf("not supported interlacing %s", values[0])
			}
		case ParameterPixelAspectRatio:
			v := strings.Split(values[0], ":")
			if len(v) != 2 {
				return fmt.Errorf("wrong pixel aspect ratio %s", values[0])
			}
			if s.properties.PixelAspectRatio.Numerator, err = strconv.Atoi(v[0]); err != nil {
				return err
			}
			if s.properties.PixelAspectRatio.Denominator, err = strconv.Atoi(v[1]); err != nil {
				return err
			}
		case ParameterColorFormat:
			if s.properties.ColorSpace, err = color.NewColorFormatFromString(values[0]); err != nil {
				return err
			}
		case ParameterExtension:
			for _, v := range values {
				extVal := strings.Split(v, "=")
				if len(extVal) >= 2 {
					switch extVal[0] {
					case "COLORRANGE":
						if extVal[1] == "FULL" {
							s.properties.FullColorRange = true
						} else if extVal[1] == "LIMITED" {
							s.properties.FullColorRange = false
						} else {
							return fmt.Errorf("not supported %s: %s", extVal[0], extVal[1])
						}
					}
				}
			}
		}
	}

	//TODO: check for missing values width, height, colorformat etc.

	s.frameSize, err = s.properties.ColorSpace.FrameSize(s.properties.Width, s.properties.Height)

	if seeker, ok := s.r.(io.Seeker); ok {
		if index, err := seeker.Seek(0, io.SeekCurrent); err == nil {
			s.frameStartOffset = index
		}
	}

	return err
}

func (s *Decoder) Close() {

}

func (s *Decoder) Version() string {
	return "y4m"
}