ozai-03/drone-detection-api

1import os2import time3import tempfile4import cv25import numpy as np6from ultralytics import YOLO7from huggingface_hub import hf_hub_download8 9MODEL_REPO = os.environ.get("MODEL_REPO", "ozai-03/yolov8m-drone-detection")10MODEL_VERSION = os.environ.get("MODEL_VERSION", "v1.0")11 12_model = None13 14 15def load_model():16    global _model17    if _model is not None:18        return _model19    model_path = hf_hub_download(20        repo_id=MODEL_REPO,21        filename="best.pt",22        repo_type="model",23    )24    _model = YOLO(model_path)25    return _model26 27 28try:29    load_model()30except Exception as e:31    print(f"[WARNING] Model failed to load at startup: {e}")32 33 34def run_inference(video_path: str, conf: float = 0.25, max_frames: int = 120, progress_callback=None, annotate: bool = False) -> dict:35    model = load_model()36    start_time = time.time()37 38    cap = cv2.VideoCapture(video_path)39    if not cap.isOpened():40        raise ValueError("Unable to open video file")41 42    fps = cap.get(cv2.CAP_PROP_FPS) or 30.043    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))44    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))45    if width == 0 or height == 0:46        cap.release()47        raise ValueError("Unable to read video dimensions — file may not be a valid video")48 49    # int(fps) can be 0 for sub-1fps or malformed files50    frame_interval = max(1, int(fps))51 52    out_path = None53    writer = None54    if annotate:55        out_fd, out_path = tempfile.mkstemp(suffix="_annotated.mp4")56        os.close(out_fd)57        output_fps = min(fps, 30.0)58        writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), output_fps, (width, height))59 60    frames_result = []61    # seen_ids counts unique tracked objects per class across the whole video62    seen_ids: dict[str, set] = {}63    confidence_sum = 0.064    detection_count = 065    frame_idx = 066    sampled = 067    # hard cap: process at most max_frames seconds of video68    max_total_frames = max_frames * frame_interval69 70    while frame_idx < max_total_frames:71        ret, frame = cap.read()72        if not ret:73            break74 75        # Run tracker on every frame to maintain Kalman filter continuity76        results = model.track(frame, conf=conf, tracker="bytetrack.yaml", verbose=False, persist=True)77 78        # Write every annotated frame so output video plays at source FPS79        if writer is not None:80            writer.write(results[0].plot())81 82        # Collect results only at 1-second sample points (keeps API response size manageable)83        if frame_idx % frame_interval == 0:84            timestamp_ms = int((frame_idx / fps) * 1000)85            detections = []86 87            for box in results[0].boxes:88                cls_id = int(box.cls[0])89                cls_name = model.names[cls_id]90                conf_val = float(box.conf[0])91                track_id = int(box.id[0]) if box.id is not None else "N/A"92                x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()]93                print(f"[TRACK] frame={sampled} ts={timestamp_ms}ms | id={track_id} cls={cls_name} conf={conf_val:.2f}")94 95                detections.append({96                    "class": cls_name,97                    "confidence": round(conf_val, 4),98                    "bbox": [x1, y1, x2, y2],99                    "track_id": track_id,100                })101 102                if track_id != "N/A":103                    seen_ids.setdefault(cls_name, set()).add(track_id)104                confidence_sum += conf_val105                detection_count += 1106 107            frames_result.append({108                "frame_id": sampled,109                "timestamp_ms": timestamp_ms,110                "detections": detections,111            })112 113            sampled += 1114            if progress_callback:115                progress_callback(sampled, max_frames)116 117        frame_idx += 1118 119    cap.release()120    if writer is not None:121        writer.release()122 123    # Reset tracker so IDs from this video don't carry into the next request124    if hasattr(model, "predictor") and model.predictor is not None:125        if hasattr(model.predictor, "trackers") and model.predictor.trackers:126            model.predictor.trackers[0].reset()127 128    # Count unique objects per class, fall back to detection count if tracker had no IDs129    summary = {cls: len(ids) for cls, ids in seen_ids.items()} if seen_ids else {}130 131    elapsed_ms = int((time.time() - start_time) * 1000)132    avg_confidence = round(confidence_sum / detection_count, 4) if detection_count > 0 else 0.0133 134    return {135        "model_version": MODEL_VERSION,136        "inference_time_ms": elapsed_ms,137        "total_frames_processed": sampled,138        "frames": frames_result,139        "summary": summary,140        "avg_confidence": avg_confidence,141        "output_video_path": out_path,142    }143