Refactor

inference.py

@@ -32,10 +32,6 @@ from core import (
 )
 
 
-# =============================================================================
-# INFERENCE
-# =============================================================================
-
 def run_inference(image_path: str, category: str = "bottle", model_name: str = "patchcore", checkpoint_path: str = None) -> dict:
     """
     Runs inference on a single image.
@@ -54,18 +50,15 @@ def run_inference(image_path: str, category: str = "bottle", model_name: str = "
     if not os.path.exists(ckpt):
         raise FileNotFoundError(f"Checkpoint not found: {ckpt}")
     
-    # Load model and dataset
     model = load_model(model_name)
     dataset = PredictDataset(path=image_path)
     
-    # Run prediction (disable automatic image saving)
     engine = Engine(
         default_root_dir="/tmp/anomalib_inference",
         callbacks=[],
     )
     predictions = engine.predict(model=model, dataset=dataset, ckpt_path=ckpt)
     
-    # Extract results
     results = {"image_path": image_path, "category": category, "model": model_name}
     
     for batch in predictions:
@@ -77,10 +70,6 @@ def run_inference(image_path: str, category: str = "bottle", model_name: str = "
     return results
 
 
-# =============================================================================
-# VISUALIZATION
-# =============================================================================
-
 def visualize_results(results: dict, output_dir: Path = None) -> str:
     """
     Visualizes heatmap and mask overlayed on original image.
@@ -92,15 +81,12 @@ def visualize_results(results: dict, output_dir: Path = None) -> str:
     Returns:
         Path of saved image
     """
-    # Get model name for subfolder
     model_name = results.get("model", "unknown")
     
-    # Create output/model_name subfolder
     base_output_dir = output_dir or DIR_OUTPUT
     output_dir = base_output_dir / model_name
     output_dir.mkdir(parents=True, exist_ok=True)
     
-    # Load data
     image_path = results["image_path"]
     anomaly_score = results["anomaly_score"]
     anomaly_map = results["anomaly_map"]
@@ -109,72 +95,54 @@ def visualize_results(results: dict, output_dir: Path = None) -> str:
     
     original = np.array(Image.open(image_path).convert("RGB"))
     
-    # Prepare anomaly map
     if anomaly_map.ndim == 3:
         anomaly_map = anomaly_map.squeeze(0)
     
-    # Adaptive normalization: use image-level normalization if values are compressed
     is_efficientad = model_name.lower() == "efficientad"
     
-    # Adaptive normalization: use image-level normalization if values are compressed
-    # STRICTLY RESTRICTED TO EFFICIENTAD
     amap_min, amap_max = anomaly_map.min(), anomaly_map.max()
     amap_range = amap_max - amap_min
     
-    if is_efficientad and amap_range < 0.1:  # Compressed values (e.g., EfficientAD output)
-        # Apply image-level min-max normalization
+    if is_efficientad and amap_range < 0.1:
         if amap_range > 1e-8:
             anomaly_map = (anomaly_map - amap_min) / amap_range
         else:
             anomaly_map = np.zeros_like(anomaly_map)
-        print(f"ℹ️  Applied image-level normalization (original range: {amap_min:.4f}-{amap_max:.4f})")
+        print(f"[INFO] Applied image-level normalization (original range: {amap_min:.4f}-{amap_max:.4f})")
     else:
-        # Standard clipping for all other models
         anomaly_map = np.clip(anomaly_map, 0, 1)
     
     anomaly_map = resize_to_match(anomaly_map, original.shape[:2])
     
-    # EfficientAD-specific: for "good" images, force blue heatmap
-    # is_efficientad already defined above
     is_good = anomaly_score is not None and anomaly_score < 0.5
-    show_mask_contours = True  # Whether to draw contours on mask panel
+    show_mask_contours = True
     
     if is_efficientad and is_good:
-        # Scale heatmap to low values but keep texture variation (0-0.3 range)
-        anomaly_map = anomaly_map * 0.3  # Keep variation but in low range
-        show_mask_contours = False  # Don't draw contours for good images
-        print(f"ℹ️  EfficientAD: Good image detected - scaling heatmap to low values")
+        anomaly_map = anomaly_map * 0.3
+        show_mask_contours = False
+        print(f"[INFO] EfficientAD: Good image detected - scaling heatmap to low values")s")
     
-    # Prepare pred mask
     if pred_mask is not None:
         if pred_mask.ndim == 3:
             pred_mask = pred_mask.squeeze(0)
         pred_mask = resize_to_match(pred_mask, original.shape[:2])
     
-    # Create figure - always 4 columns if pred_mask exists or EfficientAD good image
     show_fourth_panel = pred_mask is not None or (is_efficientad and is_good)
     num_cols = 4 if show_fourth_panel else 3
     fig, axes = plt.subplots(1, num_cols, figsize=(5 * num_cols, 5), facecolor='white')
     
-    # Set white background for all axes to avoid colored borders
     for ax in axes:
         ax.set_facecolor('white')
     
-    # 1. Original Image
     axes[0].imshow(original)
     axes[0].set_title("Original")
     axes[0].axis("off")
     
-    # 2. Heatmap - use aspect='auto' to fill subplot without borders
-    # 2. Heatmap - use aspect='auto' to fill subplot without borders
-    
     if is_efficientad:
-        # Mask 0 values (padding) to remove blue border
         anomaly_map_masked = np.ma.masked_where(anomaly_map == 0, anomaly_map)
         cmap = plt.cm.jet
-        cmap.set_bad(color='none')  # Transparent for masked values (padding)
+        cmap.set_bad(color='none')
     else:
-        # Standard visualization for other models (0 is Blue, not transparent)
         anomaly_map_masked = anomaly_map
         cmap = plt.cm.jet
     
@@ -183,14 +151,11 @@ def visualize_results(results: dict, output_dir: Path = None) -> str:
     axes[1].axis("off")
     plt.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04)
 
-    
-    # 3. Overlay - use aspect='auto' to match original image
     axes[2].imshow(original, aspect='auto')
     axes[2].imshow(anomaly_map_masked, cmap=cmap, alpha=0.5, vmin=0, vmax=1, aspect='auto')
     axes[2].set_title("Overlay")
     axes[2].axis("off")
     
-    # 4. Mask panel
     if show_fourth_panel:
         axes[3].imshow(original)
         if show_mask_contours and pred_mask is not None:
@@ -198,11 +163,8 @@ def visualize_results(results: dict, output_dir: Path = None) -> str:
         axes[3].set_title("Predicted Mask")
         axes[3].axis("off")
     
-    # Title and Save - scale score to push towards 0 or 1 for EfficientAD
     if anomaly_score is not None:
         if is_efficientad:
-            # Scale EfficientAD scores (typically ~0.5) to be more extreme
-            # Good (<0.5) -> ~0.2, Anomalous (>0.5) -> high values
             scaled_score = scale_efficientad_score(anomaly_score)
             score_str = f"{scaled_score:.4f}"
         else:
@@ -217,16 +179,11 @@ def visualize_results(results: dict, output_dir: Path = None) -> str:
     plt.savefig(output_path, dpi=150, bbox_inches="tight")
     plt.show()
     
-    # Always close the figure
     plt.close(fig)
     
     return str(output_path)
 
 
-# =============================================================================
-# MAIN
-# =============================================================================
-
 def parse_args():
     """Parse command line arguments."""
     parser = argparse.ArgumentParser(
@@ -265,14 +222,12 @@ Examples:
 def main():
     args = parse_args()
     
-    # Validate input
     if not os.path.exists(args.image_path):
         raise FileNotFoundError(f"Image not found: {args.image_path}")
     
-    print(f"🔍 Inference on: {args.image_path}")
-    print(f"📦 Model: {args.model} | Category: {args.category}")
+    print(f"Inference on: {args.image_path}")
+    print(f"Model: {args.model} | Category: {args.category}")
     
-    # Run inference
     results = run_inference(
         image_path=args.image_path,
         category=args.category,
@@ -280,13 +235,11 @@ def main():
         checkpoint_path=args.checkpoint
     )
     
-    # Print result
     score = results["anomaly_score"]
     print(f"\n{'='*50}")
     print(f"ANOMALY SCORE: {score:.4f}" if score else "ANOMALY SCORE: N/A")
     print(f"{'='*50}\n")
     
-    # Visualize and save
     output_dir = Path(args.output_dir) if args.output_dir else None
     output_path = visualize_results(results, output_dir)
     

train.py

@@ -34,12 +34,6 @@ from core import (
 
 logger = logging.getLogger(__name__)
 
-# =============================================================================
-# EFFICIENTAD MONKEY-PATCH
-# =============================================================================
-# Override the pretrained weights directory to keep everything inside
-# efficientad_resources/ instead of the hardcoded ./pre_trained/
-
 EFFICIENTAD_RESOURCES_DIR = Path(__file__).parent / "efficientad_resources"
 
 def _patched_prepare_pretrained_model(self) -> None:
@@ -66,12 +60,8 @@ def patch_efficientad():
     """Apply monkey-patch to EfficientAd to use custom pretrained weights directory."""
     from anomalib.models import EfficientAd
     EfficientAd.prepare_pretrained_model = _patched_prepare_pretrained_model
-    print(f"   ℹ️  EfficientAd: Pretrained weights will be saved to {EFFICIENTAD_RESOURCES_DIR / 'pre_trained'}")
-
+    print(f"   [INFO] EfficientAd: Pretrained weights directory: {EFFICIENTAD_RESOURCES_DIR / 'pre_trained'}")
 
-# =============================================================================
-# METRICS FUNCTIONS
-# =============================================================================
 
 def save_metrics(category_metrics, category, model_name):
     """Saves metrics in the Anomalib directory structure."""
@@ -95,7 +85,7 @@ def save_metrics(category_metrics, category, model_name):
     version_json_path = version_dir / "metrics.json"
     with open(version_json_path, 'w', encoding='utf-8') as f:
         json.dump(category_metrics, f, indent=2, ensure_ascii=False)
-    print(f"   💾 Saved: {version_json_path}")
+    print(f"   Saved: {version_json_path}")
     
     # Save in latest (only if it exists)
     latest_dir = category_base_dir / "latest"
@@ -107,7 +97,7 @@ def save_metrics(category_metrics, category, model_name):
 
 def print_category_metrics(metrics):
     """Prints metrics for a category."""
-    print(f"\n📊 Metrics:")
+    print(f"\n[METRICS]")
     print(f"   EFFICACY:   AUROC img={format_metric(metrics['image_auroc'])} | "
           f"AUROC pix={format_metric(metrics['pixel_auroc'])} | "
           f"F1={format_metric(metrics['image_f1'])}")
@@ -123,7 +113,7 @@ def print_final_report(all_metrics, model_name):
         return
         
     print(f"\n{'='*100}")
-    print(f"📊 FINAL REPORT - {model_name.upper()} PERFORMANCE METRICS")
+    print(f"FINAL REPORT - {model_name.upper()} PERFORMANCE METRICS")
     print(f"{'='*100}\n")
     
     # Header
@@ -157,14 +147,10 @@ def print_final_report(all_metrics, model_name):
     print(f"\n{'='*100}")
 
 
-# =============================================================================
-# MAIN TRAINING FUNCTION
-# =============================================================================
-
 def train_category(category, model_name):
     """Runs training, test, and calculates metrics for a category."""
     print(f"\n{'='*60}")
-    print(f"🚀 Training: {category} ({model_name})")
+    print(f"Training: {category} ({model_name})")
     print(f"{'='*60}")
     
     # Load config
@@ -181,11 +167,10 @@ def train_category(category, model_name):
     
     # EfficientAd-specific setup
     if model_name == "efficientad":
-        patch_efficientad()  # Redirect pretrained weights to efficientad_resources/
-        # Override imagenet_dir to use absolute path inside efficientad_resources/
+        patch_efficientad()
         model_params["imagenet_dir"] = str(EFFICIENTAD_RESOURCES_DIR / "imagenette")
-        print(f"   ℹ️  EfficientAd: ImageNet data will be saved to {EFFICIENTAD_RESOURCES_DIR / 'imagenette'}")
-        print("   ℹ️  EfficientAd: Disabling image visualization (saving storage)")
+        print(f"   [INFO] EfficientAd: ImageNet directory: {EFFICIENTAD_RESOURCES_DIR / 'imagenette'}")
+        print("   [INFO] EfficientAd: Image visualization disabled")
         model_params["visualizer"] = False
         
     model = model_class(**model_params)
@@ -222,15 +207,11 @@ def train_category(category, model_name):
     # Output and save
     print_category_metrics(category_metrics)
     save_metrics(category_metrics, category, model_name)
-    print(f"\n✓ Completed: {category}\n")
+    print(f"\nCompleted: {category}\n")
     
     return category_metrics
 
 
-# =============================================================================
-# MAIN
-# =============================================================================
-
 def parse_args():
     """Parse command line arguments."""
     available_models = get_available_models()
@@ -262,37 +243,32 @@ Examples:
 def main():
     args = parse_args()
     
-    # Determine categories
     if args.category == "all":
         categories = MVTEC_CATEGORIES
-        print(f"🔄 Training on ALL {len(categories)} categories")
+        print(f"Training on ALL {len(categories)} categories")
     else:
         categories = [args.category]
-        print(f"🎯 Training on: {args.category}")
+        print(f"Training on: {args.category}")
     
-    # Determine models
     if args.model == "all":
         models = get_available_models()
-        print(f"📦 Models: ALL ({', '.join(models)})")
+        print(f"Models: ALL ({', '.join(models)})")
     else:
         models = [args.model]
-        print(f"📦 Model: {args.model}")
+        print(f"Model: {args.model}")
 
-    # Create results directory
     DIR_RESULTS.mkdir(parents=True, exist_ok=True)
     
-    # Training loop
     all_metrics = []
     for model_name in models:
         if len(models) > 1:
-            print(f"\n{'#'*60}")
-            print(f"# MODEL: {model_name.upper()}")
-            print(f"{'#'*60}")
+            print(f"\n{'='*60}")
+            print(f"MODEL: {model_name.upper()}")
+            print(f"{'='*60}")
         
         model_metrics = [train_category(cat, model_name) for cat in categories]
         all_metrics.extend(model_metrics)
         
-        # Final Report per model
         print_final_report(model_metrics, model_name)