Examples/render_numbers.py

"""Draw the number model's example video: every player's box, the jersey digits read this frame,
and a "#63 Z. Ostapchuk" tag on each player whose number has been read.

    python CV_Models/Examples/render_numbers.py <clip>.mp4 <run_models output>/detections.json out.mp4

The detections must come from run_models.py with the number model (and --teams for names).
"""

import json
import sys
from pathlib import Path

import cv2

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))  # the repo root, for video_io
from video_io import FrameReader  # noqa: E402

# BGR per player model class: San Jose teal, Montreal red, referees grey.
TEAM_COLORS = {
    "team_a_player": (140, 120, 0), "goalie_a": (140, 120, 0),
    "team_b_player": (40, 30, 200), "goalie_b": (40, 30, 200),
    "referee": (130, 130, 130),
}
DIGIT_COLOR = (0, 215, 255)


def short_name(name):
    """'Zack Ostapchuk' -> 'Z. Ostapchuk'."""
    first, _, last = name.partition(" ")
    return f"{first[0]}. {last}" if last else name


def draw_tag(img, text, box, color):
    x1, y1, x2, _ = box
    scale, thickness = img.shape[1] / 2400, max(1, img.shape[1] // 960)
    (tw, th), base = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, scale, thickness)
    x = int(max(0, min(img.shape[1] - tw - 8, (x1 + x2) / 2 - tw / 2)))
    y = int(max(th + base + 8, y1 - 6))
    cv2.rectangle(img, (x - 4, y - th - base - 4), (x + tw + 4, y + 2), color, -1)
    cv2.rectangle(img, (x - 4, y - th - base - 4), (x + tw + 4, y + 2), (255, 255, 255), 1)
    cv2.putText(img, text, (x, y - base), cv2.FONT_HERSHEY_SIMPLEX, scale, (255, 255, 255), thickness,
                cv2.LINE_AA)


def draw(frame, det):
    line = max(1, frame.shape[1] // 800)
    for d in det["number"]:
        x1, y1, x2, y2 = map(int, d["box_xyxy"])
        cv2.rectangle(frame, (x1, y1), (x2, y2), DIGIT_COLOR, line)
    for p in det["player"]:
        color = TEAM_COLORS.get(p["class"], (200, 200, 200))
        x1, y1, x2, y2 = map(int, p["box_xyxy"])
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, line)
        if p.get("jersey_number"):
            name = p.get("player_name")
            draw_tag(frame, f"#{p['jersey_number']}" + (f" {short_name(name)}" if name else ""),
                     p["box_xyxy"], color)
    return frame


def main():
    video, detections, out = sys.argv[1:4]
    frames = json.loads(open(detections).read())
    reader = FrameReader(video, len(frames))
    writer = cv2.VideoWriter(out, cv2.VideoWriter_fourcc(*"mp4v"), reader.fps, (reader.width, reader.height))
    for frame, f in zip(reader, frames):
        writer.write(draw(frame, f["detections"]))
    writer.release()
    reader.close()


if __name__ == "__main__":
    main()