Spaces:
Sleeping
Sleeping
| # 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 using YOLOv8. | |
| Args: | |
| input_type (str): The selected input method ("Upload File", "Enter URL", "Enter Base64"). | |
| uploaded_image (PIL.Image.Image): The uploaded image (if input_type is "Upload File"). | |
| image_url (str): The image URL (if input_type is "Enter URL"). | |
| base64_string (str): The image base64 string (if input_type is "Enter Base64"). | |
| Returns: | |
| 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 | |
| if input_type == "Upload File" and uploaded_image is not None: | |
| image = uploaded_image # This is a PIL Image | |
| print("Using uploaded image (PIL) for object detection") # Debug print | |
| elif input_type == "Enter URL" and image_url and image_url.strip(): | |
| input_value = image_url | |
| print(f"Using URL for object detection: {input_value}") # Debug print | |
| elif input_type == "Enter Base64" and base64_string and base64_string.strip(): | |
| input_value = base64_string | |
| print(f"Using Base64 string for object detection") # Debug print | |
| else: | |
| print("No valid input provided for object detection based on selected type.") | |
| 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, 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, None # Return None for both outputs | |
| try: | |
| # Preprocess the image (convert PIL to numpy, ensure RGB) | |
| 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 | |
| except Exception as e: | |
| print(f"Error during YOLOv8 object detection: {e}") | |
| return None, None # Return None for both outputs | |