import streamlit as st
from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
from ultralytics import YOLO
import av
from PIL import Image

# Cargar el modelo YOLOv8
model = YOLO('yolov8_background1k_best.pt')

# Configurar la página
st.title("Gun/Arms Detection 🚀")

# Definir el umbral de confianza
CONFIDENCE_THRESHOLD = 0.55

# Procesar imágenes con YOLOv8
def detect_objects(image):
    results = model(image, conf=CONFIDENCE_THRESHOLD)
    annotated_image = results[0].plot()
    return annotated_image, len(results[0].boxes)

# Clase para procesar frames de la webcam
class YOLOVideoProcessor(VideoTransformerBase):
    def __init__(self):
        self.model = model

    def transform(self, frame):
        image = frame.to_ndarray(format="bgr24")  # Convertir frame a numpy array
        results = self.model(image, conf=CONFIDENCE_THRESHOLD)  # Detección
        annotated_image = results[0].plot()  # Imagen anotada
        return av.VideoFrame.from_ndarray(annotated_image, format="bgr24")

# Interfaz con Streamlit
st.sidebar.title("Sube una imagen o usa la webcam")

# Opción para cargar imagen
uploaded_file = st.sidebar.file_uploader("Sube una imagen (JPG/PNG):", type=["jpg", "png", "jpeg"])

if uploaded_file:
    image = Image.open(uploaded_file)
    st.image(image, caption='Imagen subida', use_column_width=True)

    if st.button('Detectar objetos en la imagen'):
        annotated_image, num_boxes = detect_objects(image)
        st.image(annotated_image, caption=f"{num_boxes} objetos detectados", use_column_width=True)

# Opción para usar webcam
st.sidebar.write("O usa la webcam para detección en tiempo real:")
use_webcam = st.sidebar.checkbox("Usar Webcam")

if use_webcam:
    webrtc_streamer(key="object-detection", video_processor_factory=YOLOVideoProcessor)