Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| import os | |
| import cv2 | |
| from PIL import Image | |
| class ObjectDetector: | |
| def __init__(self, model_path="models/yolov11m.pt", threshold=0.5): | |
| """Initialize the object detector with YOLO model""" | |
| try: | |
| self.model = YOLO(model_path) | |
| self.threshold = threshold | |
| print(f"[INFO] Model loaded successfully: {model_path}") | |
| except Exception as e: | |
| print(f"❌ Error loading model: {e}") | |
| raise | |
| def detect_objects_in_folder( | |
| self, input_folder="core/input_frames", output_folder="core/output_frames" | |
| ): | |
| """Detect objects in all images in the input folder and save annotated images""" | |
| os.makedirs(output_folder, exist_ok=True) | |
| # Check if input folder exists and has images | |
| if not os.path.exists(input_folder): | |
| print(f"❌ Input folder not found: {input_folder}") | |
| return | |
| image_files = [ | |
| f | |
| for f in os.listdir(input_folder) | |
| if f.lower().endswith((".jpg", ".jpeg", ".png")) | |
| ] | |
| if not image_files: | |
| print(f"❌ No image files found in: {input_folder}") | |
| return | |
| print(f"[INFO] Processing {len(image_files)} images...") | |
| for filename in image_files: | |
| image_path = os.path.join(input_folder, filename) | |
| try: | |
| # Run detection | |
| results = self.model(image_path, conf=self.threshold) | |
| if len(results) == 0: | |
| print(f"[NO DETECTIONS] {filename}") | |
| continue | |
| # Get detections | |
| result = results[0] | |
| if result.boxes is None or len(result.boxes) == 0: | |
| print(f"[NO DETECTIONS] {filename}") | |
| continue | |
| # Load original image for annotation | |
| original_image = cv2.imread(image_path) | |
| if original_image is None: | |
| print(f"❌ Could not load image: {filename}") | |
| continue | |
| # Get bounding boxes and draw them | |
| boxes = result.boxes.data.cpu().numpy() | |
| detected_count = 0 | |
| for box in boxes: | |
| x1, y1, x2, y2, conf, cls = box | |
| # Convert to integers for drawing | |
| x1, y1, x2, y2 = map(int, [x1, y1, x2, y2]) | |
| # Draw bounding box | |
| cv2.rectangle(original_image, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| # Add label | |
| label = f"Drone: {conf:.2f}" | |
| cv2.putText( | |
| original_image, | |
| label, | |
| (x1, y1 - 10), | |
| cv2.FONT_HERSHEY_SIMPLEX, | |
| 0.5, | |
| (255, 0, 0), | |
| 2, | |
| ) | |
| detected_count += 1 | |
| # Save annotated image | |
| output_path = os.path.join(output_folder, filename) | |
| cv2.imwrite(output_path, original_image) | |
| print(f"[SAVED] {output_path} ({detected_count} objects)") | |
| except Exception as e: | |
| print(f"❌ Error processing {filename}: {e}") | |
| continue | |