main.py

from pathlib import Path
from typing import cast

import cv2
import ultralytics
from ultralytics.models.yolo import YOLO

model = YOLO()


# cap = cv2.VideoCapture(0)
base = Path(__file__).resolve().parent
test_video = base / "test5.mp4"
cap = cv2.VideoCapture(str(test_video))
if not cap.isOpened():
    raise RuntimeError("Unable to open capture")


def draw_boxes(
    frame: cv2.typing.MatLike,
    boxes: ultralytics.engine.results.Boxes,
) -> None:
    xyxy = boxes.xyxy
    classes = boxes.cls
    scores = boxes.conf
    assert len(xyxy) == len(classes) == len(scores)

    for i, box in enumerate(xyxy):
        x1, y1, x2, y2 = map(int, box)
        cls = classes[i]
        score = scores[i]
        name = model.names.get(int(cls), "unknown")
        label = "cow" if name.lower() == "cow" else "not cow"
        color = (0, 255, 0) if label == "cow" else (0, 128, 255)

        txt = f"{label} {score:.2f}"
        (tw, th), _ = cv2.getTextSize(txt, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)

        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
        cv2.rectangle(frame, (x1, y1 - th - 6), (x1 + tw, y1), color, -1)
        cv2.putText(
            frame, txt, (x1, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1
        )


try:
    while True:
        ret, frame = cap.read()
        if not ret:
            cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
            continue
            # print("Failed to read frame from capture")
            # break

        results = model(frame)
        assert isinstance(results, (list, tuple))
        assert len(results) == 1

        boxes = cast(ultralytics.engine.results.Boxes, results[0].boxes)
        draw_boxes(frame, boxes)
        cv2.imshow("not-my-cows", frame)

        key = cv2.waitKey(1) & 0xFF
        if key == 27:
            break

finally:
    cap.release()
    cv2.destroyAllWindows()