import gradio as gr
import cv2
import os
import tempfile
import shutil

from ultralytics import YOLO

# Initialize YOLO model
model_path = r"best.pt"
if not os.path.exists(model_path):
    raise FileNotFoundError(f"Model file not found at {model_path}")
model = YOLO(model_path)

# Define the maximum file size (70MB)
MAX_FILE_SIZE_MB = 70  # Limit is set to 70MB
MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_MB * 1024 * 1024  # Convert MB to bytes

def process_video(video_file):
    # Check the size of the uploaded video
    video_size = os.path.getsize(video_file)
    if video_size > MAX_FILE_SIZE_BYTES:
        return f"Error: The uploaded video exceeds the {MAX_FILE_SIZE_MB}MB size limit. Please upload a smaller video.", []

    print("Video size is ", video_size/1024)
    # Create a temporary folder to save extracted frames
    temp_dir = tempfile.mkdtemp()

    # Capture the video using OpenCV
    video = cv2.VideoCapture(video_file)
    if not video.isOpened():
        return f"Error: Unable to open video file {video_file}", []
    
    fps = int(video.get(cv2.CAP_PROP_FPS))
    if fps == 0:
        return "Error: Video FPS is 0 or not detected.", []
    
    frame_rate = max(1, fps // 5)  # Extract 5 frames per second, ensure at least 1 frame/sec

    frame_count = 0
    extracted_frames = []

    try:
        while True:
            ret, frame = video.read()
            if not ret:
                break

            # Extract every nth frame according to the frame_rate
            if frame_count % frame_rate == 0:
                # Resize the frame if it's larger than 1280x720
                h, w, _ = frame.shape
                if w > 1280 or h > 720:
                    frame = cv2.resize(frame, (1280, 720))

                # Save the frame to the temporary directory
                frame_path = os.path.join(temp_dir, f"frame_{frame_count}.jpg")
                cv2.imwrite(frame_path, frame)

                # Perform detection using the YOLO model
                results = model.predict(source=frame_path, conf=0.3, save=False, show=False, 
                                        line_width=2, show_labels=True, show_boxes=True, retina_masks=True)
                
                for result in results:
                    boxes = result.boxes
                    if boxes is not None and len(boxes) > 0:
                        # Save the result image if detections are found
                        result.save(filename=frame_path)
                    else:
                        # If no detections, delete the frame
                        if os.path.exists(frame_path):
                            os.remove(frame_path)

                # Add the frame path to the list of extracted frames if it still exists
                if os.path.exists(frame_path):
                    extracted_frames.append(frame_path)

            frame_count += 1
    finally:
        video.release()
        # Clean up temp directory after processing
        if len(extracted_frames) == 0:
            shutil.rmtree(temp_dir)

    # Return the list of extracted frame paths for display
    if extracted_frames:
        return f"Success: {len(extracted_frames)} frames extracted and processed.", extracted_frames
    else:
        return "No damage found on the road.", []

# Define the Gradio interface
iface = gr.Interface(
    fn=process_video,  # Function to process the video and return extracted frames
    inputs=gr.Video(label="Upload your video"),  # Video upload input
    outputs=[
        gr.Textbox(label="Status"),  # Display status or error message
        gr.Gallery(label="Output Images Results")  # Display extracted frames in a gallery
    ],
    title="Road Damage Detection",  # Interface title
    description="Upload a video with a maximum size of 70MB."
)

# Launch the Gradio app
iface.launch()