Foxy-Roxy/Wheel_Defect_Detection

1import os
2os.environ["STREAMLIT_DISABLE_WATCHDOG_WARN"] = "1"
3os.environ["PYTHONASYNCIODEBUG"] = "1"
4import streamlit as st
5from ultralytics import YOLO
6from PIL import Image
7import tempfile
8
9# Load YOLO model
10@st.cache_resource
11def load_model():
12    return YOLO("best (2).pt")
13
14model = load_model()
15
16# Streamlit UI
17st.title("🛞 Defective Tyre Detection")
18st.write("Upload an image to detect defects using a YOLO model.")
19
20uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
21
22if uploaded_file:
23    # Load and show original image
24    image = Image.open(uploaded_file).convert("RGB")
25    st.image(image, caption="Uploaded Image", use_column_width=True)
26
27    # Save to temp file because ultralytics expects a path or ndarray
28    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
29        image.save(tmp.name)
30        results = model(tmp.name)
31
32    # Draw detections
33    result_image = Image.fromarray(results[0].plot())
34    st.image(result_image, caption="Detection Result", use_column_width=True)
35