Files
silver_pose/v2/internal/spike/preprocess.go
T

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
}