from ultralytics import YOLO
import os
import numpy as np
import cv2
import time

from tqdm import tqdm

model = YOLO('./weights/segment_water/weights/best.pt')

# Change this line to point to your video file
cap = cv2.VideoCapture('pollution_cctv1.mp4')

fps = int(cap.get(cv2.CAP_PROP_FPS))

start_time = time.time()
total_fps = 0
frame_count = 0
total_elapsed_time_ms = 0

# Before entering the loop, determine the total number of frames in the video:
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

# Wrap your loop with tqdm for a progress bar
for _ in tqdm(range(total_frames), desc="Processing frames"):
    status, frame = cap.read()

    if not status:
        break
    if frame is None:
        continue

    st = time.time()
    results = model(frame, conf=0.3, iou=0.5)

    for r in results:
        if r.masks is None:
            img = frame
        else:
            mask = np.zeros((frame.shape[0], frame.shape[1], 3), dtype=np.uint8)
            for masks in r.masks:
                segmentations = masks.xy
                segmentations = np.array(segmentations).reshape(1,-1,2)
                segmentations = segmentations.astype(np.int32)
                cv2.fillPoly(mask, segmentations, color=(0, 0, 255, 127))
            img = cv2.addWeighted(frame, 1, mask, 0.5, 0)

    end_time = time.time()
    
    elapsed_time = end_time - st
    
    elapsed_time_ms = elapsed_time * 1000  # Convert to milliseconds
    total_elapsed_time_ms += elapsed_time_ms
    
    fps = 1 / elapsed_time
    total_fps += fps
    
    frame_count += 1
    
print(f"Inference time: {total_elapsed_time_ms/frame_count:.2f} ms")
print(f"FPS: {total_fps/frame_count: .2f}")