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

def main():
    st.title("Gun Detection")

    video_source_option = st.radio("Select Video Source:", ("Video File", "RTSP Stream", "Webcam"))

    if video_source_option == "Video File":
        video_file = st.file_uploader("Upload Video", type=["mp4", "avi"])
        if video_file is not None:
            # Create a temporary file to store the uploaded video
            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                temp_file.write(video_file.read())
                temp_file_path = temp_file.name
            detect_objects(temp_file_path)
            # Remove the temporary file after processing
            os.unlink(temp_file_path)

    elif video_source_option == "RTSP Stream":
        rtsp_link = st.text_input("Enter RTSP Link:")
        if st.button("Start Detection"):
            detect_objects(rtsp_link)

    elif video_source_option == "Webcam":
        detect_objects(0)

def detect_objects(video_source):
    yolo_model = YOLO('350epochs.pt')
    cap = cv2.VideoCapture(video_source)
    placeholder = st.empty()
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # Get predictions
        results = yolo_model(frame)

        # Draw bounding boxes and labels on the frame
        annotated_frame = results[0].plot()

        # Convert the annotated frame to RGB (Streamlit uses RGB)
        annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)

        # Convert the frame to Image
        img_pil = Image.fromarray(annotated_frame_rgb)

        # Display the frame
        placeholder.image(img_pil, use_column_width=True)
        time.sleep(0.1)

if __name__ == '__main__':
    main()