Update app.py

app.py

@@ -1,55 +1,115 @@
 import gradio as gr
-import supervision as sv
 import numpy as np
 import cv2
-from inference import get_roboflow_model
-
-# Define the Roboflow model
-model = get_roboflow_model(model_id="nescafe-4base/46", api_key="Otg64Ra6wNOgDyjuhMYU")
+import supervision as sv
+from roboflow import Roboflow
+import tempfile
+import os
+import requests
+from dotenv import load_dotenv
 
-def callback(image_slice: np.ndarray) -> sv.Detections:
-    # Perform inference on the image slice
-    results = model.infer(image_slice)[0]
-    return sv.Detections.from_inference(results)  # Wrap inference results in the proper supervision format
+# 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"))
 
-# Define the slicer
-slicer = sv.InferenceSlicer(callback=callback)
+# 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):
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)  # Convert from RGB (Gradio) to BGR (OpenCV)
+    # 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
     
-    # Run inference
-    sliced_detections = slicer(image=image)
+    try:
+        # Perform inference on the uploaded image using the Roboflow model
+        predictions = model.predict(temp_file_path, confidence=60, overlap=80).json()
+
+        # 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']
+            
+            # 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
+                )
+            )
+
+        # Convert detections to a Detections object for Supervision
+        detections = sv.Detections(detections)
 
-    # Annotating the image with boxes and labels
-    label_annotator = sv.LabelAnnotator()
-    box_annotator = sv.BoxAnnotator()
+        # Annotate the image with bounding boxes and labels
+        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)
+        annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections)
 
-    annotated_image = box_annotator.annotate(scene=image.copy(), detections=sliced_detections)
-    annotated_image = label_annotator.annotate(scene=annotated_image, detections=sliced_detections)
+        # Count detected objects per class
+        class_count = {}
+        total_count = 0
 
-    # Count detected objects per class
-    class_counts = {}
+        for detection in detections:
+            class_name = detection.class_name
+            class_count[class_name] = class_count.get(class_name, 0) + 1
+            total_count += 1
+
+        # 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}"
+
+        # Save the annotated image as output
+        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
+
+    # Clean up by removing the temporary file
+    os.remove(temp_file_path)
     
-    # Loop through the detections, which should now be in the correct format
-    for detection in sliced_detections:
-        class_name = detection.class_name  # Now `detection` should be a detection object with class_name
-        class_counts[class_name] = class_counts.get(class_name, 0) + 1
+    return output_image_path, result_text
 
-    # Total objects detected
-    total_count = sum(class_counts.values())
+# 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")
     
-    # Display results: annotated image and object counts
-    result_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)  # Convert back to RGB for Gradio
-    return result_image, class_counts, total_count
-
-# Create a Gradio interface
-iface = gr.Interface(
-    fn=detect_objects,
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="pil"), gr.JSON(), gr.Number(label="Total Objects Detected")],
-    live=True
-)
+    detect_button.click(
+        fn=detect_objects,
+        inputs=input_image,
+        outputs=[output_image, output_text]
+    )
 
 # Launch the Gradio interface
 iface.launch()