Update app.py

app.py

@@ -0,0 +1,217 @@
+
+import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
+from ultralytics import YOLO
+import os
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Load YOLO Model
+# ═════════════════════════════════════════════════════════════════════════════
+
+MODEL_PATH = "best.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) -> 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=2.0, 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
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Detection Functions
+# ═════════════════════════════════════════════════════════════════════════════
+
+def detect_engine_parts(image, conf_threshold=0.25, apply_clahe_preprocessing=True):
+    """
+    Detect engine parts with YOLO.
+    
+    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)
+    
+    # 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
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 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; }
+"""
+
+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.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")
+    
+    # Examples
+    gr.Markdown("### 📸 Try these examples:")
+    gr.Examples(
+        examples=[
+            # Add example image paths if you have them
+            # ["example1.jpg", 0.25, True],
+            # ["example2.jpg", 0.3, True],
+        ],
+        inputs=[input_image, conf_slider, clahe_checkbox],
+        outputs=[output_image, output_text],
+        fn=detect_engine_parts,
+        cache_examples=False,
+    )
+    
+    # Info section
+    gr.HTML("""
+    <div style='margin-top: 20px; padding: 15px; background: #f5f5f5; border-radius: 8px;'>
+        <h3 style='margin-top: 0;'>ℹ️ How to Use</h3>
+        <ol>
+            <li>Upload an image of an engine part or use your webcam</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>
+        <p><strong>Supported Classes:</strong> bearing_saddle, piston, defect, crack, corrosion (depends on your training)</p>
+    </div>
+    """)
+    
+    # Wire up the button
+    detect_btn.click(
+        fn=detect_engine_parts,
+        inputs=[input_image, conf_slider, clahe_checkbox],
+        outputs=[output_image, output_text],
+    )
+
+
+# Launch
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+    )