import streamlit as st
from PIL import Image
from ultralytics import YOLO

st.title("Gun/Arms Detection")

model = YOLO('yolov8_background1k_best.pt')

def detect_objects(image):
    result = model(image, conf = 0.55)
    return result


uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_column_width=True)

    if st.button('Detect Objects'):
        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("Number of objects detected:", len(boxes))
        result[0].save(filename='result.jpg')
        
        st.image('result.jpg', use_column_width=True)