Added some setters + missing display matrix methods

asticode · Mar 13, 2024 · 7543fa0 · 7543fa0
1 parent a4f27be
commit 7543fa0
Show file tree

Hide file tree

Showing 6 changed files with 91 additions and 4 deletions.
diff --git a/codec_parameters.go b/codec_parameters.go
@@ -81,6 +81,10 @@ func (cp *CodecParameters) ColorRange() ColorRange {
 	return ColorRange(cp.c.color_range)
 }
 
+func (cp *CodecParameters) SetColorRange(r ColorRange) {
+	cp.c.color_range = C.enum_AVColorRange(r)
+}
+
 func (cp *CodecParameters) ColorSpace() ColorSpace {
 	return ColorSpace(cp.c.color_space)
 }
@@ -93,6 +97,10 @@ func (cp *CodecParameters) FrameSize() int {
 	return int(cp.c.frame_size)
 }
 
+func (cp *CodecParameters) SetFrameSize(i int) {
+	cp.c.frame_size = C.int(i)
+}
+
 func (cp *CodecParameters) Height() int {
 	return int(cp.c.height)
 }
@@ -109,6 +117,10 @@ func (cp *CodecParameters) MediaType() MediaType {
 	return MediaType(cp.c.codec_type)
 }
 
+func (cp *CodecParameters) SetMediaType(t MediaType) {
+	cp.c.codec_type = C.enum_AVMediaType(t)
+}
+
 func (cp *CodecParameters) PixelFormat() PixelFormat {
 	return PixelFormat(cp.c.format)
 }

diff --git a/codec_parameters_test.go b/codec_parameters_test.go
@@ -75,10 +75,16 @@ func TestCodecParameters(t *testing.T) {
 	require.Equal(t, CodecIDRawvideo, cp6.CodecID())
 	cp6.SetCodecTag(CodecTag(2))
 	require.Equal(t, CodecTag(2), cp6.CodecTag())
+	cp6.SetColorRange(ColorRangeJpeg)
+	require.Equal(t, ColorRangeJpeg, cp6.ColorRange())
 	cp6.SetCodecType(MediaTypeAudio)
 	require.Equal(t, MediaTypeAudio, cp6.CodecType())
+	cp6.SetFrameSize(1)
+	require.Equal(t, 1, cp6.FrameSize())
 	cp6.SetHeight(1)
 	require.Equal(t, 1, cp6.Height())
+	cp6.SetMediaType(MediaTypeAudio)
+	require.Equal(t, MediaTypeAudio, cp6.MediaType())
 	cp6.SetPixelFormat(PixelFormat0Bgr)
 	require.Equal(t, PixelFormat0Bgr, cp6.PixelFormat())
 	cp6.SetSampleAspectRatio(NewRational(1, 2))

diff --git a/display_matrix.go b/display_matrix.go
@@ -28,6 +28,21 @@ func NewDisplayMatrixFromBytes(b []byte) (m *DisplayMatrix, err error) {
 	return
 }
 
+func NewDisplayMatrixFromRotation(angle float64) *DisplayMatrix {
+	m := &DisplayMatrix{}
+	C.av_display_rotation_set((*C.int32_t)(unsafe.Pointer(&m[0])), C.double(angle))
+	return m
+}
+
+func (m DisplayMatrix) Bytes() []byte {
+	b := make([]byte, 0, 36)
+	for _, v := range m {
+		b = binary.LittleEndian.AppendUint32(b, v)
+	}
+	return b
+}
+
+// Rotation is a clockwise angle
 func (m DisplayMatrix) Rotation() float64 {
-	return float64(C.av_display_rotation_get((*C.int32_t)(unsafe.Pointer(&m[0]))))
+	return -float64(C.av_display_rotation_get((*C.int32_t)(unsafe.Pointer(&m[0]))))
 }
diff --git a/display_matrix_test.go b/display_matrix_test.go
@@ -9,12 +9,16 @@ import (
 func TestDisplayMatrix(t *testing.T) {
 	_, err := NewDisplayMatrixFromBytes([]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})
 	require.Error(t, err)
-	dm, err := NewDisplayMatrixFromBytes([]byte{0, 0, 0, 0, 0, 0, 255, 255, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 64})
+	b := []byte{0, 0, 0, 0, 0, 0, 255, 255, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 64}
+	dm, err := NewDisplayMatrixFromBytes(b)
 	require.NoError(t, err)
 	require.Equal(t, DisplayMatrix{0x0, 0xffff0000, 0x0, 0x10000, 0x0, 0x0, 0x0, 0x0, 0x40000000}, *dm)
-	require.Equal(t, float64(90), dm.Rotation())
+	require.Equal(t, -90.0, dm.Rotation())
+	require.Equal(t, b, dm.Bytes())
+	dm = NewDisplayMatrixFromRotation(-90)
+	require.Equal(t, -90.0, dm.Rotation())
 	dm, err = NewDisplayMatrixFromBytes([]byte{0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 64})
 	require.NoError(t, err)
 	require.Equal(t, DisplayMatrix{0x0, 0x10000, 0x0, 0xffff0000, 0x0, 0x0, 0x0, 0x0, 0x40000000}, *dm)
-	require.Equal(t, float64(-90), dm.Rotation())
+	require.Equal(t, 90.0, dm.Rotation())
 }
diff --git a/stream.go b/stream.go
@@ -3,6 +3,10 @@ package astiav
 //#cgo pkg-config: libavformat
 //#include <libavformat/avformat.h>
 import "C"
+import (
+	"errors"
+	"unsafe"
+)
 
 // https://github.com/FFmpeg/FFmpeg/blob/n5.0/libavformat/avformat.h#L937
 type Stream struct {
@@ -20,6 +24,10 @@ func (s *Stream) AvgFrameRate() Rational {
 	return newRationalFromC(s.c.avg_frame_rate)
 }
 
+func (s *Stream) SetAvgFrameRate(r Rational) {
+	s.c.avg_frame_rate = r.c
+}
+
 func (s *Stream) CodecParameters() *CodecParameters {
 	return newCodecParametersFromC(s.c.codecpar)
 }
@@ -36,10 +44,18 @@ func (s *Stream) ID() int {
 	return int(s.c.id)
 }
 
+func (s *Stream) SetID(i int) {
+	s.c.id = C.int(i)
+}
+
 func (s *Stream) Index() int {
 	return int(s.c.index)
 }
 
+func (s *Stream) SetIndex(i int) {
+	s.c.index = C.int(i)
+}
+
 func (s *Stream) Metadata() *Dictionary {
 	return newDictionaryFromC(s.c.metadata)
 }
@@ -52,16 +68,38 @@ func (s *Stream) RFrameRate() Rational {
 	return newRationalFromC(s.c.r_frame_rate)
 }
 
+func (s *Stream) SetRFrameRate(r Rational) {
+	s.c.r_frame_rate = r.c
+}
+
 func (s *Stream) SampleAspectRatio() Rational {
 	return newRationalFromC(s.c.sample_aspect_ratio)
 }
 
+func (s *Stream) SetSampleAspectRatio(r Rational) {
+	s.c.sample_aspect_ratio = r.c
+}
+
 func (s *Stream) SideData(t PacketSideDataType) []byte {
 	return bytesFromC(func(size *cUlong) *C.uint8_t {
 		return C.av_stream_get_side_data(s.c, (C.enum_AVPacketSideDataType)(t), size)
 	})
 }
 
+func (s *Stream) AddSideData(t PacketSideDataType, d []byte) error {
+	if len(d) == 0 {
+		return nil
+	}
+
+	ptr := C.av_stream_new_side_data(s.c, (C.enum_AVPacketSideDataType)(t), cUlong(len(d)))
+	if ptr == nil {
+		return errors.New("astiav: nil pointer")
+	}
+
+	C.memcpy(unsafe.Pointer(ptr), unsafe.Pointer(&d[0]), cUlong(len(d)))
+	return nil
+}
+
 func (s *Stream) StartTime() int64 {
 	return int64(s.c.start_time)
 }

diff --git a/stream_test.go b/stream_test.go
@@ -34,6 +34,18 @@ func TestStream(t *testing.T) {
 	require.Equal(t, int64(0), s2.StartTime())
 	require.Equal(t, NewRational(1, 48000), s2.TimeBase())
 
+	s1.SetAvgFrameRate(NewRational(2, 1))
+	require.Equal(t, NewRational(2, 1), s1.AvgFrameRate())
+	s1.SetID(2)
+	require.Equal(t, 2, s1.ID())
+	s1.SetIndex(1)
+	require.Equal(t, 1, s1.Index())
+	s1.SetRFrameRate(NewRational(2, 1))
+	require.Equal(t, NewRational(2, 1), s1.RFrameRate())
+	s1.SetSampleAspectRatio(NewRational(2, 1))
+	require.Equal(t, NewRational(2, 1), s1.SampleAspectRatio())
+	require.NoError(t, s1.AddSideData(PacketSideDataTypeEncryptionInfo, []byte("test")))
+	require.Equal(t, []byte("test"), s1.SideData(PacketSideDataTypeEncryptionInfo))
 	s1.SetTimeBase(NewRational(1, 1))
 	require.Equal(t, NewRational(1, 1), s1.TimeBase())
 }