88 lines
3.2 KiB
Go
88 lines
3.2 KiB
Go
|
package pose
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"math"
|
||
|
|
)
|
||
|
|
|
||
|
|
const letterboxPadding = byte(114)
|
||
|
|
|
||
|
|
// PreprocessBGR reproduces the fixed-shape Ultralytics letterbox contract:
|
||
|
|
// BGR source pixels are resized with bilinear interpolation, padded in 114
|
||
|
|
// gray, converted to RGB/CHW and normalised to [0, 1].
|
||
|
|
func PreprocessBGR(frame []byte, width, height, target int) ([]float32, LetterboxTransform, error) {
|
||
|
|
if width <= 0 || height <= 0 || target <= 0 {
|
||
|
|
return nil, LetterboxTransform{}, fmt.Errorf("frame width, height and target must be positive")
|
||
|
|
}
|
||
|
|
if len(frame) != width*height*3 {
|
||
|
|
return nil, LetterboxTransform{}, fmt.Errorf("BGR frame length = %d, want %d", len(frame), width*height*3)
|
||
|
|
}
|
||
|
|
|
||
|
|
scale := math.Min(float64(target)/float64(width), float64(target)/float64(height))
|
||
|
|
resizedWidth := int(math.Round(float64(width) * scale))
|
||
|
|
resizedHeight := int(math.Round(float64(height) * scale))
|
||
|
|
padWidth := target - resizedWidth
|
||
|
|
padHeight := target - resizedHeight
|
||
|
|
padLeft := int(math.Round(float64(padWidth)/2.0 - 0.1))
|
||
|
|
padTop := int(math.Round(float64(padHeight)/2.0 - 0.1))
|
||
|
|
transform := LetterboxTransform{
|
||
|
|
OriginalWidth: width, OriginalHeight: height, InputSize: target,
|
||
|
|
ResizedWidth: resizedWidth, ResizedHeight: resizedHeight,
|
||
|
|
PadLeft: padLeft, PadTop: padTop, Scale: float32(scale),
|
||
|
|
}
|
||
|
|
|
||
|
|
plane := target * target
|
||
|
|
input := make([]float32, 3*plane)
|
||
|
|
padding := float32(letterboxPadding) / 255.0
|
||
|
|
for index := range input {
|
||
|
|
input[index] = padding
|
||
|
|
}
|
||
|
|
for y := 0; y < resizedHeight; y++ {
|
||
|
|
sourceY := clampFloat((float64(y)+0.5)*float64(height)/float64(resizedHeight)-0.5, 0, float64(height-1))
|
||
|
|
y0 := int(math.Floor(sourceY))
|
||
|
|
y1 := minInt(y0+1, height-1)
|
||
|
|
yWeight := float32(sourceY - float64(y0))
|
||
|
|
for x := 0; x < resizedWidth; x++ {
|
||
|
|
sourceX := clampFloat((float64(x)+0.5)*float64(width)/float64(resizedWidth)-0.5, 0, float64(width-1))
|
||
|
|
x0 := int(math.Floor(sourceX))
|
||
|
|
x1 := minInt(x0+1, width-1)
|
||
|
|
xWeight := float32(sourceX - float64(x0))
|
||
|
|
leftTop := (y0*width + x0) * 3
|
||
|
|
rightTop := (y0*width + x1) * 3
|
||
|
|
leftBottom := (y1*width + x0) * 3
|
||
|
|
rightBottom := (y1*width + x1) * 3
|
||
|
|
destination := (padTop+y)*target + padLeft + x
|
||
|
|
for sourceChannel := 0; sourceChannel < 3; sourceChannel++ {
|
||
|
|
value := bilinear(
|
||
|
|
frame[leftTop+sourceChannel], frame[rightTop+sourceChannel],
|
||
|
|
frame[leftBottom+sourceChannel], frame[rightBottom+sourceChannel],
|
||
|
|
xWeight, yWeight,
|
||
|
|
)
|
||
|
|
// BGR source maps to RGB tensor planes.
|
||
|
|
tensorChannel := 2 - sourceChannel
|
||
|
|
// cv2.resize writes uint8 pixels before Ultralytics converts the
|
||
|
|
// image to float; round here to retain that observable contract.
|
||
|
|
input[tensorChannel*plane+destination] = float32(math.Round(float64(value))) / 255.0
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return input, transform, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func bilinear(topLeft, topRight, bottomLeft, bottomRight byte, xWeight, yWeight float32) float32 {
|
||
|
|
top := float32(topLeft)*(1-xWeight) + float32(topRight)*xWeight
|
||
|
|
bottom := float32(bottomLeft)*(1-xWeight) + float32(bottomRight)*xWeight
|
||
|
|
return top*(1-yWeight) + bottom*yWeight
|
||
|
|
}
|
||
|
|
|
||
|
|
func clampFloat(value, minimum, maximum float64) float64 {
|
||
|
|
return math.Max(minimum, math.Min(maximum, value))
|
||
|
|
}
|
||
|
|
|
||
|
|
func minInt(left, right int) int {
|
||
|
|
if left < right {
|
||
|
|
return left
|
||
|
|
}
|
||
|
|
return right
|
||
|
|
}
|