silver_pose/v2/internal/monitor/monitor.go

// Package monitor composes source, pose, fall and evidence work off the UI thread.
package monitor

import (
	"context"
	"fmt"
	"image"
	"sync"
	"time"

	"silverpose/v2/internal/alert"
	"silverpose/v2/internal/fall"
	"silverpose/v2/internal/pose"
	"silverpose/v2/internal/render"
	"silverpose/v2/internal/source"
)

type FrameSource interface {
	Frames() <-chan source.Frame
	Statuses() <-chan source.Update
}

type PoseRuntime interface {
	Run([]float32) ([]float32, error)
	Close()
}

type EventProcessor interface {
	Process(fall.Frame) fall.FrameResult
}

type Update struct {
	Image         *image.RGBA
	Result        fall.FrameResult
	Events        []alert.Record
	SourceStatus  source.Status
	SourceMessage string
	Sequence      uint64
}

type Monitor struct {
	source     FrameSource
	runtime    PoseRuntime
	processor  EventProcessor
	confidence float32
	dispatcher *alert.Dispatcher
	clock      func() time.Time
	updates    chan Update
	alerts     chan alert.Record
	done       chan struct{}
	closeOnce  sync.Once
}

func New(
	stream FrameSource,
	runtime PoseRuntime,
	processor EventProcessor,
	confidence float32,
	dispatcher *alert.Dispatcher,
	clock func() time.Time,
) (*Monitor, error) {
	if stream == nil || runtime == nil || processor == nil || dispatcher == nil {
		return nil, fmt.Errorf("monitor dependencies must be non-nil")
	}
	if confidence < 0 || confidence > 1 {
		return nil, fmt.Errorf("pose confidence must be between 0 and 1")
	}
	if clock == nil {
		clock = time.Now
	}
	return &Monitor{
		source: stream, runtime: runtime, processor: processor, confidence: confidence,
		dispatcher: dispatcher, clock: clock, updates: make(chan Update, 1), alerts: make(chan alert.Record, 16), done: make(chan struct{}),
	}, nil
}

func (monitor *Monitor) Updates() <-chan Update      { return monitor.updates }
func (monitor *Monitor) Done() <-chan struct{}       { return monitor.done }
func (monitor *Monitor) Alerts() <-chan alert.Record { return monitor.alerts }

// Run blocks until the source closes or the supplied context is cancelled.
// The caller owns the goroutine; all expensive work stays out of Walk's UI thread.
func (monitor *Monitor) Run(ctx context.Context) {
	defer monitor.close()
	defer monitor.runtime.Close()
	frames := monitor.source.Frames()
	statuses := monitor.source.Statuses()
	currentStatus := source.Connecting
	currentMessage := "正在连接视频流"
	for frames != nil || statuses != nil {
		select {
		case <-ctx.Done():
			return
		case status, open := <-statuses:
			if !open {
				statuses = nil
				continue
			}
			currentStatus, currentMessage = status.Status, status.Message
			monitor.publish(Update{SourceStatus: currentStatus, SourceMessage: currentMessage})
		case frame, open := <-frames:
			if !open {
				frames = nil
				continue
			}
			statuses, currentStatus, currentMessage = monitor.drainStatuses(statuses, currentStatus, currentMessage)
			monitor.processFrame(frame, currentStatus, currentMessage)
		}
	}
}

func (monitor *Monitor) drainStatuses(statuses <-chan source.Update, current source.Status, message string) (<-chan source.Update, source.Status, string) {
	for statuses != nil {
		select {
		case status, open := <-statuses:
			if !open {
				return nil, current, message
			}
			current, message = status.Status, status.Message
		default:
			return statuses, current, message
		}
	}
	return nil, current, message
}

func (monitor *Monitor) processFrame(frame source.Frame, status source.Status, message string) {
	input, transform, err := pose.PreprocessBGR(frame.BGR, frame.Width, frame.Height, 640)
	if err != nil {
		monitor.publish(Update{SourceStatus: status, SourceMessage: "视频帧格式无效", Sequence: frame.Sequence})
		return
	}
	output, err := monitor.runtime.Run(input)
	if err != nil {
		monitor.publish(Update{SourceStatus: status, SourceMessage: "姿态推理失败", Sequence: frame.Sequence})
		return
	}
	people, err := pose.ParseYOLOv8Pose(output, transform, monitor.confidence, 0.70)
	if err != nil {
		monitor.publish(Update{SourceStatus: status, SourceMessage: "姿态结果无效", Sequence: frame.Sequence})
		return
	}
	result := monitor.processor.Process(fall.Frame{
		Timestamp: frame.Timestamp.Seconds(), Width: frame.Width, Height: frame.Height, Poses: people,
	})
	capturedAt := monitor.clock().UTC()
	canvas, err := render.Render(frame.BGR, frame.Width, frame.Height, result, capturedAt)
	if err != nil {
		monitor.publish(Update{Result: result, SourceStatus: status, SourceMessage: "画面渲染失败", Sequence: frame.Sequence})
		return
	}
	records := make([]alert.Record, 0, len(result.Events))
	for _, event := range result.Events {
		record, err := monitor.dispatcher.Dispatch(event, canvas, capturedAt)
		if err != nil {
			monitor.publish(Update{Image: canvas, Result: result, SourceStatus: status, SourceMessage: "事件证据保存失败", Sequence: frame.Sequence})
			return
		}
		if record.Written {
			records = append(records, record)
			monitor.alerts <- record
		}
	}
	monitor.publish(Update{
		Image: canvas, Result: result, Events: records, SourceStatus: status, SourceMessage: message, Sequence: frame.Sequence,
	})
}

func (monitor *Monitor) publish(update Update) {
	select {
	case monitor.updates <- update:
		return
	default:
	}
	select {
	case <-monitor.updates:
	default:
	}
	select {
	case monitor.updates <- update:
	default:
	}
}

func (monitor *Monitor) close() {
	monitor.closeOnce.Do(func() {
		close(monitor.updates)
		close(monitor.alerts)
		close(monitor.done)
	})
}