silver_pose/v2/internal/render/overlay.go

// Package render converts video frames into annotated display/evidence images.
package render

import (
	"fmt"
	"image"
	"image/color"
	"math"
	"time"

	"golang.org/x/image/font"
	"golang.org/x/image/font/basicfont"
	"golang.org/x/image/math/fixed"

	"silverpose/v2/internal/fall"
)

var (
	normalColor   = color.RGBA{R: 21, G: 128, B: 61, A: 255}
	cautionColor  = color.RGBA{R: 180, G: 83, B: 9, A: 255}
	criticalColor = color.RGBA{R: 198, G: 40, B: 40, A: 255}
)

var skeletonEdges = [][2]int{
	{0, 1}, {0, 2}, {1, 3}, {2, 4}, {5, 6}, {5, 7}, {7, 9}, {6, 8}, {8, 10},
	{5, 11}, {6, 12}, {11, 12}, {11, 13}, {13, 15}, {12, 14}, {14, 16},
}

// Render returns a stand-alone RGBA frame. The renderer accepts no source URL
// or credential and can therefore be used both for the UI and screenshot evidence.
func Render(bgr []byte, width, height int, result fall.FrameResult, capturedAt time.Time) (*image.RGBA, error) {
	if width <= 0 || height <= 0 || len(bgr) != width*height*3 {
		return nil, fmt.Errorf("BGR frame length = %d, want %d", len(bgr), width*height*3)
	}
	canvas := image.NewRGBA(image.Rect(0, 0, width, height))
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			source := (y*width + x) * 3
			destination := canvas.PixOffset(x, y)
			canvas.Pix[destination] = bgr[source+2]
			canvas.Pix[destination+1] = bgr[source+1]
			canvas.Pix[destination+2] = bgr[source]
			canvas.Pix[destination+3] = 255
		}
	}
	confirmed := false
	for _, person := range result.People {
		personColor := colorForState(person.State)
		if person.State == fall.Confirmed {
			confirmed = true
		}
		drawPerson(canvas, person, personColor, capturedAt)
	}
	if confirmed {
		strokeRectangle(canvas, 0, 0, width-1, height-1, criticalColor, 6)
	}
	return canvas, nil
}

func drawPerson(canvas *image.RGBA, person fall.PersonAnalysis, lineColor color.RGBA, capturedAt time.Time) {
	pose := person.TrackedPose.Pose
	box := pose.Box
	strokeRectangle(canvas, round(box.Left), round(box.Top), round(box.Right), round(box.Bottom), lineColor, 2)
	for _, edge := range skeletonEdges {
		first := pose.Keypoints[edge[0]]
		second := pose.Keypoints[edge[1]]
		if first.Confidence > 0 && second.Confidence > 0 {
			drawLine(canvas, round(first.X), round(first.Y), round(second.X), round(second.Y), lineColor)
		}
	}
	for _, point := range pose.Keypoints {
		if point.Confidence > 0 {
			drawDot(canvas, round(point.X), round(point.Y), lineColor)
		}
	}
	label := fmt.Sprintf("%s %s", person.TrackedPose.TrackID, person.State)
	if !capturedAt.IsZero() {
		label += " " + capturedAt.UTC().Format("15:04:05")
	}
	drawText(canvas, round(box.Left), max(12, round(box.Top)-4), label, lineColor)
}

func colorForState(state fall.State) color.RGBA {
	switch state {
	case fall.Confirmed:
		return criticalColor
	case fall.Suspect, fall.Recovering:
		return cautionColor
	default:
		return normalColor
	}
}

func drawText(canvas *image.RGBA, x, y int, text string, textColor color.RGBA) {
	drawer := &font.Drawer{
		Dst:  canvas,
		Src:  image.NewUniform(textColor),
		Face: basicfont.Face7x13,
		Dot:  fixed.P(x, y),
	}
	drawer.DrawString(text)
}

func strokeRectangle(canvas *image.RGBA, left, top, right, bottom int, lineColor color.RGBA, thickness int) {
	if thickness < 1 {
		return
	}
	for offset := 0; offset < thickness; offset++ {
		drawLine(canvas, left+offset, top+offset, right-offset, top+offset, lineColor)
		drawLine(canvas, left+offset, bottom-offset, right-offset, bottom-offset, lineColor)
		drawLine(canvas, left+offset, top+offset, left+offset, bottom-offset, lineColor)
		drawLine(canvas, right-offset, top+offset, right-offset, bottom-offset, lineColor)
	}
}

func drawLine(canvas *image.RGBA, x0, y0, x1, y1 int, lineColor color.RGBA) {
	dx := abs(x1 - x0)
	sx := -1
	if x0 < x1 {
		sx = 1
	}
	dy := -abs(y1 - y0)
	sy := -1
	if y0 < y1 {
		sy = 1
	}
	err := dx + dy
	for {
		setPixel(canvas, x0, y0, lineColor)
		if x0 == x1 && y0 == y1 {
			return
		}
		twiceError := 2 * err
		if twiceError >= dy {
			err += dy
			x0 += sx
		}
		if twiceError <= dx {
			err += dx
			y0 += sy
		}
	}
}

func drawDot(canvas *image.RGBA, x, y int, dotColor color.RGBA) {
	for vertical := -2; vertical <= 2; vertical++ {
		for horizontal := -2; horizontal <= 2; horizontal++ {
			if horizontal*horizontal+vertical*vertical <= 4 {
				setPixel(canvas, x+horizontal, y+vertical, dotColor)
			}
		}
	}
}

func setPixel(canvas *image.RGBA, x, y int, pixel color.RGBA) {
	if !image.Pt(x, y).In(canvas.Bounds()) {
		return
	}
	canvas.SetRGBA(x, y, pixel)
}

func round(value float32) int { return int(math.Round(float64(value))) }
func abs(value int) int {
	if value < 0 {
		return -value
	}
	return value
}
func max(left, right int) int {
	if left > right {
		return left
	}
	return right
}