import streamlit as st
from ultralytics import YOLO
from PIL import Image
import numpy as np
import cv2
import os
import time

# Function to process video and save annotated results
def process_video(video_in_filepath, video_out_filepath, model):
    video_reader = cv2.VideoCapture(video_in_filepath)

    # Get video properties
    nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
    fps = video_reader.get(cv2.CAP_PROP_FPS)

    # Set up video writer
    video_writer = cv2.VideoWriter(
        video_out_filepath,
        cv2.VideoWriter_fourcc(*'mp4v'),
        fps,
        (frame_w, frame_h)
    )

    progress_bar = st.progress(0)
    for frame_idx in range(nb_frames):
        success, frame = video_reader.read()
        if not success:
            break

        # YOLO inference
        results = model(frame)
        annotated_frame = results[0].plot()

        # Write annotated frame to output video
        video_writer.write(annotated_frame)

        # Update progress bar
        progress_bar.progress((frame_idx + 1) / nb_frames)

    video_reader.release()
    video_writer.release()
    progress_bar.empty()  # Remove the progress bar after completion

# Streamlit App Layout
st.title("ITI107 Assignment: Taxi & License Plate Detection")
uploaded_file = st.file_uploader("Image or video", type=["jpg", "jpeg", "png", "mp4"])

# Create two columns
col1, col2 = st.columns(2)

# Left Column: Upload functionality
with col1:

    if uploaded_file:
        if uploaded_file.type.startswith("image"):
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Image", use_container_width=True)

            # YOLO inference on image
            model = YOLO("best.pt")
            results = model.predict(np.array(image))

            # Annotate and display results
            result_image = Image.fromarray(results[0].plot())
            with col2:  # Show results in the right column
                st.image(result_image, caption="Detection Results", use_container_width=True)

        elif uploaded_file.type == "video/mp4":
            # Save uploaded video locally
            input_video_path = os.path.join(os.getcwd(), "uploaded_video.mp4")
            with open(input_video_path, "wb") as f:
                f.write(uploaded_file.read())
            st.video(input_video_path)  # Display uploaded video

            # Process video
            output_video_path = os.path.join(os.getcwd(), f"processed_video_{int(time.time())}.mp4")
            model = YOLO("best.pt")
            st.write("Processing video, please wait...")  # Status message
            process_video(input_video_path, output_video_path, model)

            # Display output video
            with col2:  # Show results in the right column
                st.video(output_video_path)  # Display processed video
                st.write("Video processing complete!")  # Completion message