Skip to content

Commit

Permalink
Clean up
Browse files Browse the repository at this point in the history
  • Loading branch information
@asticode
asticode committed Oct 23, 2024
1 parent e1f7d0c commit e9e602f
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 107 deletions.
164 changes: 61 additions & 103 deletions examples/hardware_encoding/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import (
"flag"
"fmt"
"log"
"os"
"strings"

"github.com/asticode/go-astiav"
Expand All @@ -16,14 +15,8 @@ var (
hardwareDeviceName = flag.String("n", "", "the hardware device name (e.g. 0)")
hardwareDeviceTypeName = flag.String("t", "", "the hardware device type (e.g. cuda)")
hardwarePixelFormatName = flag.String("hpf", "", "the hardware pixel format name (e.g. cuda)")

width = flag.Int("w", 1920, "the width")
height = flag.Int("h", 1080, "the height")
fps = flag.Int("f", 25, "the fps")
initialPoolSize = flag.Int("p", 20, "the initial pool size")
patternGridSize = flag.Int("g", 128, "the pattern grid size")

output = flag.String("o", "", "the output path")
height = flag.Int("h", 1080, "the height")
width = flag.Int("w", 1920, "the width")
)

func main() {
Expand All @@ -43,18 +36,11 @@ func main() {
flag.Parse()

// Usage
if *hardwareDeviceTypeName == "" || *encoderCodecName == "" || *hardwarePixelFormatName == "" || *output == "" {
log.Println("Usage: <binary path> -t <hardware device type> -c <encoder codec> -hpf <hardware pixel format> -o <output path> [-n <hardware device name> -w <width> -h <height> -f <fps> -p <initial pool size> -g <pattern grid size>]")
if *hardwareDeviceTypeName == "" || *encoderCodecName == "" || *hardwarePixelFormatName == "" {
log.Println("Usage: <binary path> -t <hardware device type> -c <encoder codec> -hpf <hardware pixel format> [-n <hardware device name> -w <width> -h <height>]")
return
}

// Open output file
output, err := os.OpenFile(*output, os.O_CREATE|os.O_WRONLY, 0o644)
if err != nil {
log.Fatal(fmt.Errorf("main: opening output file failed: %w", err))
}
defer output.Close()

// Get hardware device type
hardwareDeviceType := astiav.FindHardwareDeviceTypeByName(*hardwareDeviceTypeName)
if hardwareDeviceType == astiav.HardwareDeviceTypeNone {
Expand Down Expand Up @@ -83,122 +69,94 @@ func main() {
// Set codec context
encCodecContext.SetWidth(*width)
encCodecContext.SetHeight(*height)
encCodecContext.SetTimeBase(astiav.NewRational(1, *fps))
encCodecContext.SetFramerate(astiav.NewRational(*fps, 1))
encCodecContext.SetTimeBase(astiav.NewRational(1, 25))
encCodecContext.SetFramerate(encCodecContext.TimeBase().Invert())
hardwarePixelFormatName := astiav.FindPixelFormatByName(*hardwarePixelFormatName)
if hardwarePixelFormatName == astiav.PixelFormatNone {
log.Fatal("main: hardware pixel format not found")
}
encCodecContext.SetPixelFormat(hardwarePixelFormatName)

// Set hardware frame context
hardwareFrameCtx := astiav.AllocHardwareFrameContext(hardwareDeviceContext)
if hardwareFrameCtx == nil {
// Alloc hardware frame context
hardwareFrameContext := astiav.AllocHardwareFrameContext(hardwareDeviceContext)
if hardwareFrameContext == nil {
log.Fatal("main: hardware frame context is nil")
}
hardwareFrameCtx.SetPixelFormat(hardwarePixelFormatName)
hardwareFrameCtx.SetSoftwarePixelFormat(astiav.PixelFormatNv12)
hardwareFrameCtx.SetWidth(*width)
hardwareFrameCtx.SetHeight(*height)
hardwareFrameCtx.SetInitialPoolSize(*initialPoolSize)

// Set hardware frame content
const softwarePixelFormat = astiav.PixelFormatNv12
hardwareFrameContext.SetPixelFormat(hardwarePixelFormatName)
hardwareFrameContext.SetSoftwarePixelFormat(softwarePixelFormat)
hardwareFrameContext.SetWidth(*width)
hardwareFrameContext.SetHeight(*height)
hardwareFrameContext.SetInitialPoolSize(20)

// Initialize hardware frame context
if err := hardwareFrameCtx.Initialize(); err != nil {
if err := hardwareFrameContext.Initialize(); err != nil {
log.Fatal(fmt.Errorf("main: initializing hardware frame context failed: %w", err))
}
encCodecContext.SetHardwareFrameContext(hardwareFrameCtx)

// Update hardware frame context
encCodecContext.SetHardwareFrameContext(hardwareFrameContext)

// Open codec context
if err := encCodecContext.Open(encCodec, nil); err != nil {
log.Fatal(fmt.Errorf("main: opening codec context failed: %w", err))
}

frameIndex := 0
// Alloc software frame
softwareFrame := astiav.AllocFrame()
defer softwareFrame.Free()

// Draw frames, upload them to hardware devices, and encode them
for {
// Alloc software frame
softwareFrame := astiav.AllocFrame()
// Set software frame
softwareFrame.SetWidth(*width)
softwareFrame.SetHeight(*height)
softwareFrame.SetPixelFormat(softwarePixelFormat)

// Set software frame
softwareFrame.SetWidth(*width)
softwareFrame.SetHeight(*height)
softwareFrame.SetPixelFormat(astiav.PixelFormatNv12)
// Alloc software frame buffer
if err := softwareFrame.AllocBuffer(0); err != nil {
log.Fatal(fmt.Errorf("main: allocating buffer failed: %w", err))
}

// Alloc software frame buffer
if err := softwareFrame.AllocBuffer(0); err != nil {
log.Fatal(fmt.Errorf("main: allocating buffer failed: %w", err))
}
// Fill software frame with black
if err = softwareFrame.ImageFillBlack(); err != nil {
log.Fatal(fmt.Errorf("main: filling software frame with black failed: %w", err))
}

// Fill software frame
yPlane, uvPlane := MakeNV12MovingCheckerboardPattern(*width, *height, *patternGridSize, frameIndex)
softwareFrame.SetData(0, yPlane)
softwareFrame.SetData(1, uvPlane)
// Alloc hardware frame
hardwareFrame := astiav.AllocFrame()
defer hardwareFrame.Free()

// Alloc hardware frame
hardwareFrame := astiav.AllocFrame()
// Alloc hardware frame buffer
if err := hardwareFrame.AllocHardwareBuffer(hardwareFrameContext); err != nil {
log.Fatal(fmt.Errorf("main: allocating hardware buffer failed: %w", err))
}

// Alloc hardware frame buffer
if err := hardwareFrame.AllocHardwareBuffer(hardwareFrameCtx); err != nil {
log.Fatal(fmt.Errorf("main: allocating hardware buffer failed: %w", err))
}
// Transfer software frame data to hardware frame
if err := softwareFrame.TransferHardwareData(hardwareFrame); err != nil {
log.Fatal(fmt.Errorf("main: transferring hardware data failed: %w", err))
}

// Upload software frame to hardware frame
if err := softwareFrame.TransferHardwareData(hardwareFrame); err != nil {
log.Fatal(fmt.Errorf("main: uploading from frame failed: %w", err))
}
softwareFrame.Free()
// Encode frame
if err := encCodecContext.SendFrame(hardwareFrame); err != nil {
log.Fatal(fmt.Errorf("main: sending frame failed: %w", err))
}

// Encode frame
if err := encCodecContext.SendFrame(hardwareFrame); err != nil {
log.Fatal(fmt.Errorf("main: sending frame failed: %w", err))
}
hardwareFrame.Free()
// Alloc packet
pkt := astiav.AllocPacket()
defer pkt.Free()

// Loop
for {
// Receive packet
packet := astiav.AllocPacket()
for {
if err := encCodecContext.ReceivePacket(packet); err != nil {
if err = encCodecContext.ReceivePacket(pkt); err != nil {
if errors.Is(err, astiav.ErrEof) || errors.Is(err, astiav.ErrEagain) {
break
}

// Write packet
if _, err := output.Write(packet.Data()); err != nil {
log.Fatal(fmt.Errorf("main: writing packet failed: %w", err))
}

packet.Unref()
log.Fatal(fmt.Errorf("main: receiving packet failed: %w", err))
}
packet.Free()

frameIndex++
log.Printf("Finished encoding frame %d\n", frameIndex)
// Log
log.Println("new packet")
}
}

func MakeNV12MovingCheckerboardPattern(width, height, blockSize, frame int) ([]byte, []byte) {
yPlane := make([]byte, width*height)
uvPlane := make([]byte, width*height/2)

xOffset := frame % blockSize
yOffset := frame % blockSize

// Y plane (checkerboard pattern)
for y := 0; y < height; y++ {
for x := 0; x < width; x++ {
if ((x+xOffset)/blockSize+(y+yOffset)/blockSize)%2 == 0 {
yPlane[y*width+x] = 255 // White
} else {
yPlane[y*width+x] = 0 // Black
}
}
}

// UV plane (gray)
for i := 0; i < len(uvPlane); i += 2 {
uvPlane[i] = 128 // U component (neutral)
uvPlane[i+1] = 128 // V component (neutral)
}

return yPlane, uvPlane
}
4 changes: 0 additions & 4 deletions frame.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,10 +73,6 @@ func (f *Frame) Data() *FrameData {
return newFrameData(newFrameDataFrame(f))
}

func (f *Frame) SetData(plane int, data []byte) {
f.c.data[plane] = (*C.uint8_t)(unsafe.Pointer(&data[0]))
}

func (f *Frame) Height() int {
return int(f.c.height)
}
Expand Down

0 comments on commit e9e602f

Please sign in to comment.