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"
	"io"
	"runtime"
	"strconv"
	"strings"
	"sync"
)

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

	properties frame.StreamProperties

	frameSize int

	frameStartOffset int64

	frameCounter int

	bufPool sync.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 NewDecoder(reader io.Reader, settings map[string]any) (*Decoder, error) {
	s := &Decoder{
		r:          reader,
		parameters: make(map[Parameter][]string),
	}

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

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

	s.bufPool.New = func() any {
		return make([]byte, s.frameSize)
	}

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

	return s, nil
}

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
	}

	var buf []byte
	if buf, err = s.readFrameData(); err != nil {
		s.bufPool.Put(buf)
		return nil, nil, err
	}

	if s.properties.ColorSpace.BitDepth > 8 {
		//it's copied below
		defer s.bufPool.Put(buf)
		if frameObject, err = frame.NewUint16FrameFromBytes(s.properties.FrameProperties(), int64(s.frameCounter), buf); err != nil {
			return nil, nil, err
		}
	} else {
		if f8, err := frame.NewUint8FrameFromBytes(s.properties.FrameProperties(), int64(s.frameCounter), buf); err != nil {
			s.bufPool.Put(buf)
			return nil, nil, err
		} else {
			frameObject = f8
			runtime.SetFinalizer(f8, func(f *frame.FrameUint8) {
				//return buffer to pool once top frame is not in use
				s.bufPool.Put(buf)
			})
		}
	}

	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() (buf []byte, err error) {
	//TODO: reuse buffers, maybe channel?
	buf = s.bufPool.Get().([]byte)
	_, err = io.ReadFull(s.r, buf)
	return buf, 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"
}