Ignite/y4m/y4m.go

package y4m

import (
	"errors"
	"fmt"
	"git.gammaspectra.live/S.O.N.G/Ignite/colorspace"
	"git.gammaspectra.live/S.O.N.G/Ignite/frame"
	"git.gammaspectra.live/S.O.N.G/Ignite/utilities"
	"io"
	"strconv"
	"strings"
)

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

	width            int
	height           int
	frameRate        utilities.Ratio
	pixelAspectRatio utilities.Ratio
	colorSpace       colorspace.ColorSpace
	colorRangeFull   bool

	frameSize int

	frameCounter int
}

type Parameter byte

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

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

func New(reader io.Reader) (*Stream, error) {
	s := &Stream{
		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
	}

	return s, nil
}

func (s *Stream) Resolution() (width, height int) {
	return s.width, s.height
}

func (s *Stream) FrameRate() utilities.Ratio {
	return s.frameRate
}

func (s *Stream) ColorSpace() colorspace.ColorSpace {
	return s.colorSpace
}

func (s *Stream) ColorRange() bool {
	return s.colorRangeFull
}

func (s *Stream) PixelAspectRatio() utilities.Ratio {
	return s.pixelAspectRatio
}

func (s *Stream) SeekToFrame(frameNumber int) (err error) {
	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 {
			return errors.New("frameNumber out of range")
		}
	} else {
		return errors.New("reader is not io.Seeker")
	}
}

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

	frameNumber = s.frameCounter

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

	if parameters, err = s.readFrameHeader(); err != nil {
		return 0, 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 {
		return 0, nil, nil, err
	}

	if s.colorSpace.BitDepth > 8 {
		if frameObject, err = frame.NewUint16FrameFromBytes(s.colorSpace, s.width, s.height, buf); err != nil {
			return 0, nil, nil, err
		}
	} else {
		if frameObject, err = frame.NewUint8FrameFromBytes(s.colorSpace, s.width, s.height, buf); err != nil {
			return 0, nil, nil, err
		}
	}

	s.frameCounter++

	return
}

func (s *Stream) readHeader() (err error) {
	header := make([]byte, 10)
	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 *Stream) readFrameHeader() (parameters map[Parameter][]string, err error) {
	header := make([]byte, 5)
	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
	buf := make([]byte, 1)
	for {
		if _, err = s.r.Read(buf); err != nil {
			return nil, err
		}

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

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

	parameters = make(map[Parameter][]string)

	for _, v := range strings.Split(string(data), " ") {
		if len(v) > 1 {
			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 *Stream) readFrameData() (buf []byte, err error) {
	//TODO: reuse buffers, maybe channel?
	buf = make([]byte, s.frameSize)
	_, err = io.ReadFull(s.r, buf)
	return buf, err
}

func (s *Stream) parseParameters() (err error) {
	for k, values := range s.parameters {
		switch k {
		case ParameterFrameWidth:
			if s.width, err = strconv.Atoi(values[0]); err != nil {
				return err
			}
		case ParameterFrameHeight:
			if s.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.frameRate.Numerator, err = strconv.Atoi(v[0]); err != nil {
				return err
			}
			if s.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.pixelAspectRatio.Numerator, err = strconv.Atoi(v[0]); err != nil {
				return err
			}
			if s.pixelAspectRatio.Denominator, err = strconv.Atoi(v[1]); err != nil {
				return err
			}
		case ParameterColorSpace:
			if s.colorSpace, err = colorspace.NewColorSpaceFromString(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.colorRangeFull = true
						} else if extVal[1] == "LIMITED" {
							s.colorRangeFull = false
						} else {
							return fmt.Errorf("not supported %s: %s", extVal[0], extVal[1])
						}
					}
				}
			}
		}
	}

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

	s.frameSize, err = s.colorSpace.FrameSize(s.width, s.height)

	return err
}