Update app.py

app.py

@@ -1,121 +1,63 @@
 import gradio as gr
+import supervision as sv
 import numpy as np
 import cv2
-import supervision as sv
-from roboflow import Roboflow
-import tempfile
-import os
-from sahi.predict import predict
-from dotenv import load_dotenv
-
-# Load environment variables from .env file
-load_dotenv()
-api_key = os.getenv("ROBOFLOW_API_KEY")
-workspace = os.getenv("ROBOFLOW_WORKSPACE")
-project_name = os.getenv("ROBOFLOW_PROJECT")
-model_version = int(os.getenv("ROBOFLOW_MODEL_VERSION"))
-
-# Initialize Roboflow with the API key
-rf = Roboflow(api_key=api_key)
-project = rf.workspace(workspace).project(project_name)
-model = project.version(model_version).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
+from inference import get_roboflow_model
 
-    # Read the image using OpenCV
-    original_image = cv2.imread(temp_file_path)
+# Replace with your actual Roboflow model ID and API key
+model_id = "your-model-id"  # Replace with your Roboflow model ID
+api_key = "your-api-key"    # Replace with your Roboflow API key
 
-    try:
-        # Use SAHI to slice the image (optional for large images)
-        predictions = predict(
-            detection_model=model,  # Use Roboflow model for prediction
-            image=original_image,
-            slice_height=800,  # Height of each slice
-            slice_width=800,   # Width of each slice
-            overlap_height_ratio=0.2,
-            overlap_width_ratio=0.2,
-            return_slice_result=False,  # We don't need slice results, just detections
-        )
+# Load the Roboflow model using the get_roboflow_model function
+model = get_roboflow_model(model_id=model_id, api_key=api_key)
 
-        # Initialize Supervision annotations
-        detections = []
-        for prediction in predictions:
-            bbox = prediction.bbox
-            class_name = prediction.category
-            confidence = prediction.score
+# Define the callback function for the SAHI slicer
+def callback(image_slice: np.ndarray) -> sv.Detections:
+    # Run inference on the image slice
+    results = model.infer(image_slice)[0]
+    return sv.Detections.from_inference(results)
 
-            # Add detection to Supervision Detections list
-            detections.append(
-                sv.Detection(
-                    x1=bbox[0],
-                    y1=bbox[1],
-                    x2=bbox[2],
-                    y2=bbox[3],
-                    confidence=confidence,
-                    class_name=class_name
-                )
-            )
+# Initialize the SAHI Inference Slicer
+slicer = sv.InferenceSlicer(callback=callback)
 
-        # Convert detections to a Detections object for Supervision
-        detections = sv.Detections(detections)
+# Function to handle image processing, inference, and annotation
+def process_image(image):
+    # Convert the PIL image to OpenCV format (BGR)
+    image = np.array(image)
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
 
-        # Annotate the image with bounding boxes and labels
-        label_annotator = sv.LabelAnnotator()
-        box_annotator = sv.BoxAnnotator()
-        
-        # 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)
+    # Run inference using SAHI (splitting the image into slices)
+    sliced_detections = slicer(image=image)
 
-        # Count detected objects per class
-        class_count = {}
-        total_count = 0
-
-        for detection in detections:
-            class_name = detection.class_name
-            class_count[class_name] = class_count.get(class_name, 0) + 1
-            total_count += 1
+    # Annotate the detections with bounding boxes and labels
+    label_annotator = sv.LabelAnnotator()
+    box_annotator = sv.BoxAnnotator()
+    
+    annotated_image = box_annotator.annotate(scene=image.copy(), detections=sliced_detections)
+    annotated_image = label_annotator.annotate(scene=annotated_image, detections=sliced_detections)
 
-        # Prepare result text
-        result_text = "Detected Objects:\n\n"
-        for class_name, count in class_count.items():
-            result_text += f"{class_name}: {count}\n"
-        result_text += f"\nTotal objects detected: {total_count}"
+    # Convert the annotated image back to RGB for display in Gradio
+    result_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
 
-        # Save the annotated image as output
-        output_image_path = "/tmp/prediction.jpg"
-        cv2.imwrite(output_image_path, annotated_image)
+    # Count the number of objects detected
+    class_count = {}
+    for detection in sliced_detections:
+        class_name = detection.class_name
+        class_count[class_name] = class_count.get(class_name, 0) + 1
 
-    except Exception as err:
-        result_text = f"An error occurred: {err}"
-        output_image_path = temp_file_path  # Return original image on error
+    total_count = sum(class_count.values())
 
-    # Clean up by removing the temporary file
-    os.remove(temp_file_path)
-    
-    return output_image_path, result_text
+    return result_image, class_count, total_count
 
 # Gradio interface
-with gr.Blocks() as iface:
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Input Image")
-        with gr.Column():
-            output_image = gr.Image(label="Detected Image")
-        with gr.Column():
-            output_text = gr.Textbox(label="Object Count Results")
-    
-    detect_button = gr.Button("Detect")
-    
-    detect_button.click(
-        fn=detect_objects,
-        inputs=input_image,
-        outputs=[output_image, output_text]
-    )
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=[gr.Image(type="pil", label="Annotated Image"), 
+             gr.JSON(label="Object Count"), 
+             gr.Number(label="Total Objects Detected")],
+    live=True
+)
 
 # Launch the Gradio interface
 iface.launch()