Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| from ultralytics import YOLO | |
| import os | |
| import time | |
| from collections import deque | |
| from typing import Optional, Tuple, Generator | |
| import threading | |
| from queue import Queue | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # Configuration | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| class Config: | |
| """Configuration for the detection system""" | |
| MODEL_PATH = "best (3).pt" | |
| # Performance settings | |
| DEFAULT_CAMERA_WIDTH = 640 | |
| DEFAULT_CAMERA_HEIGHT = 480 | |
| DEFAULT_FPS = 30 | |
| # Detection settings | |
| DEFAULT_CONF_THRESHOLD = 0.25 | |
| DEFAULT_IOU_THRESHOLD = 0.45 | |
| DEFAULT_IMGSZ = 640 | |
| # CLAHE settings | |
| CLAHE_CLIP_LIMIT = 2.0 | |
| CLAHE_TILE_GRID = (8, 8) | |
| # Live detection settings | |
| FRAME_BUFFER_SIZE = 2 # Number of frames to buffer | |
| MAX_STREAM_TIME = 3600 # 1 hour | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # Load YOLO Model | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| try: | |
| model = YOLO(Config.MODEL_PATH) | |
| print(f"✅ Model loaded: {Config.MODEL_PATH}") | |
| print(f"📊 Classes: {model.names}") | |
| print(f"🔢 Number of classes: {len(model.names)}") | |
| except Exception as e: | |
| print(f"❌ Error loading model: {e}") | |
| model = None | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # CLAHE Preprocessing | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| class CLAHEProcessor: | |
| """Optimized CLAHE processor with caching""" | |
| def __init__(self, clip_limit: float = Config.CLAHE_CLIP_LIMIT, | |
| tile_grid: tuple = Config.CLAHE_TILE_GRID): | |
| self.clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid) | |
| def process(self, image: np.ndarray) -> np.ndarray: | |
| """Apply CLAHE preprocessing for shadow recovery.""" | |
| if image.dtype != np.uint8: | |
| image = np.clip(image, 0, 255).astype(np.uint8) | |
| # LAB color space - only enhance L channel | |
| bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) | |
| lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB) | |
| l, a, b = cv2.split(lab) | |
| # CLAHE on L channel | |
| l_enhanced = self.clahe.apply(l) | |
| # Merge back | |
| lab_enhanced = cv2.merge([l_enhanced, a, b]) | |
| bgr_enhanced = cv2.cvtColor(lab_enhanced, cv2.COLOR_LAB2BGR) | |
| rgb_enhanced = cv2.cvtColor(bgr_enhanced, cv2.COLOR_BGR2RGB) | |
| return rgb_enhanced | |
| # Global CLAHE processor | |
| clahe_processor = CLAHEProcessor() | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # FPS Counter | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| class FPSCounter: | |
| """Track FPS with moving average""" | |
| def __init__(self, buffer_size=30): | |
| self.frame_times = deque(maxlen=buffer_size) | |
| self.last_time = time.time() | |
| self.lock = threading.Lock() | |
| def update(self): | |
| with self.lock: | |
| current_time = time.time() | |
| self.frame_times.append(current_time - self.last_time) | |
| self.last_time = current_time | |
| def get_fps(self) -> float: | |
| with self.lock: | |
| if len(self.frame_times) == 0: | |
| return 0.0 | |
| return 1.0 / (sum(self.frame_times) / len(self.frame_times)) | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # Detection Statistics | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| class DetectionStats: | |
| """Track detection statistics over time""" | |
| def __init__(self, window_size=100): | |
| self.total_frames = 0 | |
| self.total_detections = 0 | |
| self.class_history = deque(maxlen=window_size) | |
| self.confidence_history = deque(maxlen=window_size) | |
| self.lock = threading.Lock() | |
| def update(self, class_counts: dict, confidences: list): | |
| with self.lock: | |
| self.total_frames += 1 | |
| frame_detections = sum(class_counts.values()) | |
| self.total_detections += frame_detections | |
| self.class_history.append(class_counts) | |
| if confidences: | |
| avg_conf = sum(confidences) / len(confidences) | |
| self.confidence_history.append(avg_conf) | |
| def get_average_detections_per_frame(self) -> float: | |
| with self.lock: | |
| if self.total_frames == 0: | |
| return 0.0 | |
| return self.total_detections / self.total_frames | |
| def get_most_common_class(self) -> Optional[str]: | |
| with self.lock: | |
| all_classes = {} | |
| for frame_counts in self.class_history: | |
| for class_name, count in frame_counts.items(): | |
| all_classes[class_name] = all_classes.get(class_name, 0) + count | |
| if not all_classes: | |
| return None | |
| return max(all_classes, key=all_classes.get) | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # Single Image Detection | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| def detect_engine_parts(image, conf_threshold=0.25, apply_clahe_preprocessing=True): | |
| """ | |
| Detect engine parts with YOLO (single image). | |
| """ | |
| if model is None: | |
| return image, "❌ Model not loaded. Please check the model file." | |
| if isinstance(image, Image.Image): | |
| image = np.array(image) | |
| if image is None: | |
| return None, "❌ No image provided" | |
| # Apply CLAHE preprocessing | |
| if apply_clahe_preprocessing: | |
| image = clahe_processor.process(image) | |
| # YOLO inference | |
| results = model.predict( | |
| source=image, | |
| conf=conf_threshold, | |
| iou=Config.DEFAULT_IOU_THRESHOLD, | |
| imgsz=Config.DEFAULT_IMGSZ, | |
| verbose=False, | |
| ) | |
| # Get annotated image | |
| annotated = results[0].plot() | |
| annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB) | |
| # Extract detection info | |
| boxes = results[0].boxes | |
| if len(boxes) == 0: | |
| summary = f"**No detections** (threshold: {conf_threshold:.0%})" | |
| else: | |
| summary = f"**Detected {len(boxes)} object(s):**\n\n" | |
| # Group by class | |
| class_counts = {} | |
| for box in boxes: | |
| class_id = int(box.cls[0]) | |
| class_name = model.names[class_id] | |
| class_counts[class_name] = class_counts.get(class_name, 0) + 1 | |
| # Summary by class | |
| for class_name, count in sorted(class_counts.items()): | |
| summary += f"• **{class_name}**: {count}\n" | |
| summary += f"\n**Details:**\n" | |
| # Individual detections | |
| for i, box in enumerate(boxes, 1): | |
| conf = float(box.conf[0]) | |
| class_id = int(box.cls[0]) | |
| class_name = model.names[class_id] | |
| summary += f"{i}. **{class_name}** — {conf:.2%} confidence\n" | |
| return annotated_rgb, summary | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # Visualization | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| def draw_info_overlay(frame: np.ndarray, fps: float, detection_counts: dict, | |
| total_detections: int, conf_threshold: float, | |
| avg_confidence: Optional[float] = None) -> np.ndarray: | |
| """Draw comprehensive info overlay on frame""" | |
| overlay = frame.copy() | |
| height, width = frame.shape[:2] | |
| # Calculate overlay height based on number of classes | |
| num_classes = len(detection_counts) | |
| overlay_height = max(180, 120 + num_classes * 25) | |
| # Semi-transparent background | |
| cv2.rectangle(overlay, (10, 10), (380, overlay_height), (0, 0, 0), -1) | |
| cv2.addWeighted(overlay, 0.7, frame, 0.3, 0, frame) | |
| # Header | |
| cv2.putText(frame, "LIVE DETECTION", (20, 35), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2) | |
| # FPS | |
| fps_color = (0, 255, 0) if fps > 20 else (255, 165, 0) if fps > 10 else (0, 0, 255) | |
| cv2.putText(frame, f"FPS: {fps:.1f}", (20, 65), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, fps_color, 2) | |
| # Total detections | |
| cv2.putText(frame, f"Objects: {total_detections}", (150, 65), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) | |
| # Confidence threshold | |
| cv2.putText(frame, f"Threshold: {conf_threshold:.0%}", (20, 90), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1) | |
| # Average confidence | |
| if avg_confidence is not None: | |
| cv2.putText(frame, f"Avg Conf: {avg_confidence:.0%}", (200, 90), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1) | |
| # Separator line | |
| cv2.line(frame, (20, 100), (360, 100), (100, 100, 100), 1) | |
| # Class counts | |
| if detection_counts: | |
| cv2.putText(frame, "Detected Parts:", (20, 125), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1) | |
| y_offset = 150 | |
| for class_name, count in sorted(detection_counts.items()): | |
| # Color code by count | |
| color = (0, 255, 0) if count >= 3 else (255, 255, 0) if count >= 2 else (255, 165, 0) | |
| text = f" {class_name}: {count}" | |
| cv2.putText(frame, text, (20, y_offset), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) | |
| y_offset += 25 | |
| else: | |
| cv2.putText(frame, "No detections", (20, 125), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (150, 150, 150), 1) | |
| return frame | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # Live Detection Engine | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| class LiveDetectionEngine: | |
| """Optimized live detection with threading""" | |
| def __init__(self): | |
| self.running = False | |
| self.frame_queue = Queue(maxsize=2) | |
| self.result_queue = Queue(maxsize=2) | |
| self.detection_thread = None | |
| self.stats = DetectionStats() | |
| def process_frame(self, frame: np.ndarray, conf_threshold: float, | |
| apply_clahe: bool) -> Tuple[np.ndarray, dict]: | |
| """Process a single frame""" | |
| if model is None: | |
| return frame, {} | |
| # Apply CLAHE | |
| if apply_clahe: | |
| processed_frame = clahe_processor.process(frame) | |
| else: | |
| processed_frame = frame | |
| # YOLO inference | |
| results = model.predict( | |
| source=processed_frame, | |
| conf=conf_threshold, | |
| iou=Config.DEFAULT_IOU_THRESHOLD, | |
| imgsz=Config.DEFAULT_IMGSZ, | |
| verbose=False, | |
| ) | |
| # Get annotated image | |
| annotated = results[0].plot() | |
| annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB) | |
| # Extract detection info | |
| boxes = results[0].boxes | |
| class_counts = {} | |
| confidences = [] | |
| for box in boxes: | |
| class_id = int(box.cls[0]) | |
| class_name = model.names[class_id] | |
| class_counts[class_name] = class_counts.get(class_name, 0) + 1 | |
| confidences.append(float(box.conf[0])) | |
| # Update stats | |
| self.stats.update(class_counts, confidences) | |
| detection_info = { | |
| 'class_counts': class_counts, | |
| 'total': len(boxes), | |
| 'confidences': confidences, | |
| 'avg_confidence': sum(confidences) / len(confidences) if confidences else None | |
| } | |
| return annotated_rgb, detection_info | |
| # Global detection engine | |
| detection_engine = LiveDetectionEngine() | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # Live Camera Detection | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| def live_detection_stream(conf_threshold: float = 0.25, | |
| apply_clahe: bool = True, | |
| process_every_n_frames: int = 1, | |
| show_advanced_stats: bool = False) -> Generator: | |
| """ | |
| Generator for live camera detection. | |
| """ | |
| fps_counter = FPSCounter() | |
| frame_count = 0 | |
| last_detection_info = {} | |
| last_annotated = None | |
| # Open webcam | |
| cap = cv2.VideoCapture(0) | |
| # Set camera properties | |
| cap.set(cv2.CAP_PROP_FRAME_WIDTH, Config.DEFAULT_CAMERA_WIDTH) | |
| cap.set(cv2.CAP_PROP_FRAME_HEIGHT, Config.DEFAULT_CAMERA_HEIGHT) | |
| cap.set(cv2.CAP_PROP_FPS, Config.DEFAULT_FPS) | |
| cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) # Minimize buffer for real-time | |
| if not cap.isOpened(): | |
| error_frame = np.zeros((480, 640, 3), dtype=np.uint8) | |
| cv2.putText(error_frame, "❌ Camera not available", (120, 240), | |
| cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2) | |
| cv2.putText(error_frame, "Please check camera permissions", (100, 280), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 1) | |
| yield error_frame | |
| return | |
| print("🎥 Live detection started") | |
| try: | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| print("⚠️ Failed to read frame") | |
| break | |
| # Convert BGR to RGB | |
| frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| frame_count += 1 | |
| # Process every Nth frame | |
| if frame_count % process_every_n_frames == 0: | |
| annotated_frame, detection_info = detection_engine.process_frame( | |
| frame_rgb, conf_threshold, apply_clahe | |
| ) | |
| last_annotated = annotated_frame | |
| last_detection_info = detection_info | |
| else: | |
| # Reuse last detection | |
| annotated_frame = last_annotated if last_annotated is not None else frame_rgb | |
| detection_info = last_detection_info | |
| # Update FPS | |
| fps_counter.update() | |
| fps = fps_counter.get_fps() | |
| # Draw overlay | |
| final_frame = draw_info_overlay( | |
| annotated_frame, | |
| fps, | |
| detection_info.get('class_counts', {}), | |
| detection_info.get('total', 0), | |
| conf_threshold, | |
| detection_info.get('avg_confidence') | |
| ) | |
| # Add advanced stats if enabled | |
| if show_advanced_stats: | |
| avg_det = detection_engine.stats.get_average_detections_per_frame() | |
| most_common = detection_engine.stats.get_most_common_class() | |
| y_pos = final_frame.shape[0] - 60 | |
| cv2.putText(final_frame, f"Avg Det/Frame: {avg_det:.1f}", (10, y_pos), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (150, 150, 150), 1) | |
| if most_common: | |
| cv2.putText(final_frame, f"Most Common: {most_common}", (10, y_pos + 20), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (150, 150, 150), 1) | |
| yield final_frame | |
| finally: | |
| cap.release() | |
| print("🛑 Live detection stopped") | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # Gradio Interface | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap'); | |
| body { | |
| font-family: 'Inter', sans-serif; | |
| background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); | |
| } | |
| .gr-button-primary { | |
| background: linear-gradient(90deg, #4CAF50 0%, #45a049 100%) !important; | |
| border: none !important; | |
| font-weight: 600 !important; | |
| } | |
| .gr-button-secondary { | |
| background: linear-gradient(90deg, #2196F3 0%, #1976D2 100%) !important; | |
| border: none !important; | |
| } | |
| .gradio-container { | |
| max-width: 1400px !important; | |
| } | |
| """ | |
| with gr.Blocks(css=CSS, title="Engine Part Detector - Live", theme=gr.themes.Soft()) as demo: | |
| gr.HTML(""" | |
| <div style='text-align: center; padding: 30px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 10px; margin-bottom: 20px;'> | |
| <h1 style='color: white; margin-bottom: 10px; font-size: 2.5em; text-shadow: 2px 2px 4px rgba(0,0,0,0.3);'> | |
| ⚙️ Engine Part Detection System | |
| </h1> | |
| <p style='color: rgba(255,255,255,0.9); font-size: 1.2em; margin: 0;'> | |
| YOLOv8 with Live Camera Feed & CLAHE Shadow Recovery | |
| </p> | |
| </div> | |
| """) | |
| with gr.Tabs(): | |
| # ═════════════════════════════════════════════════════════════════ | |
| # TAB 1: Live Camera Detection | |
| # ═════════════════════════════════════════════════════════════════ | |
| with gr.Tab("🎥 Live Camera Detection", id="live"): | |
| gr.Markdown(""" | |
| ### 🔴 Real-time engine part detection from your camera | |
| **Instructions:** Adjust settings below and the camera will start automatically. The feed shows FPS, detection counts, and confidence scores in real-time. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| # Camera feed | |
| live_video = gr.Image( | |
| label="Live Camera Feed with Detections", | |
| streaming=True, | |
| show_label=True, | |
| height=540 | |
| ) | |
| with gr.Column(scale=1): | |
| gr.Markdown("### ⚙️ Detection Settings") | |
| live_conf_slider = gr.Slider( | |
| minimum=0.1, | |
| maximum=0.9, | |
| value=0.25, | |
| step=0.05, | |
| label="🎯 Confidence Threshold", | |
| info="Lower = more sensitive (may detect false positives)" | |
| ) | |
| live_clahe_checkbox = gr.Checkbox( | |
| value=True, | |
| label="✨ CLAHE Shadow Recovery", | |
| info="Enhances details in shadowed areas (recommended)" | |
| ) | |
| frame_skip_slider = gr.Slider( | |
| minimum=1, | |
| maximum=5, | |
| value=1, | |
| step=1, | |
| label="⚡ Frame Skip (Performance)", | |
| info="Process every Nth frame (1=all, 3=every 3rd)" | |
| ) | |
| advanced_stats_checkbox = gr.Checkbox( | |
| value=False, | |
| label="📊 Show Advanced Statistics", | |
| info="Display average detections and most common class" | |
| ) | |
| gr.Markdown("---") | |
| gr.Markdown("### 📈 Performance Guide") | |
| gr.Markdown(""" | |
| **Target FPS:** 20-30 fps (green) | |
| - **<10 fps (red):** Increase frame skip to 2-3 | |
| - **10-20 fps (orange):** Optimal for most use cases | |
| - **>20 fps (green):** Excellent performance | |
| **Tips:** | |
| - Frame skip 2-3 = ~2x faster | |
| - Disable CLAHE = ~10% faster | |
| - Lower camera resolution = faster | |
| """) | |
| gr.Markdown("### 🎨 Info Overlay Legend") | |
| gr.Markdown(""" | |
| - **FPS:** Frames processed per second | |
| - **Objects:** Total parts detected in frame | |
| - **Threshold:** Current confidence cutoff | |
| - **Avg Conf:** Average confidence of detections | |
| - **Part Counts:** Color-coded by quantity | |
| - 🟢 Green: 3+ detected | |
| - 🟡 Yellow: 2 detected | |
| - 🟠 Orange: 1 detected | |
| """) | |
| # Start live detection | |
| live_video.stream( | |
| fn=live_detection_stream, | |
| inputs=[live_conf_slider, live_clahe_checkbox, frame_skip_slider, advanced_stats_checkbox], | |
| outputs=live_video, | |
| stream_every=0.05, # 20 FPS max update rate | |
| time_limit=Config.MAX_STREAM_TIME, | |
| ) | |
| # ═════════════════════════════════════════════════════════════════ | |
| # TAB 2: Single Image Detection | |
| # ═════════════════════════════════════════════════════════════════ | |
| with gr.Tab("📷 Single Image Detection", id="single"): | |
| gr.Markdown("### Upload an image or capture from webcam for batch detection") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| input_image = gr.Image( | |
| sources=["upload", "webcam"], | |
| type="numpy", | |
| label="Upload or Capture Image" | |
| ) | |
| conf_slider = gr.Slider( | |
| minimum=0.1, | |
| maximum=0.9, | |
| value=0.25, | |
| step=0.05, | |
| label="🎯 Confidence Threshold" | |
| ) | |
| clahe_checkbox = gr.Checkbox( | |
| value=True, | |
| label="✨ Apply CLAHE Preprocessing" | |
| ) | |
| detect_btn = gr.Button("🔍 Detect Parts", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| output_image = gr.Image(label="Detection Results") | |
| output_text = gr.Markdown(label="Detection Summary") | |
| detect_btn.click( | |
| fn=detect_engine_parts, | |
| inputs=[input_image, conf_slider, clahe_checkbox], | |
| outputs=[output_image, output_text], | |
| ) | |
| # ═════════════════════════════════════════════════════════════════════ | |
| # Help & Info Section | |
| # ═════════════════════════════════════════════════════════════════════ | |
| gr.HTML(""" | |
| <div style='margin-top: 30px; padding: 20px; background: white; border-radius: 10px; box-shadow: 0 4px 6px rgba(0,0,0,0.1);'> | |
| <h3 style='margin-top: 0; color: #333;'>ℹ️ How to Use This System</h3> | |
| <div style='display: grid; grid-template-columns: 1fr 1fr; gap: 20px; margin-top: 20px;'> | |
| <div> | |
| <h4 style='color: #2196F3;'>🎥 Live Camera Mode</h4> | |
| <ol style='color: #666;'> | |
| <li>Switch to "Live Camera Detection" tab</li> | |
| <li>Allow camera access when prompted</li> | |
| <li>Adjust confidence threshold (0.25 recommended)</li> | |
| <li>Enable CLAHE for better shadow handling</li> | |
| <li>Use frame skip for performance optimization</li> | |
| <li>Watch real-time detections with FPS overlay</li> | |
| </ol> | |
| </div> | |
| <div> | |
| <h4 style='color: #4CAF50;'>📷 Single Image Mode</h4> | |
| <ol style='color: #666;'> | |
| <li>Switch to "Single Image Detection" tab</li> | |
| <li>Upload image or capture from webcam</li> | |
| <li>Adjust confidence threshold if needed</li> | |
| <li>Enable CLAHE for shadowed images</li> | |
| <li>Click "Detect Parts" button</li> | |
| <li>View detailed results and summary</li> | |
| </ol> | |
| </div> | |
| </div> | |
| <div style='margin-top: 20px; padding: 15px; background: #f0f7ff; border-left: 4px solid #2196F3; border-radius: 4px;'> | |
| <p style='margin: 0; color: #333;'><strong>💡 Pro Tips:</strong></p> | |
| <ul style='color: #666; margin: 10px 0 0 0;'> | |
| <li><strong>Lighting:</strong> Ensure good lighting for best results. CLAHE helps with shadows but bright, even lighting is ideal.</li> | |
| <li><strong>Distance:</strong> Keep parts 30-100cm from camera for optimal detection.</li> | |
| <li><strong>Performance:</strong> If FPS is low, increase frame skip to 2 or 3.</li> | |
| <li><strong>Accuracy:</strong> Lower confidence threshold (0.15-0.20) for difficult cases, raise it (0.30-0.40) to reduce false positives.</li> | |
| </ul> | |
| </div> | |
| <div style='margin-top: 15px; padding: 15px; background: #fff3e0; border-left: 4px solid #FF9800; border-radius: 4px;'> | |
| <p style='margin: 0; color: #333;'><strong>⚠️ Troubleshooting:</strong></p> | |
| <ul style='color: #666; margin: 10px 0 0 0;'> | |
| <li><strong>Camera not working:</strong> Check browser permissions (chrome://settings/content/camera)</li> | |
| <li><strong>Low FPS:</strong> Increase frame skip, reduce camera resolution, or disable CLAHE</li> | |
| <li><strong>No detections:</strong> Lower confidence threshold or improve lighting</li> | |
| <li><strong>Too many false positives:</strong> Increase confidence threshold</li> | |
| </ul> | |
| </div> | |
| <p style='margin-top: 20px; color: #999; font-size: 0.9em; text-align: center;'> | |
| <strong>Model:</strong> YOLOv8 | <strong>Classes:</strong> Based on your training data | <strong>Performance:</strong> 15-30 FPS typical | |
| </p> | |
| </div> | |
| """) | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| # Launch | |
| # ═════════════════════════════════════════════════════════════════════════════ | |
| if __name__ == "__main__": | |
| print("\n" + "="*60) | |
| print("🚀 Starting Engine Part Detection System") | |
| print("="*60) | |
| print(f"📦 Model: {Config.MODEL_PATH}") | |
| print(f"🎥 Camera Resolution: {Config.DEFAULT_CAMERA_WIDTH}x{Config.DEFAULT_CAMERA_HEIGHT}") | |
| print(f"🎯 Default Confidence: {Config.DEFAULT_CONF_THRESHOLD}") | |
| print(f"⚡ CLAHE Enabled: Yes") | |
| print("="*60 + "\n") | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False, | |
| show_error=True, | |
| ) | |
| # import gradio as gr | |
| # import cv2 | |
| # import numpy as np | |
| # from PIL import Image | |
| # from ultralytics import YOLO | |
| # import os | |
| # import time | |
| # from collections import deque | |
| # from typing import Optional, Tuple | |
| # import threading | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # # Load YOLO Model | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # MODEL_PATH = "best (1).pt" | |
| # try: | |
| # model = YOLO(MODEL_PATH) | |
| # print(f"✅ Model loaded: {MODEL_PATH}") | |
| # print(f"📊 Classes: {model.names}") | |
| # except Exception as e: | |
| # print(f"❌ Error loading model: {e}") | |
| # model = None | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # # CLAHE Preprocessing (Shadow Recovery) | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # def apply_clahe(image: np.ndarray, clip_limit: float = 2.0) -> np.ndarray: | |
| # """Apply CLAHE preprocessing for shadow recovery.""" | |
| # if image.dtype != np.uint8: | |
| # image = np.clip(image, 0, 255).astype(np.uint8) | |
| # # LAB color space - only enhance L channel | |
| # bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) | |
| # lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB) | |
| # l, a, b = cv2.split(lab) | |
| # # CLAHE on L channel | |
| # clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8)) | |
| # l_enhanced = clahe.apply(l) | |
| # # Merge back | |
| # lab_enhanced = cv2.merge([l_enhanced, a, b]) | |
| # bgr_enhanced = cv2.cvtColor(lab_enhanced, cv2.COLOR_LAB2BGR) | |
| # rgb_enhanced = cv2.cvtColor(bgr_enhanced, cv2.COLOR_BGR2RGB) | |
| # return rgb_enhanced | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # # FPS Counter | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # class FPSCounter: | |
| # """Track FPS with moving average""" | |
| # def __init__(self, buffer_size=30): | |
| # self.frame_times = deque(maxlen=buffer_size) | |
| # self.last_time = time.time() | |
| # def update(self): | |
| # current_time = time.time() | |
| # self.frame_times.append(current_time - self.last_time) | |
| # self.last_time = current_time | |
| # def get_fps(self) -> float: | |
| # if len(self.frame_times) == 0: | |
| # return 0.0 | |
| # return 1.0 / (sum(self.frame_times) / len(self.frame_times)) | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # # Detection Functions | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # def detect_engine_parts(image, conf_threshold=0.25, apply_clahe_preprocessing=True): | |
| # """ | |
| # Detect engine parts with YOLO (single image). | |
| # Args: | |
| # image: Input image (PIL or numpy) | |
| # conf_threshold: Confidence threshold (0-1) | |
| # apply_clahe_preprocessing: Whether to apply CLAHE before detection | |
| # Returns: | |
| # annotated_image: Image with bounding boxes | |
| # results_text: Detection summary | |
| # """ | |
| # if model is None: | |
| # return image, "❌ Model not loaded. Please check the model file." | |
| # # Convert to numpy array | |
| # if isinstance(image, Image.Image): | |
| # image = np.array(image) | |
| # if image is None: | |
| # return None, "❌ No image provided" | |
| # # Apply CLAHE preprocessing | |
| # if apply_clahe_preprocessing: | |
| # image = apply_clahe(image) | |
| # # YOLO inference | |
| # results = model.predict( | |
| # source=image, | |
| # conf=conf_threshold, | |
| # iou=0.45, | |
| # imgsz=640, | |
| # verbose=False, | |
| # ) | |
| # # Get annotated image | |
| # annotated = results[0].plot() # BGR format | |
| # annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB) | |
| # # Extract detection info | |
| # boxes = results[0].boxes | |
| # if len(boxes) == 0: | |
| # summary = f"**No detections** (threshold: {conf_threshold:.0%})" | |
| # else: | |
| # summary = f"**Detected {len(boxes)} object(s):**\n\n" | |
| # # Group by class | |
| # class_counts = {} | |
| # for box in boxes: | |
| # class_id = int(box.cls[0]) | |
| # class_name = model.names[class_id] | |
| # class_counts[class_name] = class_counts.get(class_name, 0) + 1 | |
| # # Summary by class | |
| # for class_name, count in sorted(class_counts.items()): | |
| # summary += f"• **{class_name}**: {count}\n" | |
| # summary += f"\n**Details:**\n" | |
| # # Individual detections | |
| # for i, box in enumerate(boxes, 1): | |
| # x1, y1, x2, y2 = box.xyxy[0].tolist() | |
| # conf = float(box.conf[0]) | |
| # class_id = int(box.cls[0]) | |
| # class_name = model.names[class_id] | |
| # summary += f"{i}. **{class_name}** — {conf:.2%} confidence\n" | |
| # return annotated_rgb, summary | |
| # def draw_info_overlay(frame: np.ndarray, fps: float, detection_counts: dict, | |
| # total_detections: int, conf_threshold: float) -> np.ndarray: | |
| # """Draw FPS and detection info overlay on frame""" | |
| # overlay = frame.copy() | |
| # height, width = frame.shape[:2] | |
| # # Semi-transparent background for text | |
| # cv2.rectangle(overlay, (10, 10), (350, 150), (0, 0, 0), -1) | |
| # cv2.addWeighted(overlay, 0.6, frame, 0.4, 0, frame) | |
| # # FPS | |
| # cv2.putText(frame, f"FPS: {fps:.1f}", (20, 35), | |
| # cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) | |
| # # Total detections | |
| # cv2.putText(frame, f"Total: {total_detections}", (20, 60), | |
| # cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) | |
| # # Confidence threshold | |
| # cv2.putText(frame, f"Conf: {conf_threshold:.0%}", (20, 85), | |
| # cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) | |
| # # Class counts | |
| # y_offset = 110 | |
| # for class_name, count in sorted(detection_counts.items()): | |
| # text = f"{class_name}: {count}" | |
| # cv2.putText(frame, text, (20, y_offset), | |
| # cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 200, 0), 2) | |
| # y_offset += 25 | |
| # return frame | |
| # def process_video_frame(frame: np.ndarray, conf_threshold: float, | |
| # apply_clahe: bool, fps_counter: FPSCounter) -> Tuple[np.ndarray, dict]: | |
| # """ | |
| # Process a single video frame for detection. | |
| # Returns: | |
| # annotated_frame: Frame with detections drawn | |
| # detection_info: Dictionary with detection counts | |
| # """ | |
| # if model is None: | |
| # return frame, {} | |
| # # Apply CLAHE preprocessing | |
| # if apply_clahe: | |
| # processed_frame = apply_clahe(frame) | |
| # else: | |
| # processed_frame = frame | |
| # # YOLO inference | |
| # results = model.predict( | |
| # source=processed_frame, | |
| # conf=conf_threshold, | |
| # iou=0.45, | |
| # imgsz=640, | |
| # verbose=False, | |
| # stream=False, | |
| # ) | |
| # # Get annotated image | |
| # annotated = results[0].plot() # BGR format | |
| # annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB) | |
| # # Count detections by class | |
| # boxes = results[0].boxes | |
| # class_counts = {} | |
| # total_count = 0 | |
| # for box in boxes: | |
| # class_id = int(box.cls[0]) | |
| # class_name = model.names[class_id] | |
| # class_counts[class_name] = class_counts.get(class_name, 0) + 1 | |
| # total_count += 1 | |
| # # Update FPS | |
| # fps_counter.update() | |
| # fps = fps_counter.get_fps() | |
| # # Draw overlay | |
| # annotated_rgb = draw_info_overlay( | |
| # annotated_rgb, | |
| # fps, | |
| # class_counts, | |
| # total_count, | |
| # conf_threshold | |
| # ) | |
| # detection_info = { | |
| # 'class_counts': class_counts, | |
| # 'total': total_count, | |
| # 'fps': fps | |
| # } | |
| # return annotated_rgb, detection_info | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # # Live Camera Detection (Generator Function for Gradio) | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # def live_detection_stream(conf_threshold: float = 0.25, | |
| # apply_clahe: bool = True, | |
| # process_every_n_frames: int = 1): | |
| # """ | |
| # Generator function for live camera detection with Gradio. | |
| # Args: | |
| # conf_threshold: Detection confidence threshold | |
| # apply_clahe: Whether to apply CLAHE preprocessing | |
| # process_every_n_frames: Process every Nth frame (1 = all frames, 2 = every other frame) | |
| # Yields: | |
| # annotated_frame: Frame with detections and info overlay | |
| # """ | |
| # fps_counter = FPSCounter() | |
| # frame_count = 0 | |
| # last_detections = {} | |
| # last_annotated = None | |
| # # Open webcam | |
| # cap = cv2.VideoCapture(0) | |
| # # Set camera properties for better performance | |
| # cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640) | |
| # cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) | |
| # cap.set(cv2.CAP_PROP_FPS, 30) | |
| # if not cap.isOpened(): | |
| # # Return error frame | |
| # error_frame = np.zeros((480, 640, 3), dtype=np.uint8) | |
| # cv2.putText(error_frame, "Camera not available", (150, 240), | |
| # cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2) | |
| # yield error_frame | |
| # return | |
| # try: | |
| # while True: | |
| # ret, frame = cap.read() | |
| # if not ret: | |
| # break | |
| # # Convert BGR to RGB | |
| # frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| # frame_count += 1 | |
| # # Process every Nth frame for performance | |
| # if frame_count % process_every_n_frames == 0: | |
| # annotated_frame, detection_info = process_video_frame( | |
| # frame_rgb, | |
| # conf_threshold, | |
| # apply_clahe, | |
| # fps_counter | |
| # ) | |
| # last_annotated = annotated_frame | |
| # last_detections = detection_info | |
| # else: | |
| # # Reuse last detection but update FPS | |
| # if last_annotated is not None: | |
| # fps_counter.update() | |
| # annotated_frame = draw_info_overlay( | |
| # frame_rgb, | |
| # fps_counter.get_fps(), | |
| # last_detections.get('class_counts', {}), | |
| # last_detections.get('total', 0), | |
| # conf_threshold | |
| # ) | |
| # else: | |
| # annotated_frame = frame_rgb | |
| # yield annotated_frame | |
| # finally: | |
| # cap.release() | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # # Gradio Interface | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # CSS = """ | |
| # @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap'); | |
| # body { font-family: 'Inter', sans-serif; } | |
| # .gr-button-primary { background: linear-gradient(90deg, #4CAF50 0%, #45a049 100%) !important; } | |
| # .gr-button-secondary { background: linear-gradient(90deg, #2196F3 0%, #1976D2 100%) !important; } | |
| # """ | |
| # with gr.Blocks(css=CSS, title="Engine Part Detector") as demo: | |
| # gr.HTML(""" | |
| # <div style='text-align: center; padding: 20px;'> | |
| # <h1 style='color: #2196F3; margin-bottom: 10px;'>⚙️ Engine Part Detection System</h1> | |
| # <p style='color: #666; font-size: 1.1em;'> | |
| # YOLOv8 object detection with CLAHE shadow recovery preprocessing | |
| # </p> | |
| # </div> | |
| # """) | |
| # with gr.Tabs(): | |
| # # ═════════════════════════════════════════════════════════════════ | |
| # # TAB 1: Single Image Detection | |
| # # ═════════════════════════════════════════════════════════════════ | |
| # with gr.Tab("📷 Single Image Detection"): | |
| # gr.Markdown("### Upload an image or capture from webcam for one-time detection") | |
| # with gr.Row(): | |
| # with gr.Column(scale=1): | |
| # input_image = gr.Image( | |
| # sources=["upload", "webcam"], | |
| # type="numpy", | |
| # label="Upload Engine Part Image" | |
| # ) | |
| # with gr.Row(): | |
| # conf_slider = gr.Slider( | |
| # minimum=0.1, | |
| # maximum=0.9, | |
| # value=0.25, | |
| # step=0.05, | |
| # label="Confidence Threshold", | |
| # info="Lower = more detections (may include false positives)" | |
| # ) | |
| # clahe_checkbox = gr.Checkbox( | |
| # value=True, | |
| # label="Apply CLAHE Preprocessing", | |
| # info="Recovers details in shadowed areas (recommended)" | |
| # ) | |
| # detect_btn = gr.Button("🔍 Detect Parts", variant="primary", size="lg") | |
| # with gr.Column(scale=1): | |
| # output_image = gr.Image(label="Detection Results") | |
| # output_text = gr.Markdown(label="Summary") | |
| # # Wire up the button | |
| # detect_btn.click( | |
| # fn=detect_engine_parts, | |
| # inputs=[input_image, conf_slider, clahe_checkbox], | |
| # outputs=[output_image, output_text], | |
| # ) | |
| # # ═════════════════════════════════════════════════════════════════ | |
| # # TAB 2: Live Camera Detection | |
| # # ═════════════════════════════════════════════════════════════════ | |
| # with gr.Tab("🎥 Live Camera Detection"): | |
| # gr.Markdown(""" | |
| # ### Real-time engine part detection from your camera | |
| # **Note:** Click "Start Detection" to begin, and "Stop" to end the stream. | |
| # """) | |
| # with gr.Row(): | |
| # with gr.Column(scale=1): | |
| # # Camera feed | |
| # live_video = gr.Image( | |
| # label="Live Camera Feed", | |
| # streaming=True, | |
| # show_label=True, | |
| # height=480 | |
| # ) | |
| # with gr.Column(scale=1): | |
| # gr.Markdown("### ⚙️ Detection Settings") | |
| # live_conf_slider = gr.Slider( | |
| # minimum=0.1, | |
| # maximum=0.9, | |
| # value=0.25, | |
| # step=0.05, | |
| # label="Confidence Threshold", | |
| # info="Adjust sensitivity" | |
| # ) | |
| # live_clahe_checkbox = gr.Checkbox( | |
| # value=True, | |
| # label="Apply CLAHE Preprocessing", | |
| # info="Better results in shadows" | |
| # ) | |
| # frame_skip_slider = gr.Slider( | |
| # minimum=1, | |
| # maximum=5, | |
| # value=1, | |
| # step=1, | |
| # label="Process Every N Frames", | |
| # info="Higher = faster but less smooth (1 = process all frames)" | |
| # ) | |
| # gr.Markdown("### 📊 Performance Tips") | |
| # gr.Markdown(""" | |
| # - **Process Every N Frames**: Set to 2-3 for better performance | |
| # - **Confidence**: Lower threshold = more detections | |
| # - **CLAHE**: Slight performance cost but better quality | |
| # """) | |
| # gr.Markdown("### 📈 Info Overlay") | |
| # gr.Markdown(""" | |
| # The video shows: | |
| # - **FPS**: Frames per second | |
| # - **Total**: Total objects detected | |
| # - **Conf**: Current confidence threshold | |
| # - **Class counts**: Number of each part type | |
| # """) | |
| # # Start live detection | |
| # live_video.stream( | |
| # fn=live_detection_stream, | |
| # inputs=[live_conf_slider, live_clahe_checkbox, frame_skip_slider], | |
| # outputs=live_video, | |
| # stream_every=0.1, # Update interval in seconds | |
| # time_limit=3600, # 1 hour max | |
| # ) | |
| # # ═════════════════════════════════════════════════════════════════════ | |
| # # Info Section (Shared) | |
| # # ═════════════════════════════════════════════════════════════════════ | |
| # gr.HTML(""" | |
| # <div style='margin-top: 20px; padding: 15px; background: #f5f5f5; border-radius: 8px;'> | |
| # <h3 style='margin-top: 0;'>ℹ️ How to Use</h3> | |
| # <h4>📷 Single Image Mode:</h4> | |
| # <ol> | |
| # <li>Upload an image of an engine part or use your webcam to capture</li> | |
| # <li>Adjust confidence threshold (default 0.25 works well)</li> | |
| # <li>Enable CLAHE preprocessing for better results on shadowed images</li> | |
| # <li>Click "Detect Parts" to run detection</li> | |
| # </ol> | |
| # <h4>🎥 Live Camera Mode:</h4> | |
| # <ol> | |
| # <li>Switch to the "Live Camera Detection" tab</li> | |
| # <li>Adjust settings (confidence, CLAHE, frame skip)</li> | |
| # <li>The camera will start automatically showing real-time detections</li> | |
| # <li>View live FPS and detection counts on the video overlay</li> | |
| # <li>Close the tab or adjust settings to stop</li> | |
| # </ol> | |
| # <p><strong>Supported Classes:</strong> Depends on your model training (e.g., bearing_saddle, piston, defect, crack, corrosion)</p> | |
| # <p><strong>Performance:</strong> Live mode processes at 15-30 FPS depending on your hardware. Use "Process Every N Frames" to optimize speed.</p> | |
| # </div> | |
| # """) | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # # Launch | |
| # # ═════════════════════════════════════════════════════════════════════════════ | |
| # if __name__ == "__main__": | |
| # demo.launch( | |
| # server_name="0.0.0.0", | |
| # server_port=7860, | |
| # share=False, | |
| # ) |