Create app.py

app.py

@@ -0,0 +1,253 @@
+# app.py
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+import numpy as np
+import pickle
+import os
+
+# class EnginePartDetector:
+#     def __init__(self):
+#         self.model = models.resnet50(weights='IMAGENET1K_V1')
+#         self.model = nn.Sequential(*list(self.model.children())[:-1])
+#         self.model.eval()
+        
+#         self.transform = transforms.Compose([
+#             transforms.Resize((224, 224)),
+#             transforms.ToTensor(),
+#             transforms.Normalize(
+#                 mean=[0.485, 0.456, 0.406],
+#                 std=[0.229, 0.224, 0.225]
+#             )
+#         ])
+        
+#         self.templates = {}
+#         self.load_templates()
+    
+#     def extract_features(self, image):
+#         if isinstance(image, np.ndarray):
+#             image = Image.fromarray(image)
+        
+#         img_tensor = self.transform(image).unsqueeze(0)
+        
+#         with torch.no_grad():
+#             features = self.model(img_tensor)
+#             features = features.squeeze().numpy()
+        
+#         return features
+
+class EnginePartDetector:
+    def __init__(
+        self,
+        clahe_clip_limit: float = 9.9,
+        clahe_tile_grid: tuple = (8, 8),
+    ):
+        # ── ResNet-50 backbone (feature extractor only) ──────────────────
+        self.model = models.resnet50(weights='IMAGENET1K_V1')
+        self.model = nn.Sequential(*list(self.model.children())[:-1])
+        self.model.eval()
+
+        # ── CLAHE (OpenCV) — applied BEFORE the torch transform ──────────
+        # Operates on grayscale to recover shadow-suppressed edges
+        # (e.g. missing bearing saddle arcs), then merged back to RGB
+        # so the 3-channel ResNet pipeline is unaffected.
+        self.clahe = cv2.createCLAHE(
+            clipLimit=clahe_clip_limit,
+            tileGridSize=clahe_tile_grid,
+        )
+
+        # ── ResNet normalisation transform (unchanged) ───────────────────
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225],
+            )
+        ])
+
+        self.templates = {}
+        self.load_templates()
+
+    # ── CLAHE preprocessing ───────────────────────────────────────────────
+
+    def apply_clahe(self, image: np.ndarray) -> np.ndarray:
+        
+        # Convert RGB (PIL/numpy) → BGR for OpenCV
+        bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+        # BGR → LAB
+        lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB)
+
+        # Split channels; apply CLAHE only to L (luminance)
+        l_channel, a_channel, b_channel = cv2.split(lab)
+        l_enhanced = self.clahe.apply(l_channel)
+
+        # Merge enhanced L back with untouched A and B
+        lab_enhanced = cv2.merge([l_enhanced, a_channel, b_channel])
+
+        # LAB → BGR → RGB
+        bgr_enhanced = cv2.cvtColor(lab_enhanced, cv2.COLOR_LAB2BGR)
+        rgb_enhanced = cv2.cvtColor(bgr_enhanced, cv2.COLOR_BGR2RGB)
+
+        return rgb_enhanced  # uint8 numpy array, same shape as input
+
+    # ── Feature extraction ────────────────────────────────────────────────
+
+    def extract_features(self, image) -> np.ndarray:
+       
+        # 1. Normalise input to numpy uint8 RGB
+        if isinstance(image, Image.Image):
+            image = np.array(image.convert("RGB"))
+        elif isinstance(image, np.ndarray) and image.dtype != np.uint8:
+            image = image.astype(np.uint8)
+
+        # 2. CLAHE — recover shadow-suppressed structural edges
+        image = self.apply_clahe(image)
+
+        # 3. Mild Gaussian blur — reduces high-freq metallic sheen noise
+        #    that CLAHE can amplify; kernel (3,3) is intentionally light
+        #    so real surface-defect texture is preserved
+        image = cv2.GaussianBlur(image, (3, 3), 0)
+
+        # 4. Convert back to PIL for torchvision transforms
+        image_pil = Image.fromarray(image)
+
+        # 5. ResNet transform → tensor
+        img_tensor = self.transform(image_pil).unsqueeze(0)
+
+        # 6. Forward pass (no grad needed — inference only)
+        with torch.no_grad():
+            features = self.model(img_tensor)
+            features = features.squeeze().numpy()
+
+        return features 
+    
+    def cosine_similarity(self, feat1, feat2):
+        return np.dot(feat1, feat2) / (np.linalg.norm(feat1) * np.linalg.norm(feat2))
+    
+    def save_template(self, image, part_name):
+        if image is None or not part_name:
+            return "Please provide both image and part name"
+        
+        features = self.extract_features(image)
+        self.templates[part_name] = features
+        
+        with open('templates.pkl', 'wb') as f:
+            pickle.dump(self.templates, f)
+        
+        return f"✅ Template '{part_name}' saved successfully!"
+    
+    def load_templates(self):
+        if os.path.exists('templates.pkl'):
+            try:
+                with open('templates.pkl', 'rb') as f:
+                    self.templates = pickle.load(f)
+            except:
+                self.templates = {}
+    
+    def match_part(self, image, threshold=0.7):
+        if image is None:
+            return "Please provide an image", None
+        
+        if not self.templates:
+            return "⚠️ No templates available. Please add templates first.", None
+        
+        query_features = self.extract_features(image)
+        
+        results = []
+        for part_name, template_features in self.templates.items():
+            similarity = self.cosine_similarity(query_features, template_features)
+            results.append((part_name, similarity))
+        
+        results.sort(key=lambda x: x[1], reverse=True)
+        
+        best_match = results[0]
+        output_text = f"🔍 **Best Match**: {best_match[0]}\n"
+        output_text += f"📊 **Confidence**: {best_match[1]:.2%}\n\n"
+        
+        if best_match[1] >= threshold:
+            output_text += "✅ **Status**: MATCHED\n\n"
+        else:
+            output_text += "❌ **Status**: NO MATCH (below threshold)\n\n"
+        
+        output_text += "**All Results:**\n"
+        for part, sim in results:
+            output_text += f"- {part}: {sim:.2%}\n"
+        
+        matched_label = best_match[0] if best_match[1] >= threshold else None
+        return output_text, matched_label
+
+detector = EnginePartDetector()
+
+def add_template(image, part_name):
+    return detector.save_template(image, part_name)
+
+def detect_part(image, threshold):
+    return detector.match_part(image, threshold)
+
+def list_templates():
+    if not detector.templates:
+        return "No templates saved yet"
+    return "\n".join([f"- {name}" for name in detector.templates.keys()])
+
+with gr.Blocks(title="Engine Part Detection System") as demo:
+    gr.Markdown("""
+    # 🔧 Engine Part Detection System
+    ### Using ResNet50 Feature Extraction & Template Matching
+    
+    **How to use:**
+    1. **Add Templates**: Upload reference images of engine parts
+    2. **Detect Parts**: Upload/capture images to identify parts
+    """)
+    
+    with gr.Tab("🔍 Detect Part"):
+        with gr.Row():
+            with gr.Column():
+                detect_input = gr.Image(sources=["upload", "webcam"], type="numpy")
+                threshold_slider = gr.Slider(0.5, 0.95, value=0.7, label="Similarity Threshold")
+                detect_btn = gr.Button("Detect Part", variant="primary")
+            with gr.Column():
+                detect_output = gr.Textbox(label="Detection Results", lines=10)
+                match_label = gr.Label(label="Matched Part")
+        
+        detect_btn.click(
+            fn=detect_part,
+            inputs=[detect_input, threshold_slider],
+            outputs=[detect_output, match_label],
+            api_name="detect"
+        )
+    
+    with gr.Tab("➕ Add Template"):
+        with gr.Row():
+            with gr.Column():
+                template_input = gr.Image(sources=["upload"], type="numpy")
+                part_name_input = gr.Textbox(label="Part Name (e.g., 'spark_plug', 'piston')")
+                add_btn = gr.Button("Save Template", variant="primary")
+            with gr.Column():
+                add_output = gr.Textbox(label="Status")
+        
+        add_btn.click(
+            fn=add_template,
+            inputs=[template_input, part_name_input],
+            outputs=add_output,
+            api_name="add_template"
+        )
+    
+    with gr.Tab("📋 View Templates"):
+        template_list = gr.Textbox(label="Saved Templates", lines=10)
+        refresh_btn = gr.Button("Refresh List")
+        refresh_btn.click(
+            fn=list_templates,
+            outputs=template_list,
+            api_name="list_templates"
+        )
+        demo.load(fn=list_templates, outputs=template_list)
+
+if __name__ == "__main__":
+    demo.launch()
+
+
+app.py