44 lines
1.4 KiB
Go
44 lines
1.4 KiB
Go
package spike
|
|
|
|
import "fmt"
|
|
|
|
// LetterboxBGRToNCHW converts one packed BGR frame into the RGB, CHW,
|
|
// float32 tensor expected by the locked 640-pixel YOLO pose ONNX model.
|
|
//
|
|
// It intentionally uses nearest-neighbour scaling for this Spike. T-303
|
|
// must replace or validate this preprocessing against the V1 implementation
|
|
// before it becomes the V2 production preprocessing path.
|
|
func LetterboxBGRToNCHW(frame []byte, width, height, target int) ([]float32, error) {
|
|
if width <= 0 || height <= 0 || target <= 0 {
|
|
return nil, fmt.Errorf("frame width, height and target must be positive")
|
|
}
|
|
if len(frame) != width*height*3 {
|
|
return nil, fmt.Errorf("BGR frame length = %d, want %d", len(frame), width*height*3)
|
|
}
|
|
|
|
scaleWidth := target
|
|
scaleHeight := height * target / width
|
|
if scaleHeight > target {
|
|
scaleHeight = target
|
|
scaleWidth = width * target / height
|
|
}
|
|
padX := (target - scaleWidth) / 2
|
|
padY := (target - scaleHeight) / 2
|
|
plane := target * target
|
|
input := make([]float32, 3*plane)
|
|
|
|
for y := 0; y < scaleHeight; y++ {
|
|
sourceY := y * height / scaleHeight
|
|
for x := 0; x < scaleWidth; x++ {
|
|
sourceX := x * width / scaleWidth
|
|
source := (sourceY*width + sourceX) * 3
|
|
destination := (padY+y)*target + padX + x
|
|
input[destination] = float32(frame[source+2]) / 255.0
|
|
input[plane+destination] = float32(frame[source+1]) / 255.0
|
|
input[2*plane+destination] = float32(frame[source]) / 255.0
|
|
}
|
|
}
|
|
|
|
return input, nil
|
|
}
|