import streamlit as st
from main import save_image
from ultralytics import YOLO
import shutil

model = YOLO('yolov8m.pt')

st.title('Object Detection')

save_image()

if st.button("Detect Object"):
    model.predict(source = 'processed_image.jpg', save = True, conf = 0.7)
    st.image('runs/detect/predict/processed_image.jpg', use_column_width = 'always')
    shutil.rmtree('runs')