Update app.py

app.py

@@ -1,15 +1,14 @@
-import streamlit as st
-from PIL import Image
 import torch
-from ultralytics import YOLO  # Make sure YOLO is correctly installed
+from PIL import Image
 import numpy as np
+import streamlit as st
+from ultralytics import YOLO
 
-# Load the YOLO model
-model_path = "best.pt"
-model = YOLO(model_path)
+# Load YOLOv10 model
+model = YOLO('best.pt')  # Load the pre-trained model
 
 # Streamlit App
-st.title("YOLO Object Detection")
+st.title("YOLO Object Detection with Confidence Threshold")
 
 st.sidebar.title("Options")
 st.sidebar.markdown("Upload an image to detect objects.")
@@ -18,34 +17,24 @@ uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
     # Load the image
-    image = Image.open(uploaded_file)
+    image = Image.open(uploaded_file).convert('RGB')
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    # Convert PIL image to numpy array
-    image_np = np.array(image)
-
-    # Perform inference
-    st.write("Detecting objects...")
-    results = model.predict(image_np)
+    # Perform inference with a confidence threshold of 0.25
+    st.write("Detecting objects with confidence threshold of 0.25...")
+    results = model.predict(source=image, conf=0.25, save=False)  # Directly pass PIL image
 
-    # Draw predictions on the image
+    # Annotate and display the image
     annotated_image = results[0].plot()  # Get annotated image with bounding boxes
-
-    # Display the results
     st.image(annotated_image, caption="Detected Objects", use_column_width=True)
 
     # Show raw predictions
     st.write("Detection Results:")
     for result in results:
         for box in result.boxes:
-            class_id = int(box.cls.item())  # Convert tensor to Python int
-            confidence = float(box.conf.item())  # Convert tensor to Python float
-
-            # Ensure bbox is processed correctly
-            if isinstance(box.xyxy, torch.Tensor):
-                bbox = box.xyxy.tolist()  # Convert tensor to list
-            else:
-                bbox = box.xyxy  # Already in list format
+            class_id = int(box.cls)  # Convert to Python int
+            confidence = float(box.conf)  # Convert to Python float
+            bbox = box.xyxy.tolist()  # Bounding box coordinates as a list
 
             st.write(
                 f"Class: {class_id}, Confidence: {confidence:.2f}, Box: {bbox}"