File size: 1,812 Bytes
ae80fa4 db7145a ae80fa4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
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.30
# 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)
|