|
import streamlit as st |
|
from PIL import Image |
|
from ultralytics import YOLO |
|
|
|
st.title("Detector de armas - LoBa") |
|
|
|
model = YOLO('yolov8_background1k_best.pt') |
|
|
|
def detect_objects(image): |
|
result = model(image, conf = 0.28) |
|
return result |
|
|
|
|
|
uploaded_file = st.file_uploader("Elige una imagen...", type=["jpg", "png", "jpeg", "avif"]) |
|
|
|
if uploaded_file is not None: |
|
image = Image.open(uploaded_file) |
|
st.image(image, caption='Cargar imagen', use_container_width=True) |
|
|
|
if st.button('Detectar objetos'): |
|
image.save('uploaded_image.jpg') |
|
|
|
result = detect_objects('uploaded_image.jpg') |
|
|
|
boxes = result[0].boxes |
|
masks = result[0].masks |
|
keypoints = result[0].keypoints |
|
probs = result[0].probs |
|
|
|
st.write("Número de objetos detectados:", len(boxes)) |
|
result[0].save(filename='result.jpg') |
|
|
|
st.image('result.jpg', use_container_width=True) |
|
|