Skip to content

Commit

Permalink
Fix dryness and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@hexbabe
hexbabe committed Jan 29, 2025
1 parent ad63a2b commit e796c3a
Show file tree
Hide file tree
Showing 3 changed files with 119 additions and 54 deletions.
78 changes: 43 additions & 35 deletions mime.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,13 @@ package viamrtsp
#include <libavutil/dict.h>
#include <libavutil/frame.h>
#include <libavutil/imgutils.h>
#include <libavutil/pixfmt.h>
#include <libswscale/swscale.h>
#include <stdlib.h>
*/
import "C"

import (
"bytes"
"encoding/binary"
"errors"
"sync"
Expand All @@ -26,12 +26,10 @@ import (
)

const (
yuyvMagicString = "YUYV"
yuyvHeaderDimBytes = 4
yuyvBytesPerPixel = 2
rgbaMagicString = "RGBA"
rgbaHeaderDimBytes = 8
rgbaBytesPerPixel = 4
yuyvMagicString = "YUYV"
yuyvBytesPerPixel = 2
rgbaMagicString = "RGBA"
rgbaBytesPerPixel = 4
)

type mimeHandler struct {
Expand Down Expand Up @@ -139,7 +137,6 @@ func (mh *mimeHandler) convertPixelFormat(
mu *sync.Mutex,
initContext func(*C.AVFrame) error,
bytesPerPixel int,
headerDimBytes int,
mimeType string,
) ([]byte, camera.ImageMetadata, error) {
if frame == nil {
Expand Down Expand Up @@ -176,7 +173,7 @@ func (mh *mimeHandler) convertPixelFormat(
}

dataSize := int((*dstFrame).width) * int((*dstFrame).height) * bytesPerPixel
header := packHeader(format, int((*dstFrame).width), int((*dstFrame).height), headerDimBytes)
header := packHeader(format, int((*dstFrame).width), int((*dstFrame).height))
data := make([]byte, len(header)+dataSize)
copy(data[0:], header)
C.memcpy(unsafe.Pointer(&data[len(header)]), unsafe.Pointer((*dstFrame).data[0]), C.size_t(dataSize))
Expand All @@ -195,7 +192,6 @@ func (mh *mimeHandler) convertYUYV(frame *C.AVFrame) ([]byte, camera.ImageMetada
&mh.yuyvMu,
mh.initYUYVContext,
yuyvBytesPerPixel,
yuyvHeaderDimBytes,
mimeTypeYUYV,
)
}
Expand All @@ -209,7 +205,6 @@ func (mh *mimeHandler) convertRGBA(frame *C.AVFrame) ([]byte, camera.ImageMetada
&mh.rgbaMu,
mh.initRGBAContext,
rgbaBytesPerPixel,
rgbaHeaderDimBytes,
rutils.MimeTypeRawRGBA,
)
}
Expand All @@ -224,35 +219,50 @@ func (mh *mimeHandler) convertRGBA(frame *C.AVFrame) ([]byte, camera.ImageMetada
// - dstFrame: Pointer to AVFrame pointer that will be initialized
//
// Returns error if any allocation or initialization fails
func (mh *mimeHandler) initPixelFormatContext(frame *C.AVFrame, pixFmt C.int, swsCtxPtr **C.struct_SwsContext, dstFrame **C.AVFrame) error {
func (mh *mimeHandler) initPixelFormatContext(
frame *C.AVFrame,
pixFmt C.enum_AVPixelFormat,
swsCtxPtr **C.struct_SwsContext,
dstFrame **C.AVFrame,
) error {
mh.logger.Infof("creating sws context with frame size: %dx%d for format %d", frame.width, frame.height, pixFmt)

if *swsCtxPtr != nil {
C.sws_freeContext(*swsCtxPtr)
*swsCtxPtr = nil
}
if *dstFrame != nil {
C.av_frame_free(dstFrame)
*dstFrame = nil
}
*dstFrame = C.av_frame_alloc()
if *dstFrame == nil {

newFrame := C.av_frame_alloc()
if newFrame == nil {
return errors.New("failed to allocate frame")
}
(*dstFrame).width = frame.width
(*dstFrame).height = frame.height
(*dstFrame).format = pixFmt
if res := C.av_frame_get_buffer(*dstFrame, 32); res < 0 {
C.av_frame_free(dstFrame)

newFrame.width = frame.width
newFrame.height = frame.height
newFrame.format = C.int(pixFmt)

if res := C.av_frame_get_buffer(newFrame, 32); res < 0 {
C.av_frame_free(&newFrame)
return newAvError(res, "failed to allocate buffer for frame")
}
*swsCtxPtr = C.sws_getContext(
frame.width, frame.height, C.AV_PIX_FMT_YUV420P,

newSwsCtx := C.sws_getContext(
frame.width, frame.height, pixFmt,
frame.width, frame.height, pixFmt,
C.SWS_FAST_BILINEAR, nil, nil, nil,
)
if *swsCtxPtr == nil {
C.av_frame_free(dstFrame)
*dstFrame = nil

if newSwsCtx == nil {
C.av_frame_free(&newFrame)
return errors.New("failed to create converter")
}

*dstFrame = newFrame
*swsCtxPtr = newSwsCtx
return nil
}

Expand Down Expand Up @@ -292,23 +302,21 @@ func (mh *mimeHandler) close() {
// - Format string (4 bytes): A fixed string indicating the format (e.g., "YUYV" or "RGBA").
// - Width (4 bytes): The width of the image, stored in big-endian format.
// - Height (4 bytes): The height of the image, stored in big-endian format.
func packHeader(format string, width, height int, dimBytes int) []byte {
var header bytes.Buffer
header.WriteString(format)
tmp := make([]byte, dimBytes)
binary.BigEndian.PutUint32(tmp, uint32(width))
header.Write(tmp)
binary.BigEndian.PutUint32(tmp, uint32(height))
header.Write(tmp)
return header.Bytes()
func packHeader(format string, width, height int) []byte {
headerSize := 12
headerBytes := make([]byte, headerSize)
copy(headerBytes[0:4], []byte(format))
binary.BigEndian.PutUint32(headerBytes[4:8], uint32(width))
binary.BigEndian.PutUint32(headerBytes[8:12], uint32(height))
return headerBytes
}

// packYUYVHeader creates a header for YUYV data with the given width and height.
func packYUYVHeader(width, height int) []byte {
return packHeader(yuyvMagicString, width, height, yuyvHeaderDimBytes)
return packHeader(yuyvMagicString, width, height)
}

// packRGBAHeader creates a header for RGBA data with the given width and height.
func packRGBAHeader(width, height int) []byte {
return packHeader(rgbaMagicString, width, height, rgbaHeaderDimBytes)
return packHeader(rgbaMagicString, width, height)
}
30 changes: 30 additions & 0 deletions mime_cgo.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,3 +86,33 @@ func fillDummyYUV420PData(frame *C.AVFrame) {
func freeFrame(frame *C.AVFrame) {
C.av_frame_free(&frame)
}

func createTestRGBAFrame(width, height int) *C.AVFrame {
return createTestFrame(width, height, C.AV_PIX_FMT_RGBA)
}

func fillDummyRGBAData(frame *C.AVFrame) {
width := int(frame.width)
height := int(frame.height)
linesize := int(frame.linesize[0])

rgbaPlane := (*[1 << 30]uint8)(unsafe.Pointer(frame.data[0]))[: height*linesize : height*linesize]

// Top half: solid red (255, 0, 0, 255)
// Bottom half: solid blue (0, 0, 255, 255)
for y := range height {
for x := range width {
idx := y*linesize + x*rgbaBytesPerPixel
if y < height/2 {
rgbaPlane[idx+0] = 255 // R
rgbaPlane[idx+1] = 0 // G
rgbaPlane[idx+2] = 0 // B
} else {
rgbaPlane[idx+0] = 0 // R
rgbaPlane[idx+1] = 0 // G
rgbaPlane[idx+2] = 255 // B
}
rgbaPlane[idx+3] = 255 // A always fully opaque
}
}
}
65 changes: 46 additions & 19 deletions mime_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,11 +93,12 @@ func TestYUYVConvert(t *testing.T) {
mh := newMimeHandler(logger)
bytes, metadata, err := mh.convertYUYV(frame)
test.That(t, err, test.ShouldNotBeNil)
test.That(t, err.Error(), test.ShouldContainSubstring, "failed to allocate buffer for YUYV")
test.That(t, err.Error(), test.ShouldContainSubstring, "failed to allocate buffer")
test.That(t, bytes, test.ShouldBeNil)
test.That(t, metadata.MimeType, test.ShouldBeEmpty)
})

//nolint:dupl
t.Run("test yuyv magic header", func(t *testing.T) {
origWidth := 640
origHeight := 480
Expand All @@ -116,34 +117,43 @@ func TestYUYVConvert(t *testing.T) {
}

func TestRGBAConvert(t *testing.T) {
t.Run("valid YUV420P frame succeeds", func(t *testing.T) {
t.Run("valid RGBA frame succeeds", func(t *testing.T) {
width, height := 640, 480
frame := createTestYUV420PFrame(width, height)
frame := createTestRGBAFrame(width, height)
test.That(t, frame, test.ShouldNotBeNil)
defer freeFrame(frame)
fillDummyYUV420PData(frame)
fillDummyRGBAData(frame)
logger := logging.NewDebugLogger("mime_test")
mh := newMimeHandler(logger)
bytes, metadata, err := mh.convertRGBA(frame)
test.That(t, err, test.ShouldBeNil)
test.That(t, bytes, test.ShouldNotBeNil)
test.That(t, len(bytes), test.ShouldEqual, width*height*rgbaBytesPerPixel+12) // header size
test.That(t, metadata.MimeType, test.ShouldEqual, rutils.MimeTypeRawRGBA)
})

t.Run("valid YUVJ420P frame succeeds", func(t *testing.T) {
width, height := 640, 480
frame := createTestYUVJ420PFrame(width, height)
test.That(t, frame, test.ShouldNotBeNil)
defer freeFrame(frame)
fillDummyYUV420PData(frame)
logger := logging.NewDebugLogger("mime_test")
mh := newMimeHandler(logger)
bytes, metadata, err := mh.convertRGBA(frame)
test.That(t, err, test.ShouldBeNil)
test.That(t, bytes, test.ShouldNotBeNil)
test.That(t, len(bytes), test.ShouldEqual, width*height*rgbaBytesPerPixel+12) // header size
test.That(t, metadata.MimeType, test.ShouldEqual, rutils.MimeTypeRawRGBA)
// Verify the header
header := bytes[:12]
test.That(t, header[0], test.ShouldEqual, 'R')
test.That(t, header[1], test.ShouldEqual, 'G')
test.That(t, header[2], test.ShouldEqual, 'B')
test.That(t, header[3], test.ShouldEqual, 'A')

// Verify the RGBA data
data := bytes[12:]

// Check top half (should be red)
idx := (height / 4) * width * rgbaBytesPerPixel // Point in top half
test.That(t, data[idx+0], test.ShouldEqual, byte(255)) // R
test.That(t, data[idx+1], test.ShouldEqual, byte(0)) // G
test.That(t, data[idx+2], test.ShouldEqual, byte(0)) // B
test.That(t, data[idx+3], test.ShouldEqual, byte(255)) // A

// Check bottom half (should be blue)
idx = (3 * height / 4) * width * rgbaBytesPerPixel // Point in bottom half
test.That(t, data[idx+0], test.ShouldEqual, byte(0)) // R
test.That(t, data[idx+1], test.ShouldEqual, byte(0)) // G
test.That(t, data[idx+2], test.ShouldEqual, byte(255)) // B
test.That(t, data[idx+3], test.ShouldEqual, byte(255)) // A
})

t.Run("invalid frame fails", func(t *testing.T) {
Expand All @@ -154,8 +164,25 @@ func TestRGBAConvert(t *testing.T) {
mh := newMimeHandler(logger)
bytes, metadata, err := mh.convertRGBA(frame)
test.That(t, err, test.ShouldNotBeNil)
test.That(t, err.Error(), test.ShouldContainSubstring, "failed to allocate buffer for RGBA")
test.That(t, err.Error(), test.ShouldContainSubstring, "failed to allocate buffer")
test.That(t, bytes, test.ShouldBeNil)
test.That(t, metadata.MimeType, test.ShouldBeEmpty)
})

//nolint:dupl
t.Run("test rgba magic header", func(t *testing.T) {
origWidth := 640
origHeight := 480
header := packRGBAHeader(origWidth, origHeight)
test.That(t, header, test.ShouldNotBeNil)
test.That(t, len(header), test.ShouldEqual, 12)
test.That(t, header[0], test.ShouldEqual, 'R')
test.That(t, header[1], test.ShouldEqual, 'G')
test.That(t, header[2], test.ShouldEqual, 'B')
test.That(t, header[3], test.ShouldEqual, 'A')
parsedWidth := int(binary.BigEndian.Uint32(header[4:8]))
test.That(t, parsedWidth, test.ShouldEqual, origWidth)
parsedHeight := int(binary.BigEndian.Uint32(header[8:12]))
test.That(t, parsedHeight, test.ShouldEqual, origHeight)
})
}

0 comments on commit e796c3a

Please sign in to comment.