import cv2
import numpy as np
import os
from ultralytics import YOLO
from natsort import natsorted

# Load the model
model = YOLO("./weights/segment_water/weights/best.pt")
seg_threshold = 0.2

# Input and output directories
input_dir = './input_frames'
output_dir = './output_vid'

# Create the output directory if it doesn't exist
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# Function to process each frame
def process_frame(frame):
    seg_results = model(frame, conf=seg_threshold, iou=0.5, verbose=False)
    pred_img = frame.copy()

    for r in seg_results:
        if r.masks is not None:
            for masks in r.masks:
                segmentations = masks.xy
                segmentations = np.array(segmentations).reshape(-1, 2)
                segmentations = segmentations.astype(np.int32)
                # Update here to fill the segmentation with the desired color
                cv2.fillPoly(pred_img, [segmentations], color=(0, 0, 255))  # Red color
    
    # Combining the original frame and the prediction with transparency
    pred_img = cv2.addWeighted(frame, 0.7, pred_img, 0.3, 0)
    return pred_img

# Function to compile frames into a video
def compile_video(frames, output_path):
    frame_height, frame_width = frames[0].shape[:2]
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, 5.0, (frame_width, frame_height))

    for frame in frames:
        out.write(frame)
    
    out.release()

# Read frames, process them, and compile into a video
frames = []
for filename in natsorted(os.listdir(input_dir)):
    if filename.endswith(".png") or filename.endswith(".jpg"):  # Check if the file is an image
        frame_path = os.path.join(input_dir, filename)
        frame = cv2.imread(frame_path)
        processed_frame = process_frame(frame)
        frames.append(processed_frame)

output_path = os.path.join(output_dir, "compiled_video.mp4")
compile_video(frames, output_path)
print("Video compilation completed.")