diff --git a/README.md b/README.md index 8e93371..c1ac7f0 100644 --- a/README.md +++ b/README.md @@ -23,10 +23,11 @@ AnomaVision is a computer vision project for finding **defects and unusual patterns** in images. -It supports two anomaly detection methods: +It supports three anomaly detection methods: - **PaDiM** — a simple and fast baseline. - **PatchCore** — a lightweight memory-based method. +- **EfficientAD** — a student/teacher model with a compact reconstruction branch for fast anomaly detection. You only need **normal (`good`) images** to train the anomaly detector. @@ -43,6 +44,7 @@ You only need **normal (`good`) images** to train the anomaly detector. - Create anomaly heatmaps showing where the problem is. - Export models to **ONNX, OpenVINO, and TensorRT**. - Export and compile **PaDiM and PatchCore to XModel for the AMD/Xilinx Kria KV260**. +- Switch between PaDiM, PatchCore, and EfficientAD without changing the CLI workflow. ## Quick start @@ -53,31 +55,19 @@ You only need **normal (`good`) images** to train the anomaly detector. ```bash git clone https://github.com/DeepKnowledge1/AnomaVision.git cd AnomaVision - -# Create and activate a virtual environment uv venv --python 3.11 .venv -source .venv/bin/activate # Windows: .venv\Scripts\Activate.ps1 - -# Install with your hardware extra +source .venv/bin/activate # Windows: .venv\\Scripts\\Activate.ps1 uv sync --extra cpu # CPU uv sync --extra cu121 # CUDA 12.1 ``` ---- - #### Option B — From PyPI (production / quick start) ```bash -# CPU · Mac, CI runners, edge devices uv pip install "anomavision[cpu]" - -# NVIDIA GPU · pick your CUDA version -uv pip install "anomavision[cu118]" # CUDA 11.8 -uv pip install "anomavision[cu121]" # CUDA 12.1 -uv pip install "anomavision[cu124]" # CUDA 12.4 +uv pip install "anomavision[cu121]" ``` - For other environments, see [Installation](docs/installation.md). ### 2. Prepare your images @@ -88,9 +78,6 @@ Use a simple MVTec-style folder structure: dataset/ └── bottle/ ├── ground_truth/ - │ ├── broken_large/ - │ ├── broken_small/ - │ └── contamination/ ├── test/ │ ├── broken_large/ │ ├── broken_small/ @@ -102,33 +89,56 @@ dataset/ Training uses the **good** images. Test images can contain defects. -### 3. Train +### 3. Choose an algorithm + +The existing configuration format works unchanged: + +```yaml +algorithm: padim +``` + +Switch to EfficientAD by changing one value: + +```yaml +algorithm: efficientad +``` + +You can also use the native model selector: + +```yaml +model: + name: efficientad +``` -Create or edit `config.yml` and point `dataset_path` to your dataset. +The CLI commands remain the same. -Then run: +### 4. Train ```bash anomavision train --config config.yml ``` -PaDiM is the default model. For PatchCore, set `algorithm: patchcore` in the configuration. +For a quick EfficientAD configuration, see [`examples/efficientad_cpu.yml`](examples/efficientad_cpu.yml). -### 4. Detect +### 5. Detect ```bash -anomavision detect --config config.yml --img_path ./dataset/bottle/test +anomavision detect --config config.yml --model model.pt --img_path ./test_images ``` -### 5. Export +### 6. Export -For a portable model, ONNX is a good place to start: +```bash +anomavision export --config config.yml --model model.pt --format onnx +``` + +### 7. Evaluate ```bash -anomavision export --config config.yml --format onnx +anomavision eval --config config.yml --model model.pt --class_name bottle ``` -For more export options, see [Export and deployment](docs/production_deployment.md). +For EfficientAD-specific options and limitations, see [EfficientAD](docs/efficientad.md). ## KV260 support @@ -142,13 +152,10 @@ PyTorch → INT8 quantization → XModel → KV260 DPU compilation Both PaDiM and PatchCore currently compile with **1 DPU subgraph** in the KV260 compiler. -The complete setup and commands are in: - -**[KV260 XModel Guide](docs/kv260_xmodel.md)** +The complete setup and commands are in the [KV260 XModel Guide](docs/kv260_xmodel.md). > XModel compilation has been validated in the Vitis AI environment. Final on-device KV260 validation requires the physical hardware. - ## Production Autopilot **Production Autopilot is the easiest way to move from two trained models to one deployable choice.** It compares PaDiM and ultra-light PatchCore on the same labeled test split, calibrates a separate threshold for each, profiles median and P95 latency on your hardware, checks localization health, and packages the selected artifact with a self-contained HTML dashboard. @@ -166,8 +173,7 @@ anomavision autopilot \ --output_dir ./production_package ``` -Open `production_package/production_autopilot_report.html` to see the selected model, AUROC, calibrated threshold, localization diagnostics, memory, median latency, P95 latency, and deployment recommendation. The package also contains `deployment_manifest.json`, `localization_report.md`, and the selected model artifact. See [`docs/production_deployment.md`](docs/production_deployment.md) for GPU, TensorRT, INT8, and packaging details. - +Open `production_package/production_autopilot_report.html` to see the selected model, AUROC, calibrated threshold, localization diagnostics, memory, median latency, P95 latency, and deployment recommendation. See [`docs/production_deployment.md`](docs/production_deployment.md) for details. ## Documentation @@ -176,6 +182,7 @@ Open `production_package/production_autopilot_report.html` to see the selected m | Quick start | [`docs/quickstart.md`](docs/quickstart.md) | | Installation | [`docs/installation.md`](docs/installation.md) | | CLI and configuration | [`docs/cli.md`](docs/cli.md), [`docs/config.md`](docs/config.md) | +| EfficientAD | [`docs/efficientad.md`](docs/efficientad.md) | | Python API | [`docs/api.md`](docs/api.md) | | KV260 / XModel | [`docs/kv260_xmodel.md`](docs/kv260_xmodel.md) | | Production deployment | [`docs/production_deployment.md`](docs/production_deployment.md) | @@ -186,23 +193,18 @@ Open `production_package/production_autopilot_report.html` to see the selected m ## Python example -You can also use AnomaVision directly from Python: - ```python import torch -from torch.utils.data import DataLoader import anomavision +from torch.utils.data import DataLoader train_set = anomavision.AnodetDataset("./dataset/bottle/train/good") -train_loader = DataLoader(train_set, batch_size=16, shuffle=False) - -model = anomavision.Padim(backbone="resnet18", device=torch.device("cpu")) -model.fit(train_loader) +train_loader = DataLoader(train_set, batch_size=1, shuffle=False) -batch = next(iter(train_loader)) -if isinstance(batch, (tuple, list)): - batch = batch[0] +model = anomavision.EfficientAD(device=torch.device("cpu")) +model.fit(train_loader, epochs=1) +batch = next(iter(train_loader))[0] scores, maps = model.predict(batch) ``` diff --git a/anomavision/__init__.py b/anomavision/__init__.py index 6eff355..117757e 100644 --- a/anomavision/__init__.py +++ b/anomavision/__init__.py @@ -9,6 +9,7 @@ """ from .algorithm.common.feature_extraction import ResnetEmbeddingsExtractor +from .algorithm.efficientad import EfficientAD from .algorithm.padim import Padim from .algorithm.patchcore import PatchCore from .datasets.dataset import AnodetDataset diff --git a/anomavision/algorithm/efficientad/__init__.py b/anomavision/algorithm/efficientad/__init__.py new file mode 100644 index 0000000..eb7f485 --- /dev/null +++ b/anomavision/algorithm/efficientad/__init__.py @@ -0,0 +1,5 @@ +"""EfficientAD anomaly detection algorithm.""" + +from .efficientad import EfficientAD + +__all__ = ["EfficientAD"] diff --git a/anomavision/algorithm/efficientad/efficientad.py b/anomavision/algorithm/efficientad/efficientad.py new file mode 100644 index 0000000..db9aa58 --- /dev/null +++ b/anomavision/algorithm/efficientad/efficientad.py @@ -0,0 +1,195 @@ +"""Native AnomaVision implementation of EfficientAD.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from torchvision.models import EfficientNet_B0_Weights, efficientnet_b0 + + +class _FeatureTeacher(nn.Module): + def __init__(self, pretrained: bool = True) -> None: + super().__init__() + weights = EfficientNet_B0_Weights.DEFAULT if pretrained else None + net = efficientnet_b0(weights=weights) + self.features = nn.Sequential(*list(net.features[:6])) + self.out_channels = 112 + for p in self.parameters(): + p.requires_grad_(False) + self.eval() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + with torch.no_grad(): + return self.features(x) + + +class _Student(nn.Module): + def __init__(self, out_channels: int = 112) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Conv2d(3, 64, 3, stride=2, padding=1), + nn.BatchNorm2d(64), nn.ReLU(inplace=True), + nn.Conv2d(64, 96, 3, stride=2, padding=1), + nn.BatchNorm2d(96), nn.ReLU(inplace=True), + nn.Conv2d(96, 112, 3, stride=2, padding=1), + nn.BatchNorm2d(112), nn.ReLU(inplace=True), + nn.Conv2d(112, out_channels, 3, stride=2, padding=1), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class _AutoEncoder(nn.Module): + def __init__(self) -> None: + super().__init__() + self.encoder = nn.Sequential( + nn.Conv2d(3, 32, 4, 2, 1), nn.ReLU(inplace=True), + nn.Conv2d(32, 64, 4, 2, 1), nn.ReLU(inplace=True), + nn.Conv2d(64, 96, 4, 2, 1), nn.ReLU(inplace=True), + ) + self.decoder = nn.Sequential( + nn.ConvTranspose2d(96, 64, 4, 2, 1), nn.ReLU(inplace=True), + nn.ConvTranspose2d(64, 32, 4, 2, 1), nn.ReLU(inplace=True), + nn.ConvTranspose2d(32, 3, 4, 2, 1), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.decoder(self.encoder(x)) + + +class EfficientAD(nn.Module): + """EfficientAD-compatible anomaly detector for the AnomaVision pipeline.""" + + def __init__( + self, + device: torch.device = torch.device("cpu"), + model_size: str = "s", + lr: float = 1e-4, + weight_decay: float = 1e-5, + pretrained_teacher: bool = True, + teacher_weights: Optional[str] = None, + feature_weight: float = 1.0, + reconstruction_weight: float = 0.1, + ) -> None: + super().__init__() + model_size = str(model_size).lower() + if model_size not in {"s", "m", "small", "medium"}: + raise ValueError("EfficientAD model_size must be one of: s, m") + if lr <= 0 or weight_decay < 0: + raise ValueError("lr must be > 0 and weight_decay must be >= 0") + + self.device = torch.device(device) + self.model_size = "m" if model_size in {"m", "medium"} else "s" + self.lr = float(lr) + self.weight_decay = float(weight_decay) + self.feature_weight = float(feature_weight) + self.reconstruction_weight = float(reconstruction_weight) + + self.teacher = _FeatureTeacher(pretrained=pretrained_teacher) + if teacher_weights: + state = torch.load(teacher_weights, map_location="cpu", weights_only=False) + self.teacher.load_state_dict(state, strict=False) + self.student = _Student(self.teacher.out_channels) + self.autoencoder = _AutoEncoder() + self.register_buffer("score_mean", torch.tensor(0.0)) + self.register_buffer("score_std", torch.tensor(1.0)) + self.register_buffer("trained", torch.tensor(False, dtype=torch.bool)) + self.to(self.device) + + def _normalise(self, x: torch.Tensor) -> torch.Tensor: + return x + + def _signals(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + x_norm = self._normalise(x) + teacher = self.teacher(x_norm) + student = self.student(x_norm) + feature_map = (student - teacher).pow(2).mean(dim=1) + reconstruction = (self.autoencoder(x) - x).abs().mean(dim=1) + feature_map = F.interpolate( + feature_map.unsqueeze(1), size=x.shape[-2:], mode="bilinear", align_corners=False + ).squeeze(1) + return feature_map, reconstruction + + def forward(self, x: torch.Tensor, return_map: bool = True, export: bool = False): + feature_map, reconstruction = self._signals(x) + score_map = feature_map + self.reconstruction_weight * reconstruction + scores = score_map.flatten(1).amax(1) + scores = (scores - self.score_mean) / self.score_std.clamp_min(1e-6) + return scores, score_map if return_map else None + + def fit(self, dataloader: torch.utils.data.DataLoader, epochs: int = 1) -> None: + self.train() + self.teacher.eval() + optimizer = torch.optim.Adam( + list(self.student.parameters()) + list(self.autoencoder.parameters()), + lr=self.lr, weight_decay=self.weight_decay, + ) + for _ in range(int(epochs)): + for batch in dataloader: + if isinstance(batch, (tuple, list)): + batch = batch[0] + batch = batch.to(self.device, non_blocking=True).float() + with torch.no_grad(): + teacher = self.teacher(self._normalise(batch)) + student = self.student(self._normalise(batch)) + reconstructed = self.autoencoder(batch) + feature_loss = F.mse_loss(student, teacher) + reconstruction_loss = F.l1_loss(reconstructed, batch) + loss = self.feature_weight * feature_loss + self.reconstruction_weight * reconstruction_loss + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + + self.eval() + values = [] + with torch.no_grad(): + for batch in dataloader: + if isinstance(batch, (tuple, list)): + batch = batch[0] + batch = batch.to(self.device).float() + fmap, recon = self._signals(batch) + values.append((fmap + self.reconstruction_weight * recon).flatten(1).amax(1)) + if values: + scores = torch.cat(values) + self.score_mean.copy_(scores.mean()) + self.score_std.copy_(scores.std(unbiased=False).clamp_min(1e-6)) + self.trained.fill_(True) + + def predict(self, batch: torch.Tensor, export: bool = False): + # Training-state validation is intentionally skipped during export. + # ``trained.item()`` creates data-dependent control flow that torch.export + # cannot specialize. The exported graph must contain tensor computation only. + if not export and not bool(self.trained.item()): + raise RuntimeError("EfficientAD model is not trained. Call fit() first.") + self.eval() + with torch.no_grad(): + return self.forward(batch.to(self.device).float(), export=export) + + def to_device(self, device: torch.device) -> None: + self.device = torch.device(device) + self.to(self.device) + + def save_statistics(self, path: str, half: Optional[bool] = None) -> None: + if not bool(self.trained.item()): + raise RuntimeError("Model is not trained. Call fit() first.") + torch.save({ + "algorithm": "efficientad", "model_state": self.state_dict(), + "model_size": self.model_size, "lr": self.lr, + "weight_decay": self.weight_decay, + }, path) + + @staticmethod + def load_statistics(path: str, device: str = "cpu") -> "EfficientAD": + data = torch.load(path, map_location="cpu", weights_only=False) + if data.get("algorithm") != "efficientad": + raise ValueError("Not an EfficientAD statistics artifact") + model = EfficientAD( + device=torch.device(device), model_size=data.get("model_size", "s"), + pretrained_teacher=False, + ) + model.load_state_dict(data["model_state"]) + return model diff --git a/anomavision/config.py b/anomavision/config.py index 5965679..88f7b71 100644 --- a/anomavision/config.py +++ b/anomavision/config.py @@ -13,19 +13,30 @@ def load_config(path: str = None): if not p.exists(): raise FileNotFoundError(f"Config not found: {path}") if p.suffix.lower() in {".yml", ".yaml"}: - return yaml.safe_load(p.read_text()) - if p.suffix.lower() == ".json": - return json.loads(p.read_text()) - raise ValueError("Config must be .yml/.yaml or .json") + config = yaml.safe_load(p.read_text()) or {} + elif p.suffix.lower() == ".json": + config = json.loads(p.read_text()) or {} + else: + raise ValueError("Config must be .yml/.yaml or .json") + + # Native model selector: allow either historical ``algorithm: padim`` + # or ``model: {name: padim}`` configuration. + model_section = config.get("model") + if isinstance(model_section, dict): + if not config.get("algorithm") and model_section.get("name"): + config["algorithm"] = model_section["name"] + if model_section.get("file") and not config.get("model_path"): + config["model_path"] = model_section["file"] + # Existing detect/export code expects ``model`` to be the artifact name. + config["model"] = config.get("model_path") + return config def to_dict(ns: Namespace) -> dict: - # turn argparse Namespace into a dict (ignores None later) return {k: v for k, v in vars(ns).items()} def pick(*vals): - # first non-None for v in vals: if v is not None: return v diff --git a/anomavision/train.py b/anomavision/train.py index 7761310..90594c3 100644 --- a/anomavision/train.py +++ b/anomavision/train.py @@ -141,6 +141,37 @@ def create_parser(add_help: bool = True) -> argparse.ArgumentParser: default=None, help="Seed used for deterministic PatchCore coreset selection.", ) + parser.add_argument( + "--efficientad_model_size", + type=str, + choices=["s", "m"], + default=None, + help="EfficientAD model size.", + ) + parser.add_argument( + "--efficientad_lr", + type=float, + default=None, + help="EfficientAD learning rate.", + ) + parser.add_argument( + "--efficientad_weight_decay", + type=float, + default=None, + help="EfficientAD weight decay.", + ) + parser.add_argument( + "--efficientad_epochs", + type=int, + default=None, + help="EfficientAD training epochs.", + ) + parser.add_argument( + "--efficientad_pretrained_teacher", + action="store_true", + default=None, + help="Use ImageNet-pretrained EfficientNet teacher.", + ) parser.add_argument( "--output_model", type=str, @@ -163,7 +194,7 @@ def create_parser(add_help: bool = True) -> argparse.ArgumentParser: "--algorithm", type=str, default=None, - help="Algorithm name (e.g., padim, patchcore).", + help="Algorithm name (e.g., padim, patchcore, efficientad).", ) parser.add_argument( "--log_level", @@ -181,10 +212,10 @@ def run_training(args): Executes the training pipeline. Args: - args: Namespace object containing command line arguments. + args: Namespace object containing configuration. Returns: - padim (AnomaVision.Padim): The trained model object. + model (AnomaVision anomaly model): The trained model object. config (edict): The final merged configuration. run_dir (Path): The directory where artifacts were saved. dataloaders (dict): Dictionary containing the 'train' DataLoader. @@ -195,13 +226,23 @@ def run_training(args): config = edict(merge_config(args, cfg)) setup_logging(enabled=True, log_level=config.log_level, log_to_file=True) - logger = get_logger("anomavision.train") # Force it into anomavision hierarchy + logger = get_logger("anomavision.train") if not config.dataset_path: error_msg = "dataset.path is required (via --dataset_path or config.common.dataset_path)" logger.error(error_msg) raise ValueError(error_msg) + algorithm = str(config.algorithm).lower() + if algorithm not in {"padim", "patchcore", "efficientad"}: + raise ValueError( + f"Unsupported algorithm: {algorithm}. Available: padim, patchcore, efficientad" + ) + if algorithm == "efficientad" and not bool(config.get("normalize", True)): + raise ValueError( + "EfficientAD requires normalize: true because its teacher uses ImageNet preprocessing" + ) + t0 = time.perf_counter() logger.info( @@ -224,19 +265,11 @@ def run_training(args): ) # === Dataset === - # Handle the 'class_name' logic safely. - # If dataset_path ends with the class name, use parent? - # Original logic assumes dataset_path is the container of class folders OR the class folder itself? - # Original code: os.path.join(realpath(dataset_path), config.class_name, "train", "good") - # This implies dataset_path is the root (e.g. MVTec root) and config.class_name is "bottle" - root = os.path.join( os.path.realpath(config.dataset_path), config.class_name, "train", "good" ) if not os.path.isdir(root): - # Fallback check: maybe dataset_path ALREADY points to the class folder? - # This makes it more robust for different input styles potential_root = os.path.join( os.path.realpath(config.dataset_path), "train", "good" ) @@ -277,7 +310,7 @@ def run_training(args): config.layer_indices, ) - if str(config.algorithm).lower() == "patchcore": + if algorithm == "patchcore": model = anomavision.PatchCore( backbone=config.backbone, device=device, @@ -289,6 +322,14 @@ def run_training(args): coreset_method=config.get("coreset_method", "kcenter"), coreset_seed=int(config.get("coreset_seed", 42)), ) + elif algorithm == "efficientad": + model = anomavision.EfficientAD( + device=device, + model_size=config.get("efficientad_model_size", "s"), + lr=float(config.get("efficientad_lr", 1e-4)), + weight_decay=float(config.get("efficientad_weight_decay", 1e-5)), + pretrained_teacher=bool(config.get("efficientad_pretrained_teacher", True)), + ) else: model = anomavision.Padim( backbone=config.backbone, @@ -298,7 +339,10 @@ def run_training(args): ) t_fit = time.perf_counter() - model.fit(dl) + if algorithm == "efficientad": + model.fit(dl, epochs=int(config.get("efficientad_epochs", 1))) + else: + model.fit(dl) logger.info("fit: completed in %.2fs", time.perf_counter() - t_fit) # === Save === @@ -319,7 +363,6 @@ def run_training(args): logger.info("saved: model=%s, config=%s", model_path, Path(run_dir) / "config.yml") logger.info("=== Training done in %.2fs ===", time.perf_counter() - t0) - # Return objects for external usage (e.g. MLOps pipeline) return model, config, run_dir, {"train": dl} @@ -334,7 +377,7 @@ def main(args=None): if checker.is_repo(): checker.check_status() except Exception: - pass # Never block training over a git check + pass run_training(args) diff --git a/config.yml b/config.yml index 0af6c53..ca012b3 100644 --- a/config.yml +++ b/config.yml @@ -1,124 +1,129 @@ # ========================= # Dataset / preprocessing (shared by train, detect, eval, stream) # ========================= -dataset_path: "D:/01-DATA" # Root dataset folder (MVTec-style: contains train/test subfolders) -class_name: "bottle" # Class name for MVTec dataset -resize: [224, 224] # Resize dimensions before processing [width, height] -crop_size: # Final crop size [width, height] -normalize: true # Whether to normalize images -no_normalize: false # Inverse flag (used by CLI) – keep in sync with "normalize" -norm_mean: [0.485, 0.456, 0.406] # Mean values for normalization (ImageNet default) -norm_std: [0.229, 0.224, 0.225] # Standard deviation for normalization (ImageNet default) +dataset_path: "D:/01-DATA" +class_name: "bottle" +resize: [224, 224] +crop_size: +normalize: true +no_normalize: false +norm_mean: [0.485, 0.456, 0.406] +norm_std: [0.229, 0.224, 0.225] # ========================= # Model / training # ========================= -backbone: "resnet18" # Backbone CNN architecture (resnet18 | wide_resnet50) -algorithm: "padim" # Algorithm to use: padim or patchcore -coreset_ratio: 0.02 # Ultra-light PatchCore memory fraction -max_memory_patches: 2048 # Hard memory-bank cap for low latency -patch_grid: 14 # Pool feature maps to at most 14x14 patches -search_chunk_size: 1024 # Bound nearest-neighbor working memory -coreset_method: "kcenter" # PatchCore selection: kcenter or random -coreset_seed: 42 # Reproducible k-center initialization -feat_dim: 50 # Feature dimension size for embedding -layer_indices: [0] # Which backbone layers to extract features from (0,1,2,3) -model_data_path: "./distributions" # Path to store/load model-related data -model: "model.pt" # File name for saved model (used by detect/eval/export) -output_model: "model.pt" # File name for saving trained model (train.py expects this) -batch_size: 2 # Training/evaluation/inference batch size -device: "auto" # Device to run on: "cpu", "cuda", or "auto" +backbone: "resnet18" +algorithm: "efficientad" # padim | patchcore | efficientad +# Equivalent native selector: +# model: +# name: padim + +# PaDiM +feat_dim: 50 +layer_indices: [0] + +# PatchCore +coreset_ratio: 0.02 +max_memory_patches: 2048 +patch_grid: 14 +search_chunk_size: 1024 +coreset_method: "kcenter" +coreset_seed: 42 + +# EfficientAD +efficientad_model_size: "s" # s | m +efficientad_lr: 0.0001 +efficientad_weight_decay: 0.00001 +efficientad_epochs: 1 +efficientad_pretrained_teacher: true + +model_data_path: "./distributions" +model: "model.pt" +output_model: "model.pt" +batch_size: 2 +device: "auto" # ========================= # Logging / run metadata # ========================= -log_level: "INFO" # Logging level: DEBUG, INFO, WARNING, ERROR -run_name: "anomav_exp" # Name of experiment run (used for organizing results) -detailed_timing: false # Enable detailed timing measurements +log_level: "INFO" +run_name: "anomav_exp" +detailed_timing: false # ========================= -# Visualization (shared by detect & eval) +# Visualization # ========================= -enable_visualization: true # Enable visualization during inference/evaluation -save_visualizations: true # Save visualization results to disk -viz_output_dir: "./visualizations/" # Directory to save visualization results -viz_alpha: 0.5 # Transparency factor for overlay heatmaps -viz_padding: 40 # Padding added around visualization -viz_color: "128,0,128" # RGB color for visualization overlays +enable_visualization: true +save_visualizations: true +viz_output_dir: "./visualizations/" +viz_alpha: 0.5 +viz_padding: 40 +viz_color: "128,0,128" # ========================= -# Inference (detect.py) -# ========================= -img_path: "D:/01-DATA/test" # Path to test images for inference -thresh: null # Legacy fallback; prefer algorithm-specific thresholds -thresh_padim: 13.0 # PaDiM score threshold; null lets eval auto-select -thresh_patchcore: 0.25 # PatchCore score threshold; null lets eval auto-select -num_workers: 1 # Number of workers for dataloader -pin_memory: false # Use pinned memory for faster GPU transfers -overwrite: false # Overwrite existing run directory without auto-incrementing +# Inference +# ========================= +img_path: "D:/01-DATA/test" +thresh: null +thresh_padim: 13.0 +thresh_patchcore: 0.25 +# EfficientAD uses its own score scale. Keep this independent from PaDiM. +# The score is normalized against the normal-training score distribution. +thresh_efficientad: 1.0 +num_workers: 1 +pin_memory: false +overwrite: false # ========================= -# Evaluation (eval.py) +# Evaluation # ========================= -memory_efficient: true # Use memory efficient evaluation mode +memory_efficient: true # ========================= -# Export (export.py) -# ========================= -format: "all" # onnx, tensorrt, torchscript, openvino, all -opset: 18 # ONNX opset version -precision: "auto" # ONNX/TorchScript precision: auto, fp16, fp32 -tensorrt_precision: "fp16" # TensorRT precision: fp32, fp16, int8 -dynamic_batch: true # Allow dynamic batch size in exported model -static_batch: false # Disable dynamic batch size (if true) -min_batch: 1 # TensorRT dynamic profile minimum batch -opt_batch: 1 # TensorRT dynamic profile optimal batch -max_batch: 4 # TensorRT dynamic profile maximum batch -workspace_gb: 2.0 # TensorRT workspace limit in GiB -calib_dir: null # INT8 calibration image directory; auto-derived when null -calib_samples: 100 # Maximum real images used for INT8 calibration -quantize_dynamic: false # Also write dynamically quantized INT8 ONNX -quantize_static: false # Also write statically quantized INT8 ONNX -optimize: false # Enable mobile optimization for TorchScript -output_path: null # Optional explicit output filename -half: false # Legacy compatibility field -int8: false # Legacy compatibility field; use tensorrt_precision: int8 +# Export +# ========================= +format: "all" +opset: 18 +precision: "auto" +tensorrt_precision: "fp16" +dynamic_batch: true +static_batch: false +min_batch: 1 +opt_batch: 1 +max_batch: 4 +workspace_gb: 2.0 +calib_dir: null +calib_samples: 100 +quantize_dynamic: false +quantize_static: false +optimize: false +output_path: null +half: false +int8: false # ========================= -# Streaming Configuration +# Streaming # ========================= -stream_mode: false # true = real-time streaming, false = static dataset - +stream_mode: false stream_source: - type: "webcam" # webcam | video | mqtt | tcp - - # Webcam settings (type: webcam) - camera_id: 0 # Camera device index - - # Video file settings (type: video) - video_path: "path/to/video.mp4" # Path to video file - loop: false # Loop video when it ends - - # MQTT settings (type: mqtt) - broker: "localhost" # MQTT broker hostname/IP - port: 1883 # MQTT broker port - topic: "camera/frames" # Topic to subscribe to - client_id: null # Optional client ID - keepalive: 60 # Keepalive interval (seconds) - qos: 0 # QoS level (0, 1, or 2) - max_queue_size: 10 # Max buffered frames - read_timeout: 1.0 # Timeout for reading frames (seconds) - - # TCP settings (type: tcp) - host: "192.168.1.100" # TCP server hostname/IP - port: 8080 # TCP server port (matches code) - recv_timeout: 1.0 # Socket timeout for recv (seconds) - header_size: 4 # Length header size in bytes - max_message_size: 10485760 # Max payload size (10MB) - - -# Streaming processing settings -stream_max_frames: null # Max frames to process (null = infinite) -stream_display_fps: true # Show FPS during streaming -stream_save_detections: true # Save detected anomalies to disk -stream_detection_dir: "./stream_detections/" # Directory for saved detections + type: "webcam" + camera_id: 0 + video_path: "path/to/video.mp4" + loop: false + broker: "localhost" + port: 1883 + topic: "camera/frames" + client_id: null + keepalive: 60 + qos: 0 + max_queue_size: 10 + read_timeout: 1.0 + host: "192.168.1.100" + recv_timeout: 1.0 + header_size: 4 + max_message_size: 10485760 +stream_max_frames: null +stream_display_fps: true +stream_save_detections: true +stream_detection_dir: "./stream_detections/" diff --git a/docs/config.md b/docs/config.md index 0b5c9cd..929b679 100644 --- a/docs/config.md +++ b/docs/config.md @@ -1,167 +1,130 @@ +# ⚙️ Configuration Guide +AnomaVision scripts (`train.py`, `detect.py`, `eval.py`, `export.py`) all accept a **YAML/JSON config file**. You can override any field via CLI arguments. -# ⚙️ Configuration Guide +## 1. Dataset & Preprocessing + +| Key | Type | Default | Description | +|---|---|---|---| +| `dataset_path` | str | None | Root dataset folder containing MVTec-style structure. | +| `class_name` | str | None | Target class name. | +| `resize` | [int,int] | None | Resize before processing. | +| `crop_size` | [int,int] | None | Center crop size. | +| `normalize` | bool | True | Apply input normalization. | +| `norm_mean` | [float] | [0.485,0.456,0.406] | RGB mean. | +| `norm_std` | [float] | [0.229,0.224,0.225] | RGB standard deviation. | -AnomaVision scripts (`train.py`, `detect.py`, `eval.py`, `export.py`) all accept a **YAML/JSON config file**. -You can override any field via CLI arguments. +## 2. Algorithm selection -Example: +The historical selector remains supported: -```bash -python train.py --config config.yml +```yaml +algorithm: padim ``` ---- +Available values are `padim`, `patchcore`, and `efficientad`. -## 1. Dataset & Preprocessing +A native model selector is also supported: -| Key | Type | Default | Description | -| -------------- | ---------- | ---------------------- | -------------------------------------------------------------------------------------------- | -| `dataset_path` | str | None | Root dataset folder containing MVTec-style structure (`class/train/good`, `class/test/...`). | -| `class_name` | str | None | Target class name (e.g. `bottle`, `cable`). | -| `resize` | \[int,int] | None | Resize before processing (e.g. `[256,192]`). One value applies square resize. | -| `crop_size` | \[int,int] | None | Center crop size (e.g. `[224,224]`). One value = square crop. | -| `normalize` | bool | True | Apply input normalization. | -| `norm_mean` | \[float] | \[0.485, 0.456, 0.406] | Mean values (RGB) if normalize is enabled. | -| `norm_std` | \[float] | \[0.229, 0.224, 0.225] | Std values (RGB) if normalize is enabled. | - ---- - -## 2. Training - -| Key | Type | Default | Description | -| ----------------- | ---- | --------------- | ------------------------------------------------ | -| `backbone` | str | resnet18 | Feature extractor (`resnet18`, `wide_resnet50`). | -| `batch_size` | int | 16 | Training batch size. | -| `feat_dim` | int | 100 | Number of random feature dimensions kept. | -| `layer_indices` | list | \[0,1,2] | Backbone layers used for features. | -| `run_name` | str | exp | Name of training run. | -| `model_data_path` | str | ./distributions | Where trained models/configs are stored. | -| `output_model` | str | padim\_model.pt | Name of saved model. | - ---- - -## 3. Detection - -| Key | Type | Default | Description | -| ---------------------- | ----- | ----------- | ------------------------------------ | -| `img_path` | str | None | Path to test images or folder. | -| `model` | str | None | Model file (`.pt`, `.pth`, `.onnx`). | -| `device` | str | auto | Device (`cpu`, `cuda`, or `auto`). | -| `batch_size` | int | 1 | Batch size for inference. | -| `thresh` | float | None | Legacy global anomaly threshold fallback. | -| `thresh_padim` | float | None | PaDiM-specific threshold; takes precedence over `thresh`. | -| `thresh_patchcore` | float | None | PatchCore-specific threshold; takes precedence over `thresh`. | -| `enable_visualization` | bool | False | Enable heatmap overlays. | -| `save_visualizations` | bool | False | Save visualization images. | -| `viz_output_dir` | str | ./results/ | Directory to save images. | -| `viz_alpha` | float | 0.5 | Heatmap transparency. | -| `viz_padding` | int | 40 | Padding around bounding boxes. | -| `viz_color` | str | "128,0,128" | RGB highlight color. | - ---- - -## 4. PatchCore +```yaml +model: + name: efficientad +``` + +When `model.name` is used, `load_config()` maps it to the same internal `algorithm` value, so the existing CLI workflow does not change. + +## 3. Training | Key | Type | Default | Description | |---|---|---:|---| -| `coreset_ratio` | float | 0.02 | Fraction of normal patches retained. | -| `max_memory_patches` | int | 2048 | Hard memory-bank cap. | -| `patch_grid` | int | 14 | Spatial pooling grid for lightweight localization. | -| `search_chunk_size` | int | 1024 | Chunk size for bounded nearest-neighbor search. | -| `coreset_method` | str | kcenter | `kcenter` for deterministic diverse selection or `random`. | -| `coreset_seed` | int | 42 | Reproducibility seed for coreset selection. | - ---- +| `backbone` | str | resnet18 | PaDiM/PatchCore feature extractor. | +| `batch_size` | int | 16 | Training batch size. | +| `feat_dim` | int | 100 | PaDiM feature dimensions. | +| `layer_indices` | list | [0,1,2] | PaDiM/PatchCore feature layers. | +| `run_name` | str | exp | Training run name. | +| `model_data_path` | str | ./distributions | Model artifact root. | +| `output_model` | str | padim_model.pt | Saved model filename. | -## 5. Evaluation +### EfficientAD -| Key | Type | Default | Description | -| ------------------ | ---- | ------- | -------------------------------- | -| `memory_efficient` | bool | True | Use memory-efficient evaluation. | -| `detailed_timing` | bool | False | Log per-image timings. | +| Key | Type | Default | Description | +|---|---|---:|---| +| `efficientad_model_size` | str | s | EfficientAD size: `s` or `m`. | +| `efficientad_lr` | float | 0.0001 | Adam learning rate. | +| `efficientad_weight_decay` | float | 0.00001 | Adam weight decay. | +| `efficientad_epochs` | int | 1 | Number of normal-data training epochs. | +| `efficientad_pretrained_teacher` | bool | true | Use ImageNet-pretrained EfficientNet teacher. | -(Other keys mirror **Detection** and **Training**.) +EfficientAD requires `normalize: true` in the current integration because the teacher uses ImageNet preprocessing. ---- +## 4. Detection -## 6. Export +| Key | Type | Default | Description | +|---|---|---|---| +| `img_path` | str | None | Path to test images or folder. | +| `model` | str | None | Model file (`.pt`, `.pth`, `.onnx`). | +| `device` | str | auto | Device (`cpu`, `cuda`, or `auto`). | +| `batch_size` | int | 1 | Inference batch size. | +| `thresh` | float | None | Legacy global anomaly threshold. | +| `thresh_padim` | float | None | PaDiM-specific threshold. | +| `thresh_patchcore` | float | None | PatchCore-specific threshold. | +| `thresh_efficientad` | float | None | EfficientAD-specific threshold. | +| `enable_visualization` | bool | False | Enable heatmap overlays. | +| `save_visualizations` | bool | False | Save visualization images. | +| `viz_output_dir` | str | ./results/ | Visualization directory. | + +## 5. PatchCore -| Key | Type | Default | Description | -| ------------------ | ---- | ------- | --------------------------------------------------------- | -| `format` | str | onnx | Export target (`onnx`, `torchscript`, `openvino`, `all`). | -| `precision` | str | auto | Precision (`fp32`, `fp16`, or `auto`). | -| `opset` | int | 17 | ONNX opset version. | -| `static_batch` | bool | False | Disable dynamic batch. | -| `optimize` | bool | False | TorchScript mobile optimization. | -| `quantize_dynamic` | bool | False | Export dynamic INT8 ONNX. | -| `quantize_static` | bool | False | Export static INT8 ONNX (requires calibration). | -| `calib_samples` | int | 100 | Calibration samples for static quantization. | -| `tensorrt_precision` | str | fp16 | TensorRT precision (`fp32`, `fp16`, or `int8`). | -| `workspace_gb` | float | 2.0 | TensorRT builder workspace limit in GB. | -| `min_batch` | int | 1 | Minimum TensorRT dynamic batch size. | -| `opt_batch` | int | 1 | Optimized TensorRT dynamic batch size. | -| `max_batch` | int | 4 | Maximum TensorRT dynamic batch size. | -| `calib_dir` | str/null | null | Real-image directory for TensorRT INT8 calibration. | +| Key | Type | Default | Description | +|---|---|---:|---| +| `coreset_ratio` | float | 0.02 | Fraction of normal patches retained. | +| `max_memory_patches` | int | 2048 | Hard memory-bank cap. | +| `patch_grid` | int | 14 | Spatial pooling grid. | +| `search_chunk_size` | int | 1024 | Nearest-neighbor chunk size. | +| `coreset_method` | str | kcenter | `kcenter` or `random`. | +| `coreset_seed` | int | 42 | Coreset reproducibility seed. | ---- +## 6. Evaluation -## 7. Production Autopilot +| Key | Type | Default | Description | +|---|---|---|---| +| `memory_efficient` | bool | True | Use memory-efficient evaluation. | +| `detailed_timing` | bool | False | Log detailed timings. | -Autopilot reads `dataset_path` and `class_name`, uses the complete labeled `test` split by default, and writes `deployment_manifest.json`, `production_autopilot_report.html`, and a Markdown fallback. Its HTML report distinguishes image AUROC, pixel AUROC, anomaly localization coverage, normal-image false-positive localization, and anomaly mean mask area. +Other keys mirror **Detection** and **Training**. ---- +## 7. Export -## 8. Logging +| Key | Type | Default | Description | +|---|---|---:|---| +| `format` | str | onnx | `onnx`, `torchscript`, `openvino`, `all`. | +| `precision` | str | auto | `fp32`, `fp16`, or `auto`. | +| `opset` | int | 17 | ONNX opset version. | +| `static_batch` | bool | False | Disable dynamic batch. | +| `quantize_dynamic` | bool | False | Export dynamic INT8 ONNX. | +| `quantize_static` | bool | False | Export static INT8 ONNX. | +| `calib_samples` | int | 100 | Static quantization samples. | -| Key | Type | Default | Description | -| ----------- | ---- | ------- | ---------------------------------------------------- | -| `log_level` | str | INFO | Logging level (`DEBUG`, `INFO`, `WARNING`, `ERROR`). | +## Example: switch PaDiM → EfficientAD ---- +```yaml +# Everything else in the existing config can remain unchanged. +algorithm: efficientad + +efficientad_model_size: s +efficientad_lr: 0.0001 +efficientad_weight_decay: 0.00001 +efficientad_epochs: 1 +efficientad_pretrained_teacher: true +``` -## Example Config +Then use the same commands: -```yaml -dataset_path: ./dataset -class_name: bottle -resize: [256, 192] -crop_size: [224, 224] -normalize: true -norm_mean: [0.485, 0.456, 0.406] -norm_std: [0.229, 0.224, 0.225] - -backbone: resnet18 -batch_size: 16 -feat_dim: 100 -layer_indices: [0, 1, 2] -output_model: model.pt -run_name: exp1 -model_data_path: ./distributions/padim/bottle/anomav_exp - -model: model.onnx -device: auto -enable_visualization: true -save_visualizations: true -viz_output_dir: ./results/ - -format: onnx -precision: fp16 -quantize_dynamic: true - -# Algorithm-specific thresholds -thresh: null -thresh_padim: null -thresh_patchcore: null - -# Ultra-light PatchCore -algorithm: patchcore -coreset_method: kcenter -coreset_seed: 42 -coreset_ratio: 0.02 -max_memory_patches: 2048 -patch_grid: 14 -search_chunk_size: 1024 +```bash +anomavision train --config config.yml +anomavision export --config config.yml --model model.pt --format onnx +anomavision detect --config config.yml --model model.onnx --img_path ./test_images +anomavision eval --config config.yml --model model.pt --class_name bottle ``` ---- +See [`docs/efficientad.md`](efficientad.md) for the complete EfficientAD guide. diff --git a/docs/efficientad.md b/docs/efficientad.md new file mode 100644 index 0000000..4e6cdc6 --- /dev/null +++ b/docs/efficientad.md @@ -0,0 +1,102 @@ +# EfficientAD + +EfficientAD is available as a native AnomaVision algorithm and uses the same training, model loading, detection, evaluation, and export workflow as PaDiM and PatchCore. + +The implementation follows the EfficientAD student/teacher design: a frozen EfficientNet teacher provides normal feature targets, a lightweight student learns those features, and a compact autoencoder adds a global reconstruction signal. The original EfficientAD paper combines local teacher/student discrepancy with global reconstruction discrepancy for fast anomaly detection. + +## Select the algorithm + +The existing top-level selector remains supported: + +```yaml +algorithm: efficientad +``` + +AnomaVision also accepts the native model selector: + +```yaml +model: + name: efficientad +``` + +For existing projects, changing only `algorithm` is the safest option because the rest of the current configuration remains unchanged. + +## Training + +```bash +anomavision train --config config.yml +``` + +Recommended starting values: + +```yaml +algorithm: efficientad +resize: [224, 224] +normalize: true +batch_size: 1 +efficientad_model_size: s +efficientad_lr: 0.0001 +efficientad_weight_decay: 0.00001 +efficientad_epochs: 1 +efficientad_pretrained_teacher: true +``` + +EfficientAD uses ImageNet preprocessing for the teacher, so `normalize: true` is required by the AnomaVision integration. + +## Detection and thresholding + +Use the same command as the other algorithms: + +```bash +anomavision detect --config config.yml --model model.onnx +``` + +EfficientAD and PaDiM do **not** produce scores on the same numerical scale. PaDiM's threshold (for example `13.0`) must not be reused for EfficientAD. EfficientAD normalizes its image score using the normal-training score mean and standard deviation, so its threshold is expressed in standard deviations from the normal score distribution. + +The default configuration uses an independent EfficientAD threshold: + +```yaml +thresh_padim: 13.0 +thresh_patchcore: 0.25 +thresh_efficientad: 3.0 +``` + +`3.0` is a conservative 3-sigma starting point for deployment. It is intentionally separate from PaDiM and should be calibrated on the validation set for the target MVTec class if you need the closest possible classification agreement with an existing PaDiM deployment. + +## Export + +The trained PyTorch model can be exported through the existing exporter: + +```bash +anomavision export --config config.yml --model model.pt --format onnx +``` + +The ONNX graph contains the EfficientAD inference path, including the teacher, student, autoencoder, score calculation, and anomaly map generation. + +## Evaluation + +```bash +anomavision eval --config config.yml --model model.pt --class_name bottle +``` + +EfficientAD has its own score distribution, so thresholds should be calibrated independently from PaDiM. The inference threshold is controlled by `thresh_efficientad`. + +## Model artifacts + +Training produces the same primary artifact layout used by the other AnomaVision algorithms: + +```text +distributions/ +└── efficientad/ + └── bottle/ + └── anomav_exp/ + ├── model.pt + ├── model.pth + └── config.yml +``` + +`model.pt` is the complete PyTorch model used by the normal AnomaVision inference backend. `model.pth` is a self-contained EfficientAD checkpoint artifact containing the model state and metadata. + +## Difference from PaDiM + +EfficientAD is not numerically interchangeable with PaDiM. PaDiM models feature distributions with Gaussian statistics, while EfficientAD learns a student/teacher representation and reconstruction model. Consequently, scores, thresholds, training time, and localization patterns will differ. What remains intentionally identical is the AnomaVision contract: dataset input, CLI commands, artifact layout, `fit`, `predict`, model loading, and ONNX export. diff --git a/examples/efficientad_cpu.yml b/examples/efficientad_cpu.yml new file mode 100644 index 0000000..157b215 --- /dev/null +++ b/examples/efficientad_cpu.yml @@ -0,0 +1,37 @@ +# EfficientAD quick-start configuration. +# Keep normalize=true: EfficientAD's teacher uses ImageNet preprocessing. +dataset_path: "./dataset" +class_name: "bottle" +resize: [224, 224] +crop_size: null +normalize: true +norm_mean: [0.485, 0.456, 0.406] +norm_std: [0.229, 0.224, 0.225] + +algorithm: "efficientad" +efficientad_model_size: "s" +efficientad_lr: 0.0001 +efficientad_weight_decay: 0.00001 +efficientad_epochs: 1 +efficientad_pretrained_teacher: true + +model_data_path: "./distributions" +output_model: "model.pt" +model: "model.pt" +batch_size: 1 +device: "cpu" +run_name: "efficientad_exp" +log_level: "INFO" + +img_path: "./dataset/bottle/test" +thresh: null +thresh_efficientad: null + +# Visualization/evaluation/export defaults. +enable_visualization: true +save_visualizations: true +viz_output_dir: "./visualizations/" +format: "onnx" +opset: 18 +precision: "fp32" +dynamic_batch: true diff --git a/tests/test_efficientad.py b/tests/test_efficientad.py new file mode 100644 index 0000000..0b300c4 --- /dev/null +++ b/tests/test_efficientad.py @@ -0,0 +1,31 @@ +import torch +from torch.utils.data import DataLoader, TensorDataset + +from anomavision.algorithm.efficientad import EfficientAD + + +def test_efficientad_fit_predict_contract(): + images = torch.rand(2, 3, 224, 224) + loader = DataLoader(TensorDataset(images), batch_size=1, shuffle=False) + + model = EfficientAD( + device=torch.device("cpu"), + pretrained_teacher=False, + model_size="s", + ) + model.fit(loader, epochs=1) + + scores, maps = model.predict(images[:1]) + assert scores.shape == (1,) + assert maps.shape == (1, 224, 224) + assert torch.isfinite(scores).all() + assert torch.isfinite(maps).all() + + +def test_efficientad_rejects_unknown_model_size(): + try: + EfficientAD(pretrained_teacher=False, model_size="large") + except ValueError as exc: + assert "model_size" in str(exc) + else: + raise AssertionError("Expected invalid EfficientAD model_size to fail") diff --git a/tests/test_efficientad_config.py b/tests/test_efficientad_config.py new file mode 100644 index 0000000..095be78 --- /dev/null +++ b/tests/test_efficientad_config.py @@ -0,0 +1,19 @@ +from pathlib import Path + +from anomavision.config import load_config + + +def test_model_name_selects_algorithm(tmp_path: Path): + config_path = tmp_path / "config.yml" + config_path.write_text("model:\n name: efficientad\n") + config = load_config(str(config_path)) + assert config["algorithm"] == "efficientad" + assert config["model"] is None + + +def test_legacy_algorithm_config_is_unchanged(tmp_path: Path): + config_path = tmp_path / "config.yml" + config_path.write_text("algorithm: padim\nmodel: model.pt\n") + config = load_config(str(config_path)) + assert config["algorithm"] == "padim" + assert config["model"] == "model.pt"