classifucation

app.py

@@ -272,6 +272,7 @@ logger = logging.getLogger(__name__)
 
 class FeatureExtractor:
     def __init__(self):
+        # Using ResNet50 for 2048-D feature vectors
         backbone = models.resnet50(weights="IMAGENET1K_V1")
         self.model = nn.Sequential(*list(backbone.children())[:-1])
         self.model.eval()
@@ -291,15 +292,24 @@ class FeatureExtractor:
             rgb = np.array(rgb.convert("RGB"))
         if rgb.dtype != np.uint8:
             rgb = rgb.astype(np.uint8)
+
         if len(rgb.shape) == 2:
             rgb = cv2.cvtColor(rgb, cv2.COLOR_GRAY2RGB)
-
+        
+        # We want the layer BEFORE the global pooling to get spatial info
+        # resnet.layer4 is the last block
+        # self.model is nn.Sequential(*list(backbone.children())[:-1])
+        # children()[:-1] = [conv1, bn1, relu, maxpool, layer1, layer2, layer3, layer4]
+        
         input_tensor = self.transform(Image.fromarray(rgb)).unsqueeze(0)
-
+        
+        # Get activations from the last conv layer (Layer 4)
         with torch.no_grad():
+            # Run through the layers up to global pooling
+            # Using the original backbone for Easier Access to sub-layers
             backbone = models.resnet50(weights="IMAGENET1K_V1")
             backbone.eval()
-
+            
             x = backbone.conv1(input_tensor)
             x = backbone.bn1(x)
             x = backbone.relu(x)
@@ -307,23 +317,27 @@ class FeatureExtractor:
             x = backbone.layer1(x)
             x = backbone.layer2(x)
             x = backbone.layer3(x)
-            features_spatial = backbone.layer4(x)  # [1, 2048, 7, 7]
-
+            features_spatial = backbone.layer4(x) # [1, 2048, 7, 7]
+            
+            # Global Average Pooling to get the vector
             feat = torch.mean(features_spatial, dim=[2, 3]).squeeze().cpu().numpy()
-
+            
+            # Create Heatmap: sum across channels to see "hot" regions
             amap = torch.sum(features_spatial, dim=1).squeeze().cpu().numpy()
             amap = np.maximum(amap, 0)
             amap /= (np.max(amap) + 1e-8)
             amap = cv2.resize(amap, (rgb.shape[1], rgb.shape[0]))
             amap = np.uint8(255 * amap)
             heatmap = cv2.applyColorMap(amap, cv2.COLORMAP_JET)
+            
+            # Overlay heatmap on original image
+            # Convert BGR heatmap to RGB
             heatmap_rgb = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
             overlay = cv2.addWeighted(rgb, 0.6, heatmap_rgb, 0.4, 0)
 
         norm = np.linalg.norm(feat)
         return (feat / norm if norm > 1e-8 else feat), overlay
 
-
 # ───────────────────────────────────────────────────────────────────────────────
 # MASTER ORCHESTRATOR │ EnginePartDetector
 # ───────────────────────────────────────────────────────────────────────────────
@@ -333,87 +347,118 @@ class EnginePartDetector:
 
     def __init__(self):
         self.feature_extractor = FeatureExtractor()
-        self.templates: dict[str, np.ndarray] = {}
-        self._load_templates()
+        # Changed from simple templates to class-based feature lists
+        self.classes: dict[str, list[np.ndarray]] = {}
+        # We also store an example ROI for each class for visualization
+        self.class_rois: dict[str, np.ndarray] = {}
+        self._load_data()
 
     # ── Persistence ───────────────────────────────────────────────────────────
 
-    def _load_templates(self) -> None:
+    def _load_data(self) -> None:
         if os.path.exists(self.TEMPLATE_FILE):
             try:
                 with open(self.TEMPLATE_FILE, "rb") as f:
-                    self.templates = pickle.load(f)
-                logger.info(f"Loaded {len(self.templates)} template(s).")
+                    data = pickle.load(f)
+                    # Support legacy format if needed, but here we assume the new format
+                    if isinstance(data, dict):
+                        # If old format was {name: {"features": feat, "roi": roi}}
+                        # we convert it to {name: [feat]}
+                        self.classes = {}
+                        self.class_rois = {}
+                        for k, v in data.items():
+                            if isinstance(v, dict) and "features" in v:
+                                self.classes[k] = [v["features"]]
+                                self.class_rois[k] = v.get("roi")
+                            else:
+                                self.classes[k] = v
+                    else:
+                        self.classes = {}
+                logger.info(f"Loaded {len(self.classes)} class(es).")
             except Exception as e:
-                logger.error(f"Template load failed: {e}")
-                self.templates = {}
+                logger.error(f"Data load failed: {e}")
+                self.classes = {}
 
-    def _persist_templates(self) -> None:
+    def _persist_data(self) -> None:
         try:
             with open(self.TEMPLATE_FILE, "wb") as f:
-                pickle.dump(self.templates, f)
+                pickle.dump(self.classes, f)
+            # Separately save ROIs if needed, but for now we just persist classes
+            # In a real app we'd save ROIs too. Let's include them in a combined dict.
+            with open("class_data.pkl", "wb") as f:
+                pickle.dump({"classes": self.classes, "rois": self.class_rois}, f)
         except Exception as e:
-            logger.error(f"Template save failed: {e}")
+            logger.error(f"Data save failed: {e}")
 
     # ── Layer 1: ROI Detection & Extraction ───────────────────────────────────
 
     @staticmethod
-    def detect_connect_and_crop(
-        image_source: np.ndarray,
-    ) -> tuple[np.ndarray, np.ndarray, str, list, list]:
-       
+    def detect_connect_and_crop(image_source: np.ndarray) -> tuple[np.ndarray, np.ndarray, str]:
+        """
+        1. Detects bolt holes.
+        2. Separates into Top and Bottom rows.
+        3. Fits horizontal reference lines.
+        4. Crops the FULL horizontal band between rows (includes regions between saddles).
+        """
         img_rgb = image_source
         img_h, img_w = img_rgb.shape[:2]
         gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
         gray = cv2.GaussianBlur(gray, (7, 7), 0)
-
+        
+        # ── Step 1: Detect Circles ────────────────────────────────────────────
         circles = cv2.HoughCircles(
             gray, cv2.HOUGH_GRADIENT, dp=1.2, minDist=60,
             param1=100, param2=35, minRadius=12, maxRadius=45
         )
-
+        
         if circles is None:
-            return img_rgb, img_rgb, "❌ No bolt holes detected.", [], []
-
+            return img_rgb, img_rgb, "❌ No bolt holes detected."
+        
         circles = np.round(circles[0]).astype(int)
-
+        
+        # ── Step 2: Row Separation ────────────────────────────────────────────
         ys = sorted([c[1] for c in circles])
         y_median = np.median(ys)
-
+        
         top_row = sorted([c for c in circles if c[1] < y_median], key=lambda x: x[0])
         bot_row = sorted([c for c in circles if c[1] >= y_median], key=lambda x: x[0])
-
+        
         if len(top_row) < 2 or len(bot_row) < 2:
-            return img_rgb, img_rgb, "⚠️ Insufficient hole rows for localization.", [], []
+            return img_rgb, img_rgb, "⚠️ Insufficient hole rows for localization."
 
+        # ── Step 3: Reference Lines ───────────────────────────────────────────
         y_top = int(np.mean([c[1] for c in top_row]))
         y_bot = int(np.mean([c[1] for c in bot_row]))
-
+        
+        # Horizontal bounds (First hole to Last hole)
         xs = [c[0] for c in circles]
         x_min, x_max = min(xs), max(xs)
         padding_h = 60
         padding_v = 20
-
+        
         x_start = max(0, x_min - padding_h)
         x_end = min(img_w, x_max + padding_h)
         y_start = max(0, min(y_top, y_bot) - padding_v)
         y_end = min(img_h, max(y_top, y_bot) + padding_v)
 
+        # ── Step 4: Visualization ─────────────────────────────────────────────
         vis_img = img_rgb.copy()
         LINE_COLOR = (0, 255, 0)
         HOLE_COLOR = (255, 0, 0)
-
+        
+        # Draw lines and detected holes
         cv2.line(vis_img, (0, y_top), (img_w, y_top), LINE_COLOR, 3)
         cv2.line(vis_img, (0, y_bot), (img_w, y_bot), LINE_COLOR, 3)
-
+        
         for (x, y, r) in circles:
             cv2.circle(vis_img, (x, y), r, HOLE_COLOR, 3)
             cv2.circle(vis_img, (x, y), 2, (255, 255, 255), -1)
 
+        # ── Step 5: Full Band Crop ────────────────────────────────────────────
         cropped_img = img_rgb[y_start:y_end, x_start:x_end]
-
+        
         if cropped_img.size == 0:
-            return vis_img, img_rgb, "⚠️ ROI selection failed.", [], []
+            return vis_img, img_rgb, "⚠️ ROI selection failed."
 
         stats_text = (
             f"✅ **Full Saddle Band Extracted**\n"
@@ -423,72 +468,22 @@ class EnginePartDetector:
             f"• ROI Size: {cropped_img.shape[1]}x{cropped_img.shape[0]} px"
         )
 
-        return vis_img, cropped_img, stats_text, list(top_row), list(bot_row)
-
-    # ─�� Vertical-line detection on structural edge map ───────────────────────
-
-    @staticmethod
-    def detect_vertical_lines_on_edge_map(
-        roi_enhanced: np.ndarray,
-        angle_tolerance_deg: float = 12.0,
-        min_line_length_ratio: float = 0.15,
-    ) -> tuple[bool, np.ndarray, str]:
-       
-        gray  = cv2.cvtColor(roi_enhanced, cv2.COLOR_RGB2GRAY)
-        clahe = cv2.createCLAHE(clipLimit=2.8, tileGridSize=(8, 8))
-        gray  = clahe.apply(gray)
-        edges = cv2.Canny(gray, 50, 150)
-
-        h, w      = edges.shape
-        min_len   = max(20, int(h * min_line_length_ratio))
-
-        lines = cv2.HoughLinesP(
-            edges, rho=1, theta=np.pi / 180,
-            threshold=20, minLineLength=min_len, maxLineGap=10,
-        )
-
-        # RGB canvas from edge map
-        canvas = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
-
-        vertical_lines = []
-        if lines is not None:
-            for seg in lines:
-                x1, y1, x2, y2 = seg[0]
-                dx    = abs(x2 - x1)
-                dy    = abs(y2 - y1)
-                angle = np.degrees(np.arctan2(dx, dy + 1e-6))
-                if angle <= angle_tolerance_deg:
-                    vertical_lines.append((x1, y1, x2, y2, dy))
-
-        # Sort by length — longest first
-        vertical_lines.sort(key=lambda v: v[4], reverse=True)
-        has_vertical = len(vertical_lines) > 0
-
-        if has_vertical:
-            for (x1, y1, x2, y2, _) in vertical_lines:
-                cv2.line(canvas, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            cv2.rectangle(canvas, (0, 0), (240, 46), (0, 150, 0), -1)
-            cv2.putText(canvas, f"PRESENT  ({len(vertical_lines)})",
-                        (6, 34), cv2.FONT_HERSHEY_DUPLEX, 0.85, (255, 255, 255), 2)
-            status = (f"✅ **Vertical lines PRESENT** — "
-                      f"{len(vertical_lines)} near-vertical line(s) detected.")
-        else:
-            cv2.rectangle(canvas, (0, 0), (190, 46), (180, 0, 0), -1)
-            cv2.putText(canvas, "ABSENT",
-                        (6, 34), cv2.FONT_HERSHEY_DUPLEX, 1.1, (255, 255, 255), 2)
-            status = "❌ **Vertical lines ABSENT** — No near-vertical lines on edge map."
-
-        return has_vertical, canvas, status
+        return vis_img, cropped_img, stats_text
 
     @staticmethod
     def enhance_roi(roi: np.ndarray) -> np.ndarray:
         """Apply high-contrast CLAHE to highlight blurred lines/features."""
         if roi is None or roi.size == 0:
             return roi
+        
+        # Convert to LAB space to apply CLAHE on L (luminance) channel
         lab = cv2.cvtColor(roi, cv2.COLOR_RGB2LAB)
         l, a, b = cv2.split(lab)
-        clahe = cv2.createCLAHE(clipLimit=3.9, tileGridSize=(8, 8))
+        
+        # ClipLimit 10.0 provides very high contrast as requested
+        clahe = cv2.createCLAHE(clipLimit=10.0, tileGridSize=(8, 8))
         cl = clahe.apply(l)
+        
         merged = cv2.merge((cl, a, b))
         enhanced = cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)
         return enhanced
@@ -504,104 +499,112 @@ class EnginePartDetector:
 
     # ── Public API ────────────────────────────────────────────────────────────
 
-    def save_template(self, image: np.ndarray, part_name: str) -> tuple[str, np.ndarray | None]:
+    def add_to_class(self, image: np.ndarray, class_name: str) -> tuple[str, np.ndarray | None]:
         if image is None:
             return "❌ No image supplied.", None
-        if not part_name or not part_name.strip():
-            return "❌ Part name is empty.", None
+        if not class_name or not class_name.strip():
+            return "❌ Class name is empty.", None
 
-        part_name = part_name.strip()
-
-        vis, roi, log, top_row, bot_row = self.detect_connect_and_crop(image)
+        class_name = class_name.strip()
+        
+        # Layer 1: Localization
+        vis, roi, log = self.detect_connect_and_crop(image)
         if "❌" in log or "⚠️" in log:
             return log, None
 
+        # Enhance ROI
         roi_enhanced = self.enhance_roi(roi)
+
+        # Layer 2: Feature Extraction
         features, _ = self.feature_extractor.extract(roi_enhanced)
-        self.templates[part_name] = {
-            "features": features,
-            "roi": roi_enhanced,
-        }
-        self._persist_templates()
+        
+        if class_name not in self.classes:
+            self.classes[class_name] = []
+        
+        self.classes[class_name].append(features)
+        self.class_rois[class_name] = roi_enhanced # Keep the latest ROI as reference
+        
+        self._persist_data()
 
-        return f"✅ Template '{part_name}' saved!\n\n{log}", roi
+        return f"✅ Image added to class '{class_name}'! (Now has {len(self.classes[class_name])} samples)\n\n{log}", roi
 
     def match_part(
         self,
         image: np.ndarray,
         threshold: float = 0.70,
-    ) -> tuple[str, dict | None, np.ndarray | None, np.ndarray | None, np.ndarray | None]:
-        """
-        Returns:
-            report_text, label_dict, field_vis, attention_map, annotated_edge_map
-        """
+    ) -> tuple[str, dict | None, np.ndarray | None, np.ndarray | None]:
         if image is None:
-            return "❌ No image supplied.", None, None, None, None
-        if not self.templates:
-            return "⚠️ No templates yet. Add at least one template first.", None, None, None, None
+            return "❌ No image supplied.", None, None, None
+        if not self.classes:
+            return "⚠️ No trained classes yet. Add samples to at least one class (e.g. 'Perfect').", None, None, None
 
-        # ── Layer 1: ROI ──────────────────────────────────────────────────────
-        vis, roi, log, top_row, bot_row = self.detect_connect_and_crop(image)
+        # Layer 1: Localization
+        vis, roi, log = self.detect_connect_and_crop(image)
         if "❌" in log or "⚠️" in log:
-            return log, None, vis, None, None
+            return log, None, vis, None
 
-        # ── Layer 2: Feature Matching ─────────────────────────────────────────
+        # Enhance ROI
         roi_enhanced = self.enhance_roi(roi)
+
+        # Layer 2: Feature Extraction
         query_feat, attention_map = self.feature_extractor.extract(roi_enhanced)
 
-        scores = []
-        for name, data in self.templates.items():
-            sim = self._cosine(query_feat, data["features"])
-            scores.append((name, sim))
-        scores.sort(key=lambda x: x[1], reverse=True)
-
-        best_name, best_score = scores[0]
-        feat_matched = best_score >= threshold
-
-        # ── Structural Edge Map + Vertical-Line Detection ─────────────────────
-        has_vertical, edge_annotated, vline_status = self.detect_vertical_lines_on_edge_map(roi_enhanced)
-
-        # ── Final verdict: feature match only (vertical line is informational) ─
-        final_pass = feat_matched
-        final_icon = "✅ PASS" if final_pass else "❌ FAIL"
-
-        report_lines = [
-            f"## 🏁 Final Verdict: {final_icon}",
-            f"",
-            f"### 🔍 Feature Match",
-            f"{'✅' if feat_matched else '❌'} **Best Match**: `{best_name}`",
-            f"📊 **Confidence**: {best_score:.2%}",
-            f"",
-            f"### 📐 Vertical Line Analysis (Edge Map)",
-            vline_status,
-            f"",
-            f"---",
-            f"",
-            f"### 📸 Field Detection",
+        # Layer 3: Latent Space Matching (Cosine Similarity to centroids)
+        class_scores = []
+        for name, vectors in self.classes.items():
+            # Calculate centroid (neighborhood center)
+            centroid = np.mean(vectors, axis=0)
+            sim = self._cosine(query_feat, centroid)
+            class_scores.append((name, sim))
+        
+        class_scores.sort(key=lambda x: x[1], reverse=True)
+
+        best_class, best_score = class_scores[0]
+        matched = best_score >= threshold
+        status  = f"✅ CLASSIFIED AS: {best_class}" if matched else "❌ UNCERTAIN (below threshold)"
+
+        lines = [
+            f"{'✅' if matched else '❌'} **Top Prediction**: `{best_class}`",
+            f"📊 **Cosine Similarity**: {best_score:.2%}",
+            f"🎯 **Status**: {status}",
+            "",
+            "### 🔍 Multi-Stage Architecture Analysis",
+            "1. **Localization**: Bolt holes detected, horizontal band cropped.",
+            "2. **Feature Extraction**: ResNet50 extracted unique mathematical fingerprint.",
+            "3. **Matching**: Nearest cluster identified in latent space via Cosine Similarity.",
+            "",
+            "The heatmap on the right shows exactly where the AI is focusing.",
+            "- **Red Regions**: Areas defining the class (e.g., surface quality, edges).",
+            "",
+            "---",
             log,
         ]
+        
+        if len(class_scores) > 1:
+            lines.append("\n**Class Probabilities (Latent Distance):**")
+            for name, sim in class_scores:
+                lines.append(f"  • `{name}`: {sim:.3f}")
 
-        if len(scores) > 1:
-            report_lines.append("\n**Other Template Scores:**")
-            for name, sim in scores[1:5]:
-                report_lines.append(f"  • `{name}`: {sim:.3f}")
-
-        label_dict = {name: float(sim) for name, sim in scores[:5]}
-
-        return "\n".join(report_lines), label_dict, vis, attention_map, edge_annotated
+        label_dict = {name: float(sim) for name, sim in class_scores}
+        
+        # Edge Map for structural analysis
+        gray_enhanced = cv2.cvtColor(roi_enhanced, cv2.COLOR_RGB2GRAY)
+        edges = cv2.Canny(gray_enhanced, 50, 150)
+        edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
+        
+        return "\n".join(lines), label_dict, vis, attention_map, edges_rgb
 
-    def get_template_roi(self, part_name: str) -> np.ndarray | None:
-        if part_name in self.templates:
-            return self.templates[part_name].get("roi")
-        return None
+    def get_template_roi(self, class_name: str) -> np.ndarray | None:
+        return self.class_rois.get(class_name)
 
     def list_templates(self) -> str:
-        if not self.templates:
-            return "No templates saved yet."
-        header = f"Total: {len(self.templates)} template(s)\n" + "─" * 30
-        body   = "\n".join(f"  • {n}" for n in sorted(self.templates.keys()))
-        return f"{header}\n{body}"
-
+        if not self.classes:
+            return "No classes trained yet."
+        header = f"Total: {len(self.classes)} class(es)\n" + "─" * 30
+        body = []
+        for name, vectors in sorted(self.classes.items()):
+            body.append(f"  • {name}: {len(vectors)} samples")
+        return f"{header}\n" + "\n".join(body)
 
 # ───────────────────────────────────────────────────────────────────────────────
 # Gradio Application
@@ -609,54 +612,44 @@ class EnginePartDetector:
 
 detector = EnginePartDetector()
 
-
 def detect_part(image, threshold):
     return detector.match_part(image, threshold)
 
+def add_sample(image, class_name):
+    return detector.add_to_class(image, class_name)
 
-def add_template(image, part_name):
-    return detector.save_template(image, part_name)
-
-
-def list_templates():
+def list_classes():
     return detector.list_templates()
 
-
+# Custom CSS for premium look
 custom_css = """
-.container { max-width: 1400px; margin: auto; }
-.header    { text-align: center; margin-bottom: 2rem; }
-.footer    { text-align: center; margin-top: 2rem; color: #666; }
+.container { max-width: 1200px; margin: auto; }
+.header { text-align: center; margin-bottom: 2rem; }
+.footer { text-align: center; margin-top: 2rem; color: #666; }
 """
 
 with gr.Blocks(title="Engine Part CV System", theme=gr.themes.Soft(), css=custom_css) as demo:
     gr.Markdown("""
     <div class="header">
-        <h1>🔧 Engine Part Detection System</h1>
-        <p>
-            <strong>Layer 1:</strong> Hough Bolt-Hole Detection &amp; Crop &nbsp;|&nbsp;
-            <strong>Layer 2:</strong> ResNet50 Feature Matching &nbsp;|&nbsp;
-            <strong>Edge Map:</strong> Vertical-Line Detection
-        </p>
+        <h1>🔧 Engine Part CV System</h1>
+        <p><strong>Multi-Stage Architecture:</strong> Localization → Feature Fingerprint (ResNet) → Latent Space Matching</p>
     </div>
     """)
 
     with gr.Tab("🔍 Match Inspection"):
         with gr.Row():
             with gr.Column(scale=1):
-                detect_input     = gr.Image(sources=["upload", "webcam"], type="numpy", label="Input Image")
+                detect_input = gr.Image(sources=["upload", "webcam"], type="numpy", label="Input Image")
                 threshold_slider = gr.Slider(0.5, 0.99, value=0.75, step=0.01, label="Matching Threshold")
-                detect_btn       = gr.Button("🔍 Run Inspection", variant="primary")
+                detect_btn = gr.Button("🔍 Run Inspection", variant="primary")
 
             with gr.Column(scale=1):
                 detect_output = gr.Markdown(label="Match Report")
-                match_label   = gr.Label(label="Top Scores", num_top_classes=5)
-
-        with gr.Row():
-            vis_output  = gr.Image(label="Field Visualization (bolt holes)")
-            attn_output = gr.Image(label="AI Attention Heatmap")
-
-        with gr.Row():
-            edge_output = gr.Image(label="Structural Edge Map  (green lines = vertical PRESENT  |  red banner = ABSENT)")
+                match_label = gr.Label(label="Top Scores", num_top_classes=5)
+                with gr.Row():
+                    vis_output = gr.Image(label="Field Visualization")
+                    attn_output = gr.Image(label="AI Attention Heatmap")
+                edge_output = gr.Image(label="Structural Edge Map (Line Detection)")
 
         detect_btn.click(
             fn=detect_part,
@@ -665,39 +658,51 @@ with gr.Blocks(title="Engine Part CV System", theme=gr.themes.Soft(), css=custom
             api_name="detect_part",
         )
 
-    with gr.Tab("💾 Add Golden Template"):
+    with gr.Tab("💾 Train Latent Space"):
         with gr.Row():
             with gr.Column(scale=1):
-                template_input  = gr.Image(sources=["upload"], type="numpy", label="Reference Image")
-                part_name_input = gr.Textbox(label="Part Name", placeholder="e.g. bearing_cap_v8_A")
-                add_btn         = gr.Button("💾 Register Template", variant="primary")
+                template_input = gr.Image(sources=["upload"], type="numpy", label="Training Image")
+                class_name_input = gr.Dropdown(
+                    choices=["Perfect", "Defected", "Unknown"], 
+                    label="Class Label", 
+                    value="Perfect",
+                    allow_custom_value=True
+                )
+                add_btn = gr.Button("💾 Add to Cluster", variant="primary")
             with gr.Column(scale=1):
-                add_status   = gr.Textbox(label="Registration Status", lines=5)
-                add_roi_view = gr.Image(label="Registered Cropped ROI", interactive=False)
+                add_status = gr.Textbox(label="Training Status", lines=5)
+                add_roi_view = gr.Image(label="Processed Training ROI", interactive=False)
 
         add_btn.click(
-            fn=add_template,
-            inputs=[template_input, part_name_input],
+            fn=add_sample,
+            inputs=[template_input, class_name_input],
             outputs=[add_status, add_roi_view],
-            api_name="add_template",
+            api_name="add_sample",
         )
 
-    with gr.Tab("📋 Library"):
+    with gr.Tab("📋 Class Library"):
         with gr.Row():
             with gr.Column(scale=1):
-                template_list = gr.Textbox(label="Current Golden Templates", lines=12)
-                refresh_btn   = gr.Button("🔄 Refresh Library")
+                template_list = gr.Textbox(label="Current Trained Classes", lines=12)
+                refresh_btn = gr.Button("🔄 Refresh Clusters")
             with gr.Column(scale=1):
-                library_roi_view = gr.Image(label="Template ROI Preview", interactive=False)
-
+                library_roi_view = gr.Image(label="Last Reference ROI", interactive=False)
+                
         def update_library_preview():
-            if detector.templates:
-                first_name = sorted(detector.templates.keys())[0]
+            if detector.classes:
+                first_name = sorted(detector.classes.keys())[0]
                 return detector.list_templates(), detector.get_template_roi(first_name)
-            return "No templates saved yet.", None
+            return "No classes trained yet.", None
 
         refresh_btn.click(fn=update_library_preview, outputs=[template_list, library_roi_view])
         demo.load(fn=update_library_preview, outputs=[template_list, library_roi_view])
 
+    gr.Markdown("""
+    ---
+    <div class="footer">
+        <p>Engine Part CV System • Powered by PyTorch & OpenCV</p>
+    </div>
+    """)
+
 if __name__ == "__main__":
-    demo.launch(share=False, show_error=True)
+    demo.launch(share=False, show_error=True)