File size: 2,727 Bytes
4a01c1b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """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()
|