Update app.py

app.py

@@ -1,23 +1,31 @@
-
+import os
 import gradio as gr
 from ultralytics import YOLO
 from PIL import Image
-import numpy as np
 
-# Define the path to the trained YOLO model weights
-# This path should be accessible in the environment where the app is run
-model_path = "yolo_training/sickle_cls_model/weights/best.pt" # Update this path if needed
+# Suppress Ultralytics write warning by setting config dir
+os.environ["YOLO_CONFIG_DIR"] = "UltralyticsConfig"
+os.makedirs("UltralyticsConfig", exist_ok=True)
+
+# Path to YOLO model weights
+model_path = "yolo_training/sickle_cls_model/weights/best.pt"
+
+# Check if model exists
+if not os.path.exists(model_path):
+    raise FileNotFoundError(
+        f"Model file not found at {model_path}. "
+        "Make sure best.pt is uploaded to this path in your repo."
+    )
 
-# Load the YOLO model
+# Load YOLO model
 model = YOLO(model_path)
 
-# Define the list of class names used during model training
+# Class names
 classes = ['sickle', 'non_sickle', 'AIN']
 
-# Create a function that takes an image and returns the prediction
+# Prediction function
 def predict_image(img):
-    # Make prediction
-    # YOLO expects a list of images for prediction
+    # YOLO expects a list of images
     results = model([img])
 
     # Extract top prediction from the first image result
@@ -34,8 +42,13 @@ interface = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs=[gr.Label(), gr.Number()],
     title="Sickle Cell Classification",
-    description="Upload a blood smear image to classify it as Sickle Cell (sickle), Non-Sickle (non_sickle), or Artifact/Impurities/Noise (AIN)."
+    description=(
+        "Upload a blood smear image to classify it as "
+        "Sickle Cell (sickle), Non-Sickle (non_sickle), "
+        "or Artifact/Impurities/Noise (AIN)."
+    )
 )
 
-# Launch the interface (for local testing)
-# interface.launch()
+# Launch interface
+if __name__ == "__main__":
+    interface.launch()