mediaproc/ui/detection-app/src/cv/tracks.ts

/**
 * Motion track matching and rendering.
 *
 * Matches bounding boxes across frames by IoU (Intersection over Union).
 * Uses Hungarian-style optimal assignment so every track gets its best match.
 * Renders VFX-style tracks: × at start/end, dotted trace in between.
 * Returns ImageData (RGBA, transparent background) for use as a FrameOverlay.
 */

export type TrackPoint = { seq: number; cx: number; cy: number; w: number; h: number }
export type Track = { id: number; points: TrackPoint[] }

type Box = { x: number; y: number; w: number; h: number }

function iou(a: Box, b: Box): number {
  const x1 = Math.max(a.x, b.x)
  const y1 = Math.max(a.y, b.y)
  const x2 = Math.min(a.x + a.w, b.x + b.w)
  const y2 = Math.min(a.y + a.h, b.y + b.h)
  const inter = Math.max(0, x2 - x1) * Math.max(0, y2 - y1)
  if (inter === 0) return 0
  const union = a.w * a.h + b.w * b.h - inter
  return union > 0 ? inter / union : 0
}

/**
 * Match bounding boxes across frames by IoU.
 *
 * For each pair of consecutive frames, computes an IoU cost matrix and
 * assigns tracks to boxes using greedy-best-first on IoU score (descending).
 * minIoU: minimum overlap to consider a match (default 0.15).
 */
export function matchTracks(
  regionsByFrame: Record<number, Box[]>,
  minIoU = 0.15,
): Track[] {
  const seqs = Object.keys(regionsByFrame).map(Number).sort((a, b) => a - b)
  if (seqs.length === 0) return []

  let nextId = 0
  const active: Track[] = []
  const finished: Track[] = []

  for (const seq of seqs) {
    const boxes = regionsByFrame[seq]

    if (active.length === 0) {
      for (const b of boxes) {
        active.push({
          id: nextId++,
          points: [{ seq, cx: b.x + b.w / 2, cy: b.y + b.h / 2, w: b.w, h: b.h }],
        })
      }
      continue
    }

    // Build cost matrix: IoU between each active track's last box and each new box
    const costs: Array<{ trackIdx: number; boxIdx: number; score: number }> = []
    for (let t = 0; t < active.length; t++) {
      const last = active[t].points[active[t].points.length - 1]
      const trackBox: Box = { x: last.cx - last.w / 2, y: last.cy - last.h / 2, w: last.w, h: last.h }
      for (let b = 0; b < boxes.length; b++) {
        const score = iou(trackBox, boxes[b])
        if (score >= minIoU) {
          costs.push({ trackIdx: t, boxIdx: b, score })
        }
      }
    }

    // Greedy-best-first assignment (sorted by descending IoU)
    costs.sort((a, b) => b.score - a.score)
    const matchedTracks = new Set<number>()
    const matchedBoxes = new Set<number>()
    const next: Track[] = []

    for (const { trackIdx, boxIdx } of costs) {
      if (matchedTracks.has(trackIdx) || matchedBoxes.has(boxIdx)) continue
      matchedTracks.add(trackIdx)
      matchedBoxes.add(boxIdx)
      const b = boxes[boxIdx]
      active[trackIdx].points.push({
        seq, cx: b.x + b.w / 2, cy: b.y + b.h / 2, w: b.w, h: b.h,
      })
      next.push(active[trackIdx])
    }

    // Unmatched tracks → finished
    for (let t = 0; t < active.length; t++) {
      if (!matchedTracks.has(t)) finished.push(active[t])
    }

    // Unmatched boxes → new tracks
    for (let b = 0; b < boxes.length; b++) {
      if (!matchedBoxes.has(b)) {
        const box = boxes[b]
        next.push({
          id: nextId++,
          points: [{ seq, cx: box.x + box.w / 2, cy: box.y + box.h / 2, w: box.w, h: box.h }],
        })
      }
    }

    active.length = 0
    active.push(...next)
  }

  return [...finished, ...active]
}

/**
 * Render tracks to ImageData (RGBA, transparent background).
 *
 * Only draws tracks that span ≥2 frames (single-frame detections are skipped).
 * × at start, × at end, dotted trace in between.
 * Filled dot at currentSeq.
 */
export function renderTracksToImageData(
  tracks: Track[],
  width: number,
  height: number,
  currentSeq: number,
  color = '#00e5ff',
): ImageData {
  const canvas = new OffscreenCanvas(width, height)
  const ctx = canvas.getContext('2d')!

  ctx.clearRect(0, 0, width, height)
  ctx.strokeStyle = color
  ctx.fillStyle = color

  for (const track of tracks) {
    const pts = track.points
    // Skip single-frame detections — no temporal signal
    if (pts.length < 2) continue

    // Dotted trace connecting all centers
    ctx.globalAlpha = 0.75
    ctx.lineWidth = 1.5
    ctx.setLineDash([4, 4])
    ctx.beginPath()
    ctx.moveTo(pts[0].cx, pts[0].cy)
    for (let i = 1; i < pts.length; i++) {
      ctx.lineTo(pts[i].cx, pts[i].cy)
    }
    ctx.stroke()
    ctx.setLineDash([])

    // × at first point
    ctx.globalAlpha = 0.9
    ctx.lineWidth = 1.5
    drawX(ctx, pts[0].cx, pts[0].cy, 7)

    // × at last point
    drawX(ctx, pts[pts.length - 1].cx, pts[pts.length - 1].cy, 7)

    // Filled dot at current frame
    const curr = pts.find(p => p.seq === currentSeq)
    if (curr) {
      ctx.globalAlpha = 1
      ctx.beginPath()
      ctx.arc(curr.cx, curr.cy, 4, 0, Math.PI * 2)
      ctx.fill()
    }
  }

  return ctx.getImageData(0, 0, width, height)
}

function drawX(
  ctx: OffscreenCanvasRenderingContext2D,
  cx: number,
  cy: number,
  size: number,
): void {
  ctx.beginPath()
  ctx.moveTo(cx - size, cy - size)
  ctx.lineTo(cx + size, cy + size)
  ctx.moveTo(cx + size, cy - size)
  ctx.lineTo(cx - size, cy + size)
  ctx.stroke()
}