import gradio as gr import cv2 import numpy as np from PIL import Image import torch from transformers import DetrImageProcessor, DetrForObjectDetection from collections import defaultdict import time import psutil import os # Load DETR model (optimized for CPU) print("Loading DETR model...") processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") model.eval() # COCO class labels COCO_CLASSES = [ 'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] def get_available_memory(): """Get available system memory in GB""" return psutil.virtual_memory().available / (1024 ** 3) def auto_adjust_confidence(image_size, num_objects_hint=None): """Dynamically adjust confidence based on image complexity""" pixels = image_size[0] * image_size[1] # Base confidence on image size if pixels < 500000: # Small image base_confidence = 0.6 elif pixels < 2000000: # Medium image base_confidence = 0.65 else: # Large image base_confidence = 0.7 return base_confidence def auto_calculate_frame_interval(total_frames, video_duration, available_memory_gb): """Dynamically calculate optimal frame interval based on video properties and system resources""" # Base calculations fps = total_frames / video_duration if video_duration > 0 else 30 # Memory-based adjustment if available_memory_gb < 2: memory_factor = 3 elif available_memory_gb < 4: memory_factor = 2 else: memory_factor = 1 # Duration-based adjustment if video_duration < 10: duration_factor = 1 elif video_duration < 30: duration_factor = 2 elif video_duration < 60: duration_factor = 3 else: duration_factor = 4 # Calculate optimal frames to process target_frames = min(150, max(30, total_frames // (memory_factor * duration_factor))) # Calculate interval interval = max(1, total_frames // target_frames) return interval, target_frames def detect_objects(image, confidence_threshold=None): """Detect objects in a single image with dynamic confidence""" # Convert to RGB if needed if isinstance(image, np.ndarray): image = Image.fromarray(image) # Auto-adjust confidence if not provided if confidence_threshold is None: confidence_threshold = auto_adjust_confidence(image.size) # Prepare image inputs = processor(images=image, return_tensors="pt") # Inference with torch.no_grad(): outputs = model(**inputs) # Post-process target_sizes = torch.tensor([image.size[::-1]]) results = processor.post_process_object_detection( outputs, target_sizes=target_sizes, threshold=confidence_threshold )[0] return results, image, confidence_threshold def draw_boxes(image, results): """Draw bounding boxes on image""" img_array = np.array(image) detections = [] object_counts = defaultdict(int) for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): box = [round(i, 2) for i in box.tolist()] label_name = COCO_CLASSES[label.item()] if label_name != 'N/A': # Draw rectangle x1, y1, x2, y2 = map(int, box) cv2.rectangle(img_array, (x1, y1), (x2, y2), (0, 255, 0), 2) # Add label label_text = f"{label_name}: {score:.2f}" cv2.putText(img_array, label_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) detections.append(f"{label_name} ({score:.2%})") object_counts[label_name] += 1 return img_array, detections, object_counts def process_static_image(image): """Process static image mode with auto-detection""" if image is None: return None, "Please upload an image" start_time = time.time() # Detect objects with auto-adjusted confidence results, pil_image, used_confidence = detect_objects(image, confidence_threshold=None) # Draw boxes annotated_image, detections, object_counts = draw_boxes(pil_image, results) processing_time = time.time() - start_time # Create detailed summary if detections: summary = f"### 🎯 Detection Results\n\n" summary += f"**Found {len(detections)} objects in {processing_time:.2f} seconds**\n\n" summary += f"*Auto-adjusted confidence threshold: {used_confidence:.2f}*\n\n" summary += "#### Detected Objects:\n" # Group by object type for obj_name, count in sorted(object_counts.items(), key=lambda x: x[1], reverse=True): summary += f"- **{obj_name}**: {count} instance(s)\n" summary += f"\n#### All Detections:\n" for i, d in enumerate(detections, 1): summary += f"{i}. {d}\n" else: summary = f"### ⚠️ No objects detected\n\n" summary += f"*Confidence threshold used: {used_confidence:.2f}*\n\n" summary += "Try uploading a different image with more visible objects." return annotated_image, summary def process_video(video_path, progress=gr.Progress()): """Process video mode with full auto-adjustment""" if video_path is None: return None, "Please upload a video" progress(0, desc="Analyzing video...") cap = cv2.VideoCapture(video_path) # Get video properties fps = int(cap.get(cv2.CAP_PROP_FPS)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) duration = total_frames / fps if fps > 0 else 0 # Get available memory available_memory = get_available_memory() # Auto-calculate optimal frame interval frame_interval, estimated_frames = auto_calculate_frame_interval( total_frames, duration, available_memory ) progress(0.1, desc=f"Processing video (sampling every {frame_interval} frames)...") # Auto-adjust confidence based on video properties frame_size = width * height confidence_threshold = auto_adjust_confidence((width, height)) # Output video writer output_path = "output_video.mp4" fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) frame_count = 0 processed_count = 0 object_tracker = defaultdict(int) start_time = time.time() while cap.isOpened(): ret, frame = cap.read() if not ret: break # Update progress if frame_count % 30 == 0: progress_pct = (frame_count / total_frames) * 0.8 + 0.1 progress(progress_pct, desc=f"Processing frame {frame_count}/{total_frames}") # Process every nth frame if frame_count % frame_interval == 0: # Convert BGR to RGB rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Detect objects results, _, _ = detect_objects(rgb_frame, confidence_threshold) # Draw boxes for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): box = [round(i, 2) for i in box.tolist()] label_name = COCO_CLASSES[label.item()] if label_name != 'N/A': x1, y1, x2, y2 = map(int, box) cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) label_text = f"{label_name}: {score:.2f}" cv2.putText(frame, label_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) object_tracker[label_name] += 1 processed_count += 1 out.write(frame) frame_count += 1 cap.release() out.release() processing_time = time.time() - start_time progress(1.0, desc="Complete!") # Create detailed summary summary = f"### 🎬 Video Processing Complete\n\n" summary += f"**Processing Time**: {processing_time:.2f} seconds\n\n" summary += "#### Video Information:\n" summary += f"- Duration: {duration:.2f} seconds\n" summary += f"- Total frames: {total_frames}\n" summary += f"- FPS: {fps}\n" summary += f"- Resolution: {width}x{height}\n\n" summary += "#### Auto-Optimization Settings:\n" summary += f"- Confidence threshold: {confidence_threshold:.2f} *(auto-adjusted)*\n" summary += f"- Frame interval: Every {frame_interval} frame(s) *(auto-calculated)*\n" summary += f"- Frames processed: {processed_count}/{total_frames}\n" summary += f"- Available memory: {available_memory:.2f} GB\n\n" if object_tracker: summary += "### 📊 Detected Objects Across Video:\n\n" for obj, count in sorted(object_tracker.items(), key=lambda x: x[1], reverse=True): summary += f"- **{obj}**: {count} detection(s)\n" else: summary += "⚠️ No objects detected in the video.\n" summary += "This might be due to low lighting, fast motion, or absence of recognizable objects." return output_path, summary # Create Gradio interface with gr.Blocks(title="AI Object Recognition System", theme=gr.themes.Soft()) as demo: gr.Markdown(""" # 🤖 AI Object Recognition System ### Intelligent Auto-Adjusting Detection & Tracking This system **automatically optimizes** detection parameters based on: - Image/video size and complexity - Available system resources - Video duration and frame rate **No manual tuning required!** """) with gr.Tabs(): # Static Image Tab with gr.Tab("📸 Static Mode - Image Detection"): gr.Markdown(""" ### Automatic Image Analysis Upload any image and the system will: - Auto-adjust confidence thresholds - Detect all visible objects - Provide detailed statistics """) with gr.Row(): with gr.Column(): static_input = gr.Image(type="numpy", label="Upload Image") static_btn = gr.Button("🔍 Auto-Detect Objects", variant="primary", size="lg") gr.Markdown("*The system will automatically optimize detection settings*") with gr.Column(): static_output = gr.Image(label="Detected Objects") static_summary = gr.Markdown(label="Detection Results") static_btn.click( fn=process_static_image, inputs=[static_input], outputs=[static_output, static_summary] ) gr.Examples( examples=[], inputs=static_input, label="Try these examples (upload your own images)" ) # Dynamic Video Tab with gr.Tab("🎥 Dynamic Mode - Video Detection"): gr.Markdown(""" ### Automatic Video Analysis Upload a video and the system will: - Auto-calculate optimal frame sampling - Adjust confidence based on video quality - Optimize for available CPU resources - Track objects across frames **Supports videos of any length!** The system automatically scales processing. """) with gr.Row(): with gr.Column(): video_input = gr.Video(label="Upload Video") video_btn = gr.Button("🎬 Auto-Process Video", variant="primary", size="lg") gr.Markdown(""" *The system will automatically:* - Analyze video properties - Calculate optimal frame sampling - Adjust detection thresholds - Monitor system resources """) with gr.Column(): video_output = gr.Video(label="Processed Video with Detections") video_summary = gr.Markdown(label="Processing Results") video_btn.click( fn=process_video, inputs=[video_input], outputs=[video_output, video_summary] ) gr.Markdown(""" --- ## 🧠 How Auto-Adjustment Works ### Image Mode: - **Small images** (< 500K pixels): Lower confidence threshold for more detections - **Large images** (> 2M pixels): Higher threshold to reduce false positives ### Video Mode: - **Short videos** (< 10s): Process more frames for detail - **Long videos** (> 60s): Smart sampling to maintain performance - **Memory-aware**: Adjusts based on available RAM - **Quality-adaptive**: Balances speed vs accuracy automatically ### 📊 Technical Details: - **Model**: DETR ResNet-50 (Detection Transformer) - **Dataset**: COCO (80+ object categories) - **Optimization**: CPU-friendly with intelligent resource management - **Supported Objects**: People, vehicles, animals, furniture, electronics, food, and more ### 💡 Tips: - The system works best with clear, well-lit images/videos - All adjustments happen automatically - just upload and click! - Processing time varies based on video length and system resources """) if __name__ == "__main__": demo.launch()