74 lines
2.1 KiB
Go
74 lines
2.1 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"time"
|
|
|
|
"silverpose/v2/internal/pose"
|
|
"silverpose/v2/internal/spike"
|
|
)
|
|
|
|
func main() {
|
|
ffmpegPath := flag.String("ffmpeg", "ffmpeg", "path to ffmpeg.exe")
|
|
videoPath := flag.String("video", "..\\demo\\1.mp4", "local non-sensitive validation video")
|
|
width := flag.Int("width", 848, "source video frame width")
|
|
height := flag.Int("height", 480, "source video frame height")
|
|
modelPath := flag.String("onnx", "assets\\best.onnx", "locked ONNX pose model path")
|
|
runtimeDLLPath := flag.String("ort-dll", "", "absolute onnxruntime.dll path")
|
|
flag.Parse()
|
|
|
|
if *runtimeDLLPath == "" {
|
|
fail("--ort-dll is required")
|
|
}
|
|
if *width <= 0 || *height <= 0 {
|
|
fail("--width and --height must be positive")
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
raw, err := exec.CommandContext(ctx, *ffmpegPath,
|
|
"-v", "error", "-i", *videoPath,
|
|
"-frames:v", "1", "-f", "rawvideo", "-pix_fmt", "bgr24", "-",
|
|
).Output()
|
|
if err != nil {
|
|
fail("decode first BGR frame: %v", err)
|
|
}
|
|
expectedBytes := *width * *height * 3
|
|
if len(raw) != expectedBytes {
|
|
fail("decoded BGR frame has %d bytes, want %d; verify --width/--height", len(raw), expectedBytes)
|
|
}
|
|
|
|
input, transform, err := pose.PreprocessBGR(raw, *width, *height, 640)
|
|
if err != nil {
|
|
fail("preprocess frame: %v", err)
|
|
}
|
|
output, err := spike.RunPose(*modelPath, *runtimeDLLPath, input)
|
|
if err != nil {
|
|
fail("run ONNX Pose: %v", err)
|
|
}
|
|
people, err := pose.ParseYOLOv8Pose(output, transform, 0.25, 0.70)
|
|
if err != nil {
|
|
fail("parse ONNX Pose: %v", err)
|
|
}
|
|
|
|
max := float32(0)
|
|
for _, value := range output {
|
|
if value > max {
|
|
max = value
|
|
}
|
|
}
|
|
fmt.Printf("SPIKE_OK frame=%dx%d input=%d output=%d people=%d max=%.4f\n", *width, *height, len(input), len(output), len(people), max)
|
|
for index, person := range people {
|
|
fmt.Printf("PERSON index=%d conf=%.6f box=%.3f,%.3f,%.3f,%.3f\n", index, person.Confidence, person.Box.Left, person.Box.Top, person.Box.Right, person.Box.Bottom)
|
|
}
|
|
}
|
|
|
|
func fail(format string, args ...any) {
|
|
fmt.Fprintf(os.Stderr, "spike: "+format+"\n", args...)
|
|
os.Exit(1)
|
|
}
|