commit 84908c0ef205d7bd1042d71b4e2ca3f273b9d350 Author: WeebDataHoarder <57538841+WeebDataHoarder@users.noreply.github.com> Date: Wed Sep 14 21:42:57 2022 +0200 Initial commit: utilities, y4m parser, x264 encoder diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..62c8935 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.idea/ \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..44831fe --- /dev/null +++ b/LICENSE @@ -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. diff --git a/colorspace/colorspace.go b/colorspace/colorspace.go new file mode 100644 index 0000000..7d28a65 --- /dev/null +++ b/colorspace/colorspace.go @@ -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") +) diff --git a/encoder/x264/x264.go b/encoder/x264/x264.go new file mode 100644 index 0000000..925b1ec --- /dev/null +++ b/encoder/x264/x264.go @@ -0,0 +1,244 @@ +package x264 + +/* +#cgo pkg-config: x264 +#include +#include +#include +#include +*/ +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 +} diff --git a/frame/frame.go b/frame/frame.go new file mode 100644 index 0000000..4fdef8d --- /dev/null +++ b/frame/frame.go @@ -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 +} diff --git a/frame/frame_uint16.go b/frame/frame_uint16.go new file mode 100644 index 0000000..d780f8c --- /dev/null +++ b/frame/frame_uint16.go @@ -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 +} diff --git a/frame/frame_uint8.go b/frame/frame_uint8.go new file mode 100644 index 0000000..0eca4a2 --- /dev/null +++ b/frame/frame_uint8.go @@ -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 +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..1821e95 --- /dev/null +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..32d6375 --- /dev/null +++ b/go.sum @@ -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= diff --git a/utilities/ratio.go b/utilities/ratio.go new file mode 100644 index 0000000..b562ed7 --- /dev/null +++ b/utilities/ratio.go @@ -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) +} diff --git a/y4m/y4m.go b/y4m/y4m.go new file mode 100644 index 0000000..2cef954 --- /dev/null +++ b/y4m/y4m.go @@ -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 +}