video_detection.py

from ultralytics import YOLO
import cv2
from PIL import Image
import time
import numpy as np
import uuid

model = YOLO("model/yolo11n_6-2-25.pt")

def draw_boxes(frame, results):
    for r in results:
        boxes = r.boxes

        for box in boxes:
            x1, y1, x2, y2 = box.xyxy[0]
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)

            cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 255), 3)

            cls = r.names[box.cls[0].item()]

            # object details
            org = [x1, y1]
            font = cv2.FONT_HERSHEY_SIMPLEX
            fontScale = 1
            color = (255, 0, 0)
            thickness = 2

            cv2.putText(frame, cls, org, font, fontScale, color, thickness)

    return frame
    
def video_detection(cap):
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))

    cap = cv2.VideoCapture(video)

    video_codec = cv2.VideoWriter_fourcc(*"mp4v") # type: ignore
    fps = int(cap.get(cv2.CAP_PROP_FPS))

    desired_fps = fps // SUBSAMPLE
    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2

    iterating, frame = cap.read()

    n_frames = 0

    name = f"output_{uuid.uuid4()}.mp4"
    segment_file = cv2.VideoWriter(name, video_codec, desired_fps, (width, height)) # type: ignore
    batch = []

"""
#@spaces.GPU
def video_detection(cap):
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))

    out = cv2.VideoWriter('output_video.mp4', cv2.VideoWriter_fourcc(*'h264'), fps, (frame_width, frame_height))
    
    count = 0
    while cap.isOpened():
        success, frame = cap.read()

        if not success:
            break

        #results = model(frame, stream=True, device='cuda', verbose=False)
        results = model(frame, stream=True)

        frame = draw_boxes(frame, results)

        out.write(frame)
        #if not count % 10:
        yield frame, None
        # print(count)
        count += 1

    cap.release()
    out.release()
    cv2.destroyAllWindows()

    yield None, 'output_video.mp4'
"""