vari fix

README.md

@@ -25,6 +25,9 @@ A comprehensive benchmark for anomaly detection models on the [MVTec AD dataset]
 - **Multiple Models**: PatchCore, EfficientAD, FastFlow, STFPM, PaDiM
 - **Full Benchmark**: Train and evaluate on all 15 MVTec categories
 - **Interactive Demo**: [Gradio UI for real-time anomaly detection](https://huggingface.co/spaces/micguida1/mvtec-anomaly-benchmark)
+- **Sample Image Gallery**: Browse and select sample images from MVTec dataset with automatic category detection
+- **Draw Defects**: Draw artificial defects on images and see how models detect them
+- **Model Comparison**: Compare multiple models side-by-side on the same image
 - **Easy Configuration**: YAML-based model configs
 
 ## 📦 Installation
@@ -105,6 +108,16 @@ python app.py
 
 The demo will be available at `http://localhost:7860`.
 
+#### Demo Features
+
+- **📤 Upload Image**: Upload any image and analyze it for anomalies
+- **✏️ Draw Defects**: Load a sample image and draw artificial defects to test detection
+- **🔄 Compare Models**: Compare multiple models side-by-side on the same image
+- **📚 Learn**: Educational content about each anomaly detection model
+- **📊 Metrics**: View detailed performance metrics for each model
+
+**Sample Image Gallery**: Each tab includes a gallery of sample images from the MVTec dataset. Click on any image to load it and the category will be automatically selected.
+
 ## 📁 Project Structure
 
 ```

gradio_ui/content.py

@@ -20,26 +20,11 @@ Test how well the anomaly detection model captures defects:
 1. **Upload a GOOD image** (normal, without defects)
 2. **Use the brush to draw** artificial defects (scratches, stains, cracks, etc.)
 3. **Click Analyze** to see if the heatmap detects your drawn defects
-
-> 💡 **Tip:** Use different brush sizes and colors to simulate various defect types!
 """
 
-BRUSH_COLORS_INFO = """
----
-**Brush Colors:**
-- ⚫ Black - Dark scratches, contamination
-- 🔴 Red - Highlighted defects
-- 🟤 Brown - Rust, stains
-- ⚪ Gray/White - Light marks
-"""
+BRUSH_COLORS_INFO = ""
 
-HEATMAP_INTERPRETATION = """
----
-**How to interpret:**
-- 🔴 **Red areas** in heatmap = High anomaly probability (model detected something)
-- 🔵 **Blue areas** = Normal regions
-- Compare drawn defects with heatmap response!
-"""
+HEATMAP_INTERPRETATION = ""
 
 # Compare Models Tab Instructions
 COMPARE_MODELS_INSTRUCTIONS = """
@@ -47,10 +32,8 @@ COMPARE_MODELS_INSTRUCTIONS = """
 
 Compare how different anomaly detection models perform on the same image:
 1. **Upload an image** (normal or with defects)
-2. **Select 2-4 models** to compare
+2. **Select 2-5 models** to compare
 3. **Click Compare** to see heatmaps side-by-side
-
-> 💡 **Tip:** This helps identify which model is most sensitive to certain defect types!
 """
 
 # Learn About Models Tab Content

gradio_ui/handlers.py

@@ -164,8 +164,8 @@ def predict_compare(image, selected_models: list, category: str):
     if not selected_models or len(selected_models) < 2:
         return None, "⚠️ Please select at least 2 models to compare"
     
-    if len(selected_models) > 4:
-        return None, "⚠️ Please select at most 4 models"
+    if len(selected_models) > 5:
+        return None, "⚠️ Please select at most 5 models"
     
     try:
         # Save temp image if needed
@@ -204,29 +204,7 @@ def predict_compare(image, selected_models: list, category: str):
         # Create comparison visualization
         viz_image = create_comparison_visualization(original, results_list)
         
-        # Create summary text
-        summary_lines = ["### 📊 Comparison Results\n"]
-        summary_lines.append("| Model | Score | Status |")
-        summary_lines.append("|-------|-------|--------|")
-        
-        for result in results_list:
-            if result.get('error'):
-                summary_lines.append(f"| {result['model_name']} | ❌ Error | {result['error'][:30]}... |")
-            else:
-                score = result['score']
-                model_name = result['model_name']
-                
-                # Scale if EfficientAD
-                display_score = score
-                if "efficientad" in model_name.lower():
-                    display_score = scale_efficientad_score(score)
-                
-                status = "🔴 Anomaly" if display_score > 0.5 else "🟢 Normal"
-                summary_lines.append(f"| {model_name} | {display_score:.4f} | {status} |")
-        
-        summary = "\n".join(summary_lines)
-        
-        return viz_image, summary
+        return viz_image, ""
         
     except Exception as e:
         return None, f"❌ Error: {str(e)}"

gradio_ui/tabs/compare.py

@@ -35,7 +35,7 @@ def create_compare_tab(available_models: list, initial_categories: list):
                     height=300
                 )
                 
-                gr.Markdown("### 🧠 Select Models (2-4)")
+                gr.Markdown("### 🧠 Select Models (2-5)")
                 
                 compare_model_checkboxes = gr.CheckboxGroup(
                     choices=available_models,

gradio_ui/visualization.py

@@ -264,10 +264,12 @@ def create_comparison_visualization(original: np.ndarray, results_list: list) ->
             axes[row, 3].set_title(col_titles[3], fontsize=12, fontweight='bold')
         
         # Add model name as row label on the left using annotation
-        # Removing numerical score as requested
-        axes[row, 0].annotate(f"{model_name}\n{status}", 
+        # Show model name, score and status
+        status_color = 'red' if display_score > 0.5 else 'green'
+        axes[row, 0].annotate(f"{model_name}\nScore: {display_score:.2f}\n{status}", 
                               xy=(-0.1, 0.5), xycoords='axes fraction',
-                              fontsize=11, fontweight='bold', ha='right', va='center')
+                              fontsize=11, fontweight='bold', ha='right', va='center',
+                              color=status_color)
     
     # Add colorbar
     cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])

scripts/download_checkpoints.py

@@ -61,6 +61,21 @@ def get_checkpoint_hf_path(model_name: str, category: str) -> str:
     return f"{dirname}/MVTecAD/{category}/latest/weights/lightning/model.ckpt"
 
 
+def get_metrics_hf_path(model_name: str, category: str) -> str:
+    """
+    Returns the path of the metrics.json file in the HF repository.
+    
+    Args:
+        model_name: Name of the model
+        category: MVTec category
+        
+    Returns:
+        Path string relative to HF repo root
+    """
+    dirname = MODEL_TO_DIRNAME.get(model_name, model_name.capitalize())
+    return f"{dirname}/MVTecAD/{category}/latest/metrics.json"
+
+
 def get_local_checkpoint_path(model_name: str, category: str) -> Path:
     """
     Returns the local path where the checkpoint should be stored.
@@ -76,6 +91,21 @@ def get_local_checkpoint_path(model_name: str, category: str) -> Path:
     return DIR_RESULTS / dirname / "MVTecAD" / category / "latest" / "weights" / "lightning" / "model.ckpt"
 
 
+def get_local_metrics_path(model_name: str, category: str) -> Path:
+    """
+    Returns the local path where the metrics.json should be stored.
+    
+    Args:
+        model_name: Name of the model
+        category: MVTec category
+        
+    Returns:
+        Path object for local metrics file
+    """
+    dirname = MODEL_TO_DIRNAME.get(model_name, model_name.capitalize())
+    return DIR_RESULTS / dirname / "MVTecAD" / category / "latest" / "metrics.json"
+
+
 def download_checkpoint(model_name: str, category: str, force: bool = False) -> bool:
     """
     Downloads a single checkpoint from HuggingFace Hub.
@@ -115,13 +145,52 @@ def download_checkpoint(model_name: str, category: str, force: bool = False) ->
         return False
 
 
+def download_metrics(model_name: str, category: str, force: bool = False) -> bool:
+    """
+    Downloads metrics.json for a model/category from HuggingFace Hub.
+    
+    Args:
+        model_name: Name of the model
+        category: MVTec category
+        force: If True, re-download even if exists
+        
+    Returns:
+        True if downloaded/exists, False if failed
+    """
+    local_path = get_local_metrics_path(model_name, category)
+    
+    # Skip if already exists
+    if local_path.exists() and not force:
+        return True
+    
+    hf_path = get_metrics_hf_path(model_name, category)
+    
+    try:
+        # Create parent directories
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        # Download from HF Hub
+        downloaded_path = hf_hub_download(
+            repo_id=HF_REPO_ID,
+            filename=hf_path,
+            local_dir=DIR_RESULTS,
+            local_dir_use_symlinks=False,
+        )
+        
+        return True
+        
+    except Exception as e:
+        # Metrics file is optional, don't print error
+        return False
+
+
 def download_all_checkpoints(
     models: list[str] = None,
     categories: list[str] = None,
     force: bool = False
 ) -> dict:
     """
-    Downloads checkpoints for specified models and categories.
+    Downloads checkpoints and metrics for specified models and categories.
     
     Args:
         models: List of model names (None = all available)
@@ -136,11 +205,11 @@ def download_all_checkpoints(
     if categories is None:
         categories = MVTEC_CATEGORIES
     
-    stats = {"downloaded": 0, "existed": 0, "failed": 0}
+    stats = {"downloaded": 0, "existed": 0, "failed": 0, "metrics_downloaded": 0}
     
     total = len(models) * len(categories)
     
-    print(f"📦 Downloading checkpoints from: {HF_REPO_ID}")
+    print(f"📦 Downloading checkpoints and metrics from: {HF_REPO_ID}")
     print(f"   Models: {', '.join(models)}")
     print(f"   Categories: {len(categories)} total")
     print()
@@ -156,6 +225,10 @@ def download_all_checkpoints(
                     stats["downloaded"] += 1
                 else:
                     stats["failed"] += 1
+                
+                # Also download metrics.json if available
+                if download_metrics(model, category, force):
+                    stats["metrics_downloaded"] += 1
                     
                 pbar.update(1)
     
@@ -179,6 +252,7 @@ def check_checkpoint_exists(model_name: str, category: str) -> bool:
 def ensure_checkpoint(model_name: str, category: str) -> Path:
     """
     Ensures a checkpoint exists, downloading if necessary.
+    Also downloads metrics.json if available.
     
     This is the main function to call from inference/app code.
     
@@ -195,11 +269,15 @@ def ensure_checkpoint(model_name: str, category: str) -> Path:
     local_path = get_local_checkpoint_path(model_name, category)
     
     if local_path.exists():
+        # Also try to download metrics if not present
+        download_metrics(model_name, category)
         return local_path
     
     print(f"⬇ Checkpoint not found locally. Downloading {model_name}/{category}...")
     
     if download_checkpoint(model_name, category):
+        # Also download metrics
+        download_metrics(model_name, category)
         if local_path.exists():
             print(f"✓ Downloaded successfully")
             return local_path
@@ -258,8 +336,9 @@ def main():
     # Report
     print()
     print("=" * 50)
-    print(f"✓ Downloaded: {stats['downloaded']}")
+    print(f"✓ Checkpoints downloaded: {stats['downloaded']}")
     print(f"○ Already existed: {stats['existed']}")
+    print(f"📊 Metrics downloaded: {stats['metrics_downloaded']}")
     if stats['failed'] > 0:
         print(f"✗ Failed: {stats['failed']}")
     print("=" * 50)