| 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") |