src/detect.py

from tracking.tracker import Annotation, Box


def detect_for_video(model, video_file_path, confidence):
    raw_results = model.predict(video_file_path, verbose=False, conf=confidence)
    return [
        parse_yolo_annotations_for_single_frame(raw_results_for_frame)
        for raw_results_for_frame in raw_results
    ]


def parse_yolo_annotations_for_single_frame(raw_results):
    return [
        Annotation(Box(*box.xyxy[0].tolist()), int(box.cls[0]), float(box.conf[0]))
        for box in raw_results.boxes
    ]