wip

detection/object_detection.py

@@ -10,20 +10,21 @@ from ultralytics import YOLO
 # Local imports
 from utils.image_utils import load_image, preprocess_image
 
+YOLO_MODEL = "yolo11n.pt"
 
-# Load the YOLOv8 model globally to avoid reloading on each function call
-# Using a common pre-trained YOLOv8 nano model ('yolov8n.pt')
+# Load the YOLO model globally to avoid reloading on each function call
+# Using a common pre-trained YOLO nano model ('yolov8n.pt')
 try:
-    model = YOLO('yolov8n.pt')
-    print("YOLOv8 model loaded successfully.")
+    model = YOLO(YOLO_MODEL)
+    print("YOLO model loaded successfully.")
 except Exception as e:
-    print(f"Error loading YOLOv8 model: {e}")
+    print(f"Error loading YOLO 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 using YOLOv8.
+    Performs object detection on the image from various input types using YOLO (YOLOv11 nano).
 
     Args:
         input_type (str): The selected input method ("Upload File", "Enter URL", "Enter Base64").
@@ -37,7 +38,7 @@ def object_detection(input_type, uploaded_image, image_url, base64_string):
             - 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.")
+        print("YOLO model is not loaded. Cannot perform object detection.")
         return None, None # Return None for both outputs
 
     image = None
@@ -84,7 +85,7 @@ def object_detection(input_type, uploaded_image, image_url, base64_string):
                 # 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
+                x_center, y_center, width, height = [round(float(coord)) 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
@@ -103,5 +104,5 @@ def object_detection(input_type, uploaded_image, image_url, base64_string):
         return result_image_np, raw_data # Return both the image and raw data
 
     except Exception as e:
-        print(f"Error during YOLOv8 object detection: {e}")
+        print(f"Error during YOLO object detection: {e}")
         return None, None # Return None for both outputs