Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,7 +6,7 @@ from ultralytics import YOLO
 
 @st.cache_resource(show_spinner=False)
 def load_model():
-    return YOLO('src/best.pt') 
+    return YOLO('best.pt') 
 
 def predict(image, model):
     img_np = np.array(image)
@@ -17,9 +17,13 @@ def predict(image, model):
     return results[0]
 
 def draw_results(image, results):
-    img_annotated = results.plot()
-    img_annotated_rgb = cv2.cvtColor(img_annotated, cv2.COLOR_BGR2RGB)
-    return Image.fromarray(img_annotated_rgb)
+    try:
+        img_annotated = results.plot()
+        img_annotated_rgb = cv2.cvtColor(img_annotated, cv2.COLOR_BGR2RGB)
+        return Image.fromarray(img_annotated_rgb)
+    except Exception as e:
+        st.error(f"Error during annotation: {e}")
+        return image
 
 def crop_box(image, box):
     img_np = np.array(image)