import streamlit as st
import cv2
import tempfile
from ultralytics import YOLO
import numpy as np
import time

alerting_classes = {
    0: 'People',
    2: 'Car',
    7: 'Truck',
    24: 'Backpack',
    65: 'Suspicious handheld device',
    26: 'Handbag',
    28: 'Suitcase',
}

red_tint = np.array([[[0, 0, 255]]], dtype=np.uint8)

model1 = YOLO('yolov8n.pt')

# Set the theme to light mode
# st.set_theme("light")

# Set page config
st.set_page_config(page_title="Object Detection App", page_icon="🚗")

st.title("Object Detection and Recognition")

st.write("""
This web app performs object detection and recognition on a video using YOLOv8. 
It detects various objects, such as people, cars, trucks, backpacks, suspicious handheld devices, handbags, and suitcases. 
The processed video is displayed with alerts highlighted, and you can stop the inference at any time.
""")

video_file = st.file_uploader("Choose a video file", type=["mp4"])

if video_file is not None:
    # Create temporary file for uploaded video
    tfile = tempfile.NamedTemporaryFile(delete=False)
    tfile.write(video_file.read())

    # Open video capture using temporary file path
    cap = cv2.VideoCapture(tfile.name)
    alert_set = set(alerting_classes.keys())
    alert_set.remove(0)

    # Create red-tinted overlay
    red_tinted_overlay = np.tile(red_tint, (1, 1, 1))

    stop_button = st.button("Stop Inference")
    processing_interrupted = False

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    progress_bar_processing_slot = st.empty()

    # Collect frames in a list
    frames = []

    # Set the desired size as a percentage of the original size
    target_size_percentage = 0.65

    while cap.isOpened() and not processing_interrupted:
        alert_flag = False
        alert_reason = []

        success, frame = cap.read()

        # if the frame is read correctly ret is True
        if not success:
            # st.warning("Can't receive frame (stream end?). Exiting ...")
            break

        # Calculate the target size
        target_width = int(frame.shape[1] * target_size_percentage)
        target_height = int(frame.shape[0] * target_size_percentage)

        # Resize the frame
        frame = cv2.resize(frame, (target_width, target_height))

        if success:
            # Check if the stop button is clicked
            if stop_button:
                processing_interrupted = True
                break

            # Perform YOLO object detection
            results = model1(frame, conf=0.35, verbose=False, classes=list(alerting_classes.keys()))

            class_ids = results[0].boxes.cls.tolist()
            class_counts = {cls: class_ids.count(cls) for cls in set(class_ids)}

            for cls in alert_set:
                if cls in class_counts and class_counts[cls] > 0:
                    alert_flag = True
                    alert_reason.append((cls, class_counts[cls]))

            if class_counts.get(0, 0) > 5:
                alert_flag = True
                alert_reason.append((0, class_counts[0]))

            text = 'ALERT!'
            font = cv2.FONT_HERSHEY_DUPLEX
            font_scale = 0.75
            thickness = 2

            size = cv2.getTextSize(text, font, font_scale, thickness)
            x = 0
            y = int((2 + size[0][1]))

            img = results[0].plot()
            if alert_flag:
                # Resize the red-tinted overlay to match the image size
                red_tinted_overlay = cv2.resize(red_tinted_overlay, (img.shape[1], img.shape[0]))
                img = cv2.addWeighted(img, 0.7, red_tinted_overlay, 0.3, 0)
                cv2.putText(img, text, (x, y), font, font_scale, (0, 0, 0), thickness)

                y += int(size[0][1]) + 10  # Move to the next line

                for cls, count in alert_reason:
                    alert_text = f'{count} {alerting_classes[cls]}'
                    cv2.putText(img, alert_text, (x, y), font, font_scale, (0, 0, 0), thickness)
                    y += int(size[0][1]) + 10  # Move to the next line

            # Append the frame to the list
            frames.append(img)

            del results

            # Update processing progress bar
            current_frame_processing = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
            progress = current_frame_processing / total_frames
            progress_bar_processing_slot.progress(progress)
            progress_bar_processing_slot.text(f"Processing... {int(progress * 100)}%")

    # Release resources
    cap.release()
    tfile.close()

    # Display frames one by one as a video with 24 FPS
    if processing_interrupted:
        st.text("User interrupted processing.")
    else:
        progress_bar_processing_slot.text("Done!")
        video_placeholder = st.image([])
        progress_bar_display = st.progress(0)
        fps_delay = 1 / 24  # Delay to achieve 24 FPS

        for i, frame in enumerate(frames):
            video_placeholder.image(frame, channels="BGR", caption="YOLOv8 Inference")
            # Update display progress bar
            progress_bar_display.progress((i + 1) / len(frames))
            # Introduce a delay to achieve 24 FPS
            time.sleep(fps_delay)

        # Display completion message
        progress_bar_processing_slot.text("Video Playback Finished!")