import cv2
import numpy as np
from ultralytics import YOLO

model = YOLO("weights/best.pt")
cap = cv2.VideoCapture("test_video.mp4")

w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

fourcc = cv2.VideoWriter_fourcc(*"avc1")
out = cv2.VideoWriter("comparison_output.mp4", fourcc, fps, (w * 2, h))

print("Generating side-by-side video...")

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    results = model(frame)

    # LEFT SIDE: Original frame
    left_side = frame

    # RIGHT SIDE: Masks on a black background
    black_bg = np.zeros_like(frame)
    # img=black_bg tells YOLO to draw segments onto the black canvas
    right_side = results[0].plot(img=black_bg, boxes=False, labels=True)

    # Combine them horizontally
    combined_frame = np.hstack((left_side, right_side))

    out.write(combined_frame)

cap.release()
out.release()
print("Done! Check comparison_output.mp4")