Initial commit: utilities, y4m parser, x264 encoder

This commit is contained in:
DataHoarder 2022-09-14 21:42:57 +02:00
commit 84908c0ef2
Signed by: DataHoarder
SSH key fingerprint: SHA256:EnPQOqPpbCa7nzalCEJY2sd9iPluFIBuJu2rDFalACI
11 changed files with 942 additions and 0 deletions

1
.gitignore vendored Normal file
View file

@ -0,0 +1 @@
.idea/

9
LICENSE Normal file
View file

@ -0,0 +1,9 @@
Copyright (c) 2022 Ignite Contributors All rights reserved.
Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

177
colorspace/colorspace.go Normal file
View file

@ -0,0 +1,177 @@
package colorspace
import (
"encoding/binary"
"errors"
"fmt"
"strconv"
"strings"
)
// SubsamplingScheme The subsampling scheme is commonly expressed as a three-part ratio J : A : B (e.g. 4:2:2), that describe the number of luminance and chrominance samples in a conceptual region that is J pixels wide and 2 pixels high.
type SubsamplingScheme struct {
// J horizontal sampling reference (width of the conceptual region). Usually, 4.
J byte
// A number of chrominance samples (Cb, Cr) in the first row of J pixels.
A byte
// B number of changes of chrominance samples (Cb, Cr) between first and second row of J pixels. Note that B has to be either zero or equal to A
B byte
}
type ColorSpace struct {
Subsampling SubsamplingScheme
BitDepth byte
}
func (c ColorSpace) ToInteger() uint32 {
return binary.BigEndian.Uint32([]byte{c.Subsampling.J, c.Subsampling.A, c.Subsampling.B, c.BitDepth})
}
// ElementPixels returns the number of pixels for a full element
func (c ColorSpace) ElementPixels() int {
return int(c.Subsampling.J) * 2
}
// ElementSamples returns the number of actual samples (total Y, Cb, Cr) for an encoded element
func (c ColorSpace) ElementSamples() int {
return int(c.Subsampling.J)*2 + c.ElementChromaSamples()
}
func (c ColorSpace) ElementChromaSamples() int {
return int(c.Subsampling.A)*2 + int(c.Subsampling.B)*2
}
func (c ColorSpace) PlaneLumaSamples(width, height int) int {
return width * height
}
func (c ColorSpace) PlaneCrSamples(width, height int) int {
return int((int64(width) * int64(height) * (int64(c.Subsampling.A) + int64(c.Subsampling.B))) / int64(c.ElementPixels()))
}
func (c ColorSpace) PlaneCbSamples(width, height int) int {
return int((int64(width) * int64(height) * (int64(c.Subsampling.A) + int64(c.Subsampling.B))) / int64(c.ElementPixels()))
}
func (c ColorSpace) FrameSizePacked(width, height int) (int, error) {
a1 := int64(width) * int64(height)
a2 := int64(c.ElementPixels())
if a1%a2 != 0 {
return 0, errors.New("not divisible pixels")
}
a3 := (a1 / a2) * int64(c.ElementSamples()) * int64(c.BitDepth)
if a3%8 != 0 {
return 0, errors.New("not divisible size")
}
return int(a3 / 8), nil
}
func (c ColorSpace) FrameSize(width, height int) (int, error) {
actualBitDepth := int64(c.BitDepth)
if actualBitDepth&0b111 != 0 {
actualBitDepth = ((actualBitDepth + 8) >> 3) << 3
}
a1 := int64(width) * int64(height)
a2 := int64(c.ElementPixels())
if a1%a2 != 0 {
return 0, errors.New("not divisible pixels")
}
a3 := (a1 / a2) * int64(c.ElementSamples()) * actualBitDepth
if a3%8 != 0 {
return 0, errors.New("not divisible size")
}
return int(a3 / 8), nil
}
func (c ColorSpace) check() error {
if c.Subsampling.J <= 0 {
return fmt.Errorf("unsupported J %d", c.Subsampling.J)
}
if c.Subsampling.A < 0 && c.Subsampling.A > c.Subsampling.J {
return fmt.Errorf("unsupported A %d", c.Subsampling.A)
}
if c.Subsampling.B != 0 && c.Subsampling.B != c.Subsampling.A {
return fmt.Errorf("unsupported B %d", c.Subsampling.B)
}
if c.BitDepth != 8 && c.BitDepth != 10 && c.BitDepth != 12 && c.BitDepth != 14 && c.BitDepth != 16 {
return fmt.Errorf("unsupported BitDepth %d", c.BitDepth)
}
return nil
}
func NewColorSpaceFromInteger(colorSpace uint32) (ColorSpace, error) {
buf := make([]byte, 4)
binary.BigEndian.PutUint32(buf, colorSpace)
space := ColorSpace{
Subsampling: SubsamplingScheme{
J: buf[0],
A: buf[1],
B: buf[2],
},
BitDepth: buf[3],
}
return space, space.check()
}
func NewColorSpaceFromString(colorSpace string) (ColorSpace, error) {
colorSpace = strings.ToLower(colorSpace)
if colorSpace == "420mpeg2" {
return NewColorSpaceFromString("420p8")
}
splits := strings.Split(colorSpace, "p")
if len(splits) == 1 { //default 8 bit
splits = append(splits, "8")
}
space := ColorSpace{}
switch strings.ReplaceAll(splits[0], ":", "") {
case "420":
space.Subsampling.J = 4
space.Subsampling.A = 2
space.Subsampling.B = 0
case "422":
space.Subsampling.J = 4
space.Subsampling.A = 2
space.Subsampling.B = 2
case "444":
space.Subsampling.J = 4
space.Subsampling.A = 4
space.Subsampling.B = 4
default:
return space, fmt.Errorf("unsupported Chroma Subsampling %s", splits[0])
}
n, err := strconv.Atoi(splits[1])
if err != nil {
return space, err
}
space.BitDepth = byte(n)
return space, space.check()
}
func newColorSpaceFromStringInternal(colorSpace string) ColorSpace {
space, _ := NewColorSpaceFromString(colorSpace)
return space
}
var (
Space420 ColorSpace = newColorSpaceFromStringInternal("420")
Space422 ColorSpace = newColorSpaceFromStringInternal("422")
Space444 ColorSpace = newColorSpaceFromStringInternal("444")
Space420P10 ColorSpace = newColorSpaceFromStringInternal("420p10")
Space422P10 ColorSpace = newColorSpaceFromStringInternal("422p10")
Space444P10 ColorSpace = newColorSpaceFromStringInternal("444p10")
Space420P12 ColorSpace = newColorSpaceFromStringInternal("420p12")
Space422P12 ColorSpace = newColorSpaceFromStringInternal("422p12")
Space444P12 ColorSpace = newColorSpaceFromStringInternal("444p12")
Space420P16 ColorSpace = newColorSpaceFromStringInternal("420p16")
Space422P16 ColorSpace = newColorSpaceFromStringInternal("422p16")
Space444P16 ColorSpace = newColorSpaceFromStringInternal("444p16")
)

244
encoder/x264/x264.go Normal file
View file

@ -0,0 +1,244 @@
package x264
/*
#cgo pkg-config: x264
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <x264.h>
*/
import "C"
import (
"errors"
"fmt"
"git.gammaspectra.live/S.O.N.G/Ignite/frame"
"git.gammaspectra.live/S.O.N.G/Ignite/y4m"
"io"
"runtime"
"strconv"
"sync/atomic"
"unsafe"
)
type Encoder struct {
w io.Writer
cleaned atomic.Bool
params C.x264_param_t
pictureIn *C.x264_picture_t
pictureOut *C.x264_picture_t
h *C.x264_t
}
func NewEncoder(w io.Writer, stream *y4m.Stream, settings map[string]any) (*Encoder, error) {
e := &Encoder{
w: w,
}
preset := C.CString(getSettingString(settings, "preset", "medium"))
defer C.free(unsafe.Pointer(preset))
if C.x264_param_default_preset(&e.params, preset, nil) < 0 {
return nil, errors.New("error setting preset")
}
defaultProfile := "high"
switch true {
case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 4 && stream.ColorSpace().Subsampling.B == 4:
e.params.i_csp = C.X264_CSP_I444
defaultProfile = "high444"
case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 2 && stream.ColorSpace().Subsampling.B == 2:
e.params.i_csp = C.X264_CSP_I422
defaultProfile = "high422"
case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 2 && stream.ColorSpace().Subsampling.B == 0:
e.params.i_csp = C.X264_CSP_I420
defaultProfile = "high"
default:
return nil, errors.New("unsupported input chroma subsampling")
}
profile := C.CString(getSettingString(settings, "profile", defaultProfile))
defer C.free(unsafe.Pointer(profile))
if stream.ColorSpace().BitDepth > 8 {
e.params.i_csp |= C.X264_CSP_HIGH_DEPTH
if defaultProfile == "high" {
defaultProfile = "high10"
}
}
e.params.i_bitdepth = C.int(stream.ColorSpace().BitDepth)
width, height := stream.Resolution()
e.params.i_width = C.int(width)
e.params.i_height = C.int(height)
e.params.b_vfr_input = 0
e.params.b_repeat_headers = 1
e.params.b_annexb = 1
e.params.i_fps_num = C.uint32_t(stream.FrameRate().Numerator)
e.params.i_fps_den = C.uint32_t(stream.FrameRate().Denominator)
if stream.ColorRange() {
e.params.vui.b_fullrange = 1
} else {
e.params.vui.b_fullrange = 0
}
if C.x264_param_apply_profile(&e.params, profile) < 0 {
return nil, errors.New("error setting profile")
}
for k, v := range settings {
if err := func() error {
var strVal *C.char
if val, ok := v.(string); ok {
strVal = C.CString(val)
} else if val, ok := v.(int); ok {
strVal = C.CString(strconv.Itoa(val))
} else if val, ok := v.(int64); ok {
strVal = C.CString(strconv.Itoa(int(val)))
}
if strVal != nil {
defer C.free(unsafe.Pointer(strVal))
} else {
return fmt.Errorf("could not get parameter %s", k)
}
strKey := C.CString(k)
defer C.free(unsafe.Pointer(strKey))
if ret := C.x264_param_parse(&e.params, strKey, strVal); ret != 0 {
if ret == C.X264_PARAM_BAD_NAME {
return fmt.Errorf("bad parameter name %s", k)
} else if ret == C.X264_PARAM_BAD_VALUE {
return fmt.Errorf("bad parameter value %s for %s", C.GoString(strVal), k)
} else {
return fmt.Errorf("error setting parameter %s", k)
}
}
return nil
}(); err != nil {
return nil, err
}
}
runtime.SetFinalizer(e, func(encoder *Encoder) {
encoder.Close()
})
if e.pictureIn = (*C.x264_picture_t)(C.malloc(C.size_t(unsafe.Sizeof(C.x264_picture_t{})))); e.pictureIn == nil {
return nil, errors.New("error allocating memory")
}
if e.pictureOut = (*C.x264_picture_t)(C.malloc(C.size_t(unsafe.Sizeof(C.x264_picture_t{})))); e.pictureOut == nil {
return nil, errors.New("error allocating memory")
}
if C.x264_picture_alloc(e.pictureIn, e.params.i_csp, e.params.i_width, e.params.i_height) < 0 {
return nil, errors.New("error allocating input picture")
}
if e.h = C.x264_encoder_open(&e.params); e.h == nil {
return nil, errors.New("error opening encoder")
}
return e, nil
}
func (e *Encoder) Encode(pts int64, f frame.Frame) error {
//TODO: make this a Source channel
var nal *C.x264_nal_t
var i_nal C.int
var frame_size C.int
if int8F, ok := f.(frame.TypedFrame[uint8]); ok {
e.pictureIn.img.plane[0] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeLuma()[0]))
e.pictureIn.img.plane[1] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeCb()[0]))
e.pictureIn.img.plane[2] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeCr()[0]))
} else if int16F, ok := f.(frame.TypedFrame[uint16]); ok {
e.pictureIn.img.plane[0] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeLuma()[0]))
e.pictureIn.img.plane[1] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeCb()[0]))
e.pictureIn.img.plane[2] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeCr()[0]))
}
//cleanup pointers
defer func() {
e.pictureIn.img.plane[0] = nil
e.pictureIn.img.plane[1] = nil
e.pictureIn.img.plane[2] = nil
}()
defer runtime.KeepAlive(f)
e.pictureIn.i_pts = C.int64_t(pts)
if frame_size = C.x264_encoder_encode(e.h, &nal, &i_nal, e.pictureIn, e.pictureOut); frame_size < 0 {
return errors.New("error encoding frame")
}
if frame_size == 0 {
return nil
}
buf := unsafe.Slice((*byte)(nal.p_payload), int(frame_size))
if _, err := e.w.Write(buf); err != nil {
return err
}
return nil
}
func (e *Encoder) Flush() error {
var nal *C.x264_nal_t
var i_nal C.int
var frame_size C.int
for C.x264_encoder_delayed_frames(e.h) > 0 {
if frame_size = C.x264_encoder_encode(e.h, &nal, &i_nal, nil, e.pictureOut); frame_size <= 0 {
return errors.New("error encoding frame")
}
buf := unsafe.Slice((*byte)(nal.p_payload), int(frame_size))
if _, err := e.w.Write(buf); err != nil {
return err
}
}
return nil
}
func (e *Encoder) Close() {
if e.cleaned.Swap(true) == false {
if e.h != nil {
C.x264_encoder_close(e.h)
e.h = nil
}
if e.pictureIn != nil {
C.x264_picture_clean(e.pictureIn)
C.free(unsafe.Pointer(e.pictureIn))
e.pictureIn = nil
}
if e.pictureOut != nil {
C.free(unsafe.Pointer(e.pictureOut))
e.pictureOut = nil
}
}
}
func getSettingString(m map[string]any, name string, fallback string) string {
if v, ok := m[name]; ok {
if val, ok := v.(string); ok {
return val
}
if val, ok := v.(int); ok {
return strconv.Itoa(val)
}
if val, ok := v.(int64); ok {
return strconv.Itoa(int(val))
}
}
return fallback
}
func getSettingInt(m map[string]any, name string, fallback int) int {
return fallback
}

23
frame/frame.go Normal file
View file

@ -0,0 +1,23 @@
package frame
import "git.gammaspectra.live/S.O.N.G/Ignite/colorspace"
type AllowedFrameTypes interface {
uint8 | uint16
}
type Frame interface {
Width() int
Height() int
ColorSpace() colorspace.ColorSpace
Get16(x, y int) (Y uint16, Cb uint16, Cr uint16)
Get8(x, y int) (Y uint8, Cb uint8, Cr uint8)
}
type TypedFrame[T AllowedFrameTypes] interface {
Frame
GetNative(x, y int) (Y T, Cb T, Cr T)
GetNativeLuma() []T
GetNativeCb() []T
GetNativeCr() []T
}

94
frame/frame_uint16.go Normal file
View file

@ -0,0 +1,94 @@
package frame
import (
"errors"
"git.gammaspectra.live/S.O.N.G/Ignite/colorspace"
"runtime"
"unsafe"
)
type uint16Frame struct {
colorSpace colorspace.ColorSpace
width int
height int
Y []uint16
Cb []uint16
Cr []uint16
}
func NewUint16FrameFromBytes(space colorspace.ColorSpace, width, height int, data []byte) (*uint16Frame, error) {
if frameLength, _ := space.FrameSize(width, height); frameLength != len(data) {
return nil, errors.New("wrong length of data")
}
if space.BitDepth >= 16 {
return nil, errors.New("wrong bit depth")
}
buf := make([]uint16, len(data)/2)
copy(buf, unsafe.Slice((*uint16)(unsafe.Pointer(&data[0])), len(data)/2))
runtime.KeepAlive(data)
iY := space.PlaneLumaSamples(width, height)
iCb := space.PlaneCbSamples(width, height)
iCr := space.PlaneCrSamples(width, height)
return &uint16Frame{
colorSpace: space,
height: height,
width: width,
Y: buf[:iY],
Cb: buf[iY : iY+iCb],
Cr: buf[iY+iCb : iY+iCb+iCr],
}, nil
}
func (i *uint16Frame) Get16(x, y int) (Y uint16, Cb uint16, Cr uint16) {
cy, cb, cr := i.GetNative(x, y)
return cy << (16 - i.colorSpace.BitDepth), cb << (16 - i.colorSpace.BitDepth), cr << (16 - i.colorSpace.BitDepth)
}
func (i *uint16Frame) Get8(x, y int) (Y uint8, Cb uint8, Cr uint8) {
cy, cb, cr := i.GetNative(x, y)
return uint8(cy >> (i.colorSpace.BitDepth - 8)), uint8(cb >> (i.colorSpace.BitDepth - 8)), uint8(cr >> (i.colorSpace.BitDepth - 8))
}
func (i *uint16Frame) Width() int {
return i.width
}
func (i *uint16Frame) Height() int {
return i.height
}
func (i *uint16Frame) ColorSpace() colorspace.ColorSpace {
return i.colorSpace
}
func (i *uint16Frame) GetNative(x, y int) (Y uint16, Cb uint16, Cr uint16) {
Yindex := x + y*i.width
Cwidth := (i.width * int(i.colorSpace.Subsampling.A)) / int(i.colorSpace.Subsampling.J)
if i.colorSpace.Subsampling.B == 0 {
y /= 2
}
Cindex := (x*int(i.colorSpace.Subsampling.A))/int(i.colorSpace.Subsampling.J) + y*Cwidth
Y = i.Y[Yindex]
Cb = i.Cb[Cindex]
Cr = i.Cr[Cindex]
return
}
func (i *uint16Frame) GetNativeLuma() []uint16 {
return i.Y
}
func (i *uint16Frame) GetNativeCb() []uint16 {
return i.Cb
}
func (i *uint16Frame) GetNativeCr() []uint16 {
return i.Cr
}

86
frame/frame_uint8.go Normal file
View file

@ -0,0 +1,86 @@
package frame
import (
"errors"
"git.gammaspectra.live/S.O.N.G/Ignite/colorspace"
)
type uint8Frame struct {
colorSpace colorspace.ColorSpace
width int
height int
Y []uint8
Cb []uint8
Cr []uint8
}
func NewUint8FrameFromBytes(space colorspace.ColorSpace, width, height int, data []byte) (*uint8Frame, error) {
if frameLength, _ := space.FrameSize(width, height); frameLength != len(data) {
return nil, errors.New("wrong length of data")
}
if space.BitDepth > 8 {
return nil, errors.New("wrong bit depth")
}
iY := space.PlaneLumaSamples(width, height)
iCb := space.PlaneCbSamples(width, height)
iCr := space.PlaneCrSamples(width, height)
return &uint8Frame{
colorSpace: space,
height: height,
width: width,
Y: data[:iY],
Cb: data[iY : iY+iCb],
Cr: data[iY+iCb : iY+iCb+iCr],
}, nil
}
func (i *uint8Frame) Get16(x, y int) (Y uint16, Cb uint16, Cr uint16) {
cy, cb, cr := i.GetNative(x, y)
return uint16(cy) << (16 - i.colorSpace.BitDepth), uint16(cb) << (16 - i.colorSpace.BitDepth), uint16(cr) << (16 - i.colorSpace.BitDepth)
}
func (i *uint8Frame) Get8(x, y int) (Y uint8, Cb uint8, Cr uint8) {
return i.GetNative(x, y)
}
func (i *uint8Frame) Width() int {
return i.width
}
func (i *uint8Frame) Height() int {
return i.height
}
func (i *uint8Frame) ColorSpace() colorspace.ColorSpace {
return i.colorSpace
}
func (i *uint8Frame) GetNative(x, y int) (Y uint8, Cb uint8, Cr uint8) {
Yindex := x + y*i.width
Cwidth := (i.width * int(i.colorSpace.Subsampling.A)) / int(i.colorSpace.Subsampling.J)
if i.colorSpace.Subsampling.B == 0 {
y /= 2
}
Cindex := (x*int(i.colorSpace.Subsampling.A))/int(i.colorSpace.Subsampling.J) + y*Cwidth
Y = i.Y[Yindex]
Cb = i.Cb[Cindex]
Cr = i.Cr[Cindex]
return
}
func (i *uint8Frame) GetNativeLuma() []uint8 {
return i.Y
}
func (i *uint8Frame) GetNativeCb() []uint8 {
return i.Cb
}
func (i *uint8Frame) GetNativeCr() []uint8 {
return i.Cr
}

5
go.mod Normal file
View file

@ -0,0 +1,5 @@
module git.gammaspectra.live/S.O.N.G/Ignite
go 1.19
require golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561

2
go.sum Normal file
View file

@ -0,0 +1,2 @@
golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561 h1:MDc5xs78ZrZr3HMQugiXOAkSZtfTpbJLDr/lwfgO53E=
golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE=

10
utilities/ratio.go Normal file
View file

@ -0,0 +1,10 @@
package utilities
type Ratio struct {
Numerator int
Denominator int
}
func (r Ratio) Float64() float64 {
return float64(r.Numerator) / float64(r.Denominator)
}

291
y4m/y4m.go Normal file
View file

@ -0,0 +1,291 @@
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
}