firedetection.py

from collections import defaultdict
import cv2
from ultralytics import YOLO

def detect_fire_in_video(
    input_video_path: str,
    output_video_path: str,
    model_path: str,
    model.to('cuda' if torch.cuda.is_available() else 'cpu')

) -> str:
    """
    Detects fire in the given video using a YOLO model.
    It draws annotations on each frame and saves the output video.

    Args:
        input_video_path (str): Path to the input video file.
        output_video_path (str): Path to save the annotated output video.
        model_path (str): Path to the YOLO .pt file.
        device (str): 'cpu', 'cuda', or 'mps' for processing.

    Returns:
        str: The path to the output annotated video.
    """

    # Tracking history - optional usage
    track_history = defaultdict(lambda: [])

    # Load the YOLO model
    model = YOLO(model_path, device=device)

    # Open the video
    cap = cv2.VideoCapture(input_video_path)

    # Retrieve video properties
    w, h, fps = (int(cap.get(prop)) for prop in [
        cv2.CAP_PROP_FRAME_WIDTH,
        cv2.CAP_PROP_FRAME_HEIGHT,
        cv2.CAP_PROP_FPS
    ])

    # Prepare output video writer
    out = cv2.VideoWriter(
        output_video_path,
        cv2.VideoWriter_fourcc(*"MJPG"),
        fps,
        (w, h)
    )

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Reached end of video or no frame retrieved.")
            break

        # Create an annotator to draw on the frame
        annotator = Annotator(frame, line_width=2)

        # Perform object tracking
        results = model.track(frame, persist=True)

        # If there are boxes with IDs and masks, annotate them
        if results[0].boxes.id is not None and results[0].masks is not None:
            masks = results[0].masks.xy
            track_ids = results[0].boxes.id.int().cpu().tolist()

            for mask, track_id in zip(masks, track_ids):
                annotator.seg_bbox(
                    mask=mask,
                    mask_color=colors(int(track_id), True),
                    label=f"ID:{track_id}"
                )

        # Write the annotated frame to output
        out.write(frame)

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

    return output_video_path