Ignite/frame/stream.go

package frame

import (
	"git.gammaspectra.live/S.O.N.G/Ignite/color"
	"git.gammaspectra.live/S.O.N.G/Ignite/utilities"
	"sync/atomic"
)

type Stream struct {
	properties StreamProperties
	channel    <-chan Frame
	lock       atomic.Bool
}

func NewStream(properties StreamProperties) (*Stream, chan<- Frame) {
	c := make(chan Frame)
	s := &Stream{
		properties: properties,
		channel:    c,
	}
	return s, c
}

type StreamProperties struct {
	// Width could be not populated until the first frame is read. Frame can contain different settings.
	Width int `json:"width" yaml:"width"`
	// Height could be not populated until the first frame is read. Frame can contain different settings.
	Height int `json:"height" yaml:"height"`
	// PixelAspectRatio could be not populated until the first frame is read. Frame can contain different settings.
	PixelAspectRatio utilities.Ratio `json:"par" yaml:"par"`
	// ColorSpace could be not populated until the first frame is read. Frame can contain different settings.
	ColorSpace     color.Space     `json:"colorspace" yaml:"colorspace"`
	FrameRate      utilities.Ratio `json:"framerate" yaml:"framerate"`
	FullColorRange bool            `json:"fullrange" yaml:"fullrange"`
	VFR            bool            `json:"vfr" yaml:"vfr"`
}

func (p StreamProperties) TimeBase() utilities.Ratio {
	timeBase := p.FrameRate.Reciprocal()
	if p.VFR {
		timeBase = utilities.NewRatio(1, 1000)
	}

	return timeBase
}

func (p StreamProperties) FrameProperties() Properties {
	return Properties{
		Width:            p.Width,
		Height:           p.Height,
		PixelAspectRatio: p.PixelAspectRatio,
		ColorSpace:       p.ColorSpace,
		FullColorRange:   p.FullColorRange,
	}
}

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

// Channel gets the Frame channel, and locks the input
func (s *Stream) Channel() <-chan Frame {
	if s.Lock() {
		return s.channel
	}

	return nil
}

// Lock locks the input, and returns success
func (s *Stream) Lock() bool {
	return !s.lock.Swap(true)
}

// Unlock unlocks the input, and returns the old value
func (s *Stream) Unlock() bool {
	return s.lock.Swap(false)
}

// Copy copies the stream n times
func (s *Stream) Copy(n int) []*Stream {
	if !s.Lock() {
		return nil
	}
	slice := make([]*Stream, n)
	sliceChannels := make([]chan<- Frame, n)
	for i := range slice {
		slice[i], sliceChannels[i] = NewStream(s.Properties())
	}

	go func() {
		defer func() {
			for _, sliceItem := range sliceChannels {
				close(sliceItem)
			}
		}()

		for f := range s.channel {
			for _, sliceItem := range sliceChannels {
				sliceItem <- f
			}
		}
	}()

	return slice
}

// Slice produces a slice of the stream, and locks the input
func (s *Stream) Slice(start, end int) *Stream {
	if end < start {
		return nil
	}
	if !s.Lock() {
		return nil
	}
	slice, channel := NewStream(s.Properties())

	var index int
	go func() {
		defer func() {
			//empty source channel
			for range s.channel {

			}
		}()
		defer close(channel)
		for f := range s.channel {
			if index >= end {
				break
			}

			if index >= start {
				channel <- f
			}
			index++
		}
	}()

	return slice
}

// Split produces several Stream across each split point, and locks the input.
// For example, splits = [2, 5, 20] will produce 4 Stream, [0-1], [2-4], [5-19], [20-...]
func (s *Stream) Split(splits ...int) []*Stream {
	if len(splits) == 0 {
		return []*Stream{s}
	} else if !s.Lock() {
		return nil
	}
	slice := make([]*Stream, len(splits)+1)
	sliceChannels := make([]chan<- Frame, len(splits)+1)
	for i := range slice {
		slice[i], sliceChannels[i] = NewStream(s.Properties())
	}

	var index int
	go func() {
		defer func() {
			//empty source channel
			for range s.channel {

			}
		}()
		defer func() {
			for _, sC := range sliceChannels {
				close(sC)
			}
		}()

		for f := range s.channel {
			for len(splits) > 0 && index >= splits[0] {
				close(sliceChannels[0])
				splits = splits[1:]
				sliceChannels = sliceChannels[1:]
				slice = slice[1:]
			}

			sliceChannels[0] <- f

			index++
		}
	}()

	return slice
}

// Sample samples frames every each Frame, and locks the input
func (s *Stream) Sample(each int) *Stream {
	if !s.Lock() {
		return nil
	}
	slice, channel := NewStream(s.Properties())

	var index int
	go func() {
		defer close(channel)
		for f := range s.channel {
			if index%each == 0 {
				channel <- f
			}
			index++
		}
	}()

	return slice
}

// Monochrome makes stream luma-only
func (s *Stream) Monochrome() *Stream {
	if !s.Lock() {
		return nil
	}

	props := s.Properties()
	props.ColorSpace.ChromaSampling.A = 0
	props.ColorSpace.ChromaSampling.B = 0

	slice, channel := NewStream(props)

	go func() {
		defer close(channel)
		for f := range s.channel {
			frameProps := f.Properties()
			frameProps.ColorSpace.ChromaSampling.A = 0
			frameProps.ColorSpace.ChromaSampling.B = 0

			switch typedFrame := f.(type) {
			case TypedFrame[uint8]:
				channel <- &fUint8{
					properties: frameProps,
					Pts:        typedFrame.PTS(),
					NextPts:    typedFrame.NextPTS(),
					Y:          typedFrame.GetNativeLuma(),
					Cb:         nil,
					Cr:         nil,
				}
			case TypedFrame[uint16]:
				channel <- &fUint16{
					properties: frameProps,
					Pts:        typedFrame.PTS(),
					NextPts:    typedFrame.NextPTS(),
					Y:          typedFrame.GetNativeLuma(),
					Cb:         nil,
					Cr:         nil,
				}
			default:
				panic("unknown frame type")
			}
		}
	}()

	return slice
}