wip

app.py

@@ -138,8 +138,9 @@ with gr.Blocks() as demo:
         # Process Button
         obj_process_btn = gr.Button("Detect Objects")
 
-        # Output Component
-        obj_image_output = gr.Image(label="Detected Objects")
+        # Output Components
+        obj_image_output = gr.Image(label="Detected Objects Image") # Updated label for clarity
+        obj_raw_output = gr.JSON(label="Raw Object Detection Data") # Added JSON output
 
         # Link radio button change to visibility update function
         obj_input_type.change(
@@ -151,11 +152,11 @@ with gr.Blocks() as demo:
         )
 
         # Link process button to the object detection function
-        # The object_detection function will need to be updated to handle these inputs
+        # The object_detection function now returns a tuple (image, raw_data)
         obj_process_btn.click(
             fn=object_detection,
             inputs=[obj_input_type, obj_img_upload, obj_url_input, obj_base64_input],
-            outputs=obj_image_output,
+            outputs=[obj_image_output, obj_raw_output], # Updated outputs
         )
 
     # Launch the Gradio demo

detection/object_detection.py

@@ -1,17 +1,29 @@
 # Standard library imports
 # (Add any necessary imports for future object detection implementation)
+import json
 
 # Third-party imports
 from PIL import Image
 import numpy as np
+from ultralytics import YOLO
 
 # Local imports
 from utils.image_utils import load_image, preprocess_image
 
 
+# Load the YOLOv8 model globally to avoid reloading on each function call
+# Using a common pre-trained YOLOv8 nano model ('yolov8n.pt')
+try:
+    model = YOLO('yolov8n.pt')
+    print("YOLOv8 model loaded successfully.")
+except Exception as e:
+    print(f"Error loading YOLOv8 model: {e}")
+    model = None # Set model to None if loading fails
+
+
 def object_detection(input_type, uploaded_image, image_url, base64_string):
     """
-    Performs object detection on the image from various input types.
+    Performs object detection on the image from various input types using YOLOv8.
 
     Args:
         input_type (str): The selected input method ("Upload File", "Enter URL", "Enter Base64").
@@ -20,8 +32,14 @@ def object_detection(input_type, uploaded_image, image_url, base64_string):
         base64_string (str): The image base64 string (if input_type is "Enter Base64").
 
     Returns:
-        numpy.ndarray: The image with detected objects, or None if an error occurred.
+        tuple: A tuple containing:
+            - numpy.ndarray: The image with detected objects drawn on it, or None if an error occurred or model not loaded.
+            - dict: A dictionary containing the raw detection data (bounding boxes, classes, scores), or None.
     """
+    if model is None:
+        print("YOLOv8 model is not loaded. Cannot perform object detection.")
+        return None, None # Return None for both outputs
+
     image = None
     input_value = None
 
@@ -39,28 +57,51 @@ def object_detection(input_type, uploaded_image, image_url, base64_string):
 
     else:
         print("No valid input provided for object detection based on selected type.")
-        return None  # No valid input
+        return None, None # Return None for both outputs
 
     # If input_value is set (URL or Base64), use load_image
     if input_value:
         image = load_image(input_value)
         if image is None:
-            return None  # load_image failed
+            return None, None # load_image failed
 
     # Now 'image' should be a PIL Image or None
     if image is None:
         print("Image is None after loading/selection for object detection.")
-        return None
+        return None, None # Return None for both outputs
 
     try:
         # Preprocess the image (convert PIL to numpy, ensure RGB)
-        # preprocess_image expects a PIL Image or something convertible by Image.fromarray
-        processed_image = preprocess_image(image)
+        processed_image_np = preprocess_image(image)
+
+        # Perform inference
+        results = model.predict(processed_image_np)
+
+        # Extract raw detection data
+        raw_data = []
+        if results and results[0].boxes:
+            for box in results[0].boxes:
+                # box.xywh contains [x_center, y_center, width, height]
+                # box.conf contains confidence score
+                # box.cls contains class index
+                x_center, y_center, width, height = [round(float(coord), 2) for coord in box.xywh[0].tolist()] # Changed to xywh
+                confidence = round(float(box.conf[0]), 4)
+                class_id = int(box.cls[0])
+                class_name = model.names[class_id] if model.names else str(class_id) # Get class name if available
+
+                raw_data.append({
+                    "box": {"x": x_center, "y": y_center, "w": width, "h": height}, # Updated keys
+                    "confidence": confidence,
+                    "class_id": class_id,
+                    "class_name": class_name
+                })
+
+        # Draw results on the image
+        result_image_np = results[0].plot() if results else processed_image_np # Plot if results exist
+
+        print("Object detection performed successfully.")
+        return result_image_np, raw_data # Return both the image and raw data
 
-        # TODO: Implement object detection logic here
-        # Currently just returns the processed image
-        print("Object detection logic placeholder executed.")
-        return processed_image
     except Exception as e:
-        print(f"Error in object detection processing: {e}")
-        return None
+        print(f"Error during YOLOv8 object detection: {e}")
+        return None, None # Return None for both outputs