Update app.py

app.py

@@ -5,7 +5,9 @@ import supervision as sv
 from roboflow import Roboflow
 import tempfile
 import os
-import requests
+from sahi import AutoDetectionModel
+from sahi.utils import read_image
+from sahi.predict import predict
 from dotenv import load_dotenv
 
 # Load environment variables from .env file
@@ -20,24 +22,37 @@ rf = Roboflow(api_key=api_key)
 project = rf.workspace(workspace).project(project_name)
 model = project.version(model_version).model
 
+# Initialize SAHI model for inference
+detection_model = AutoDetectionModel.from_pretrained(model)
+
 def detect_objects(image):
     # Save the uploaded image to a temporary file
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
         image.save(temp_file, format="JPEG")
         temp_file_path = temp_file.name
-    
+
+    # Read the image using OpenCV
+    original_image = cv2.imread(temp_file_path)
+
     try:
-        # Perform inference on the uploaded image using the Roboflow model
-        predictions = model.predict(temp_file_path, confidence=60, overlap=80).json()
+        # Perform inference using SAHI Slicer
+        predictions = predict(
+            detection_model=detection_model,
+            image=original_image,
+            slice_height=800,  # Height of the slice
+            slice_width=800,   # Width of the slice
+            overlap_height_ratio=0.2,
+            overlap_width_ratio=0.2,
+            return_slice_result=False,
+        )
 
         # Initialize Supervision annotations
         detections = []
-        for prediction in predictions['predictions']:
-            # Get bounding box and class for each prediction
-            bbox = prediction['bbox']
-            class_name = prediction['class']
-            confidence = prediction['confidence']
-            
+        for prediction in predictions:
+            bbox = prediction.bbox
+            class_name = prediction.category
+            confidence = prediction.score
+
             # Add detection to Supervision Detections list
             detections.append(
                 sv.Detection(
@@ -57,9 +72,8 @@ def detect_objects(image):
         label_annotator = sv.LabelAnnotator()
         box_annotator = sv.BoxAnnotator()
         
-        # Read the image back for OpenCV processing
-        image_cv = cv2.imread(temp_file_path)
-        annotated_image = box_annotator.annotate(scene=image_cv.copy(), detections=detections)
+        # Annotate and create the final result
+        annotated_image = box_annotator.annotate(scene=original_image.copy(), detections=detections)
         annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections)
 
         # Count detected objects per class
@@ -81,9 +95,6 @@ def detect_objects(image):
         output_image_path = "/tmp/prediction.jpg"
         cv2.imwrite(output_image_path, annotated_image)
 
-    except requests.exceptions.HTTPError as http_err:
-        result_text = f"HTTP error occurred: {http_err}"
-        output_image_path = temp_file_path  # Return original image on error
     except Exception as err:
         result_text = f"An error occurred: {err}"
         output_image_path = temp_file_path  # Return original image on error