| """ |
| Unit tests for TrackChain Enhanced PatchCore visual anomaly detection, dataset expansion, and calibration. |
| """ |
| import os |
| import sys |
| import json |
| import tempfile |
| from pathlib import Path |
| import numpy as np |
| import pytest |
| from PIL import Image |
|
|
| |
| repo_root = Path(__file__).resolve().parent.parent.parent |
| if str(repo_root) not in sys.path: |
| sys.path.insert(0, str(repo_root)) |
|
|
| from ml.scripts.expand_datasets import ( |
| generate_track_pattern, |
| create_normal_augmentation_pipeline, |
| augment_normal_image, |
| ensure_diversity, |
| expand_patchcore_dataset, |
| ) |
| from ml.models.vision.anomaly import PatchCoreAnomalyDetector, get_default_transform |
| from ml.training.train_anomaly import greedy_coreset_subsampling, train_patchcore |
| from ml.calibration.patchcore_scale import SigmoidDistanceCalibrator |
| from ml.core.schema import CalibratedSignal, SignalType, DefectClass |
|
|
|
|
| def test_procedural_synthetic_track_generation(): |
| """Verify procedural generation creates a valid 3-channel track image.""" |
| img = generate_track_pattern(width=320, height=320) |
| assert img is not None |
| assert img.shape == (320, 320, 3) |
| assert img.dtype == np.uint8 |
| |
| assert img.std() > 5.0 |
|
|
|
|
| def test_normal_augmentation_pipeline(): |
| """Verify normal track augmentation produces augmented variants without crashing.""" |
| pipelines = create_normal_augmentation_pipeline() |
| assert pipelines is not None |
| assert "lighting" in pipelines |
| assert "weather" in pipelines |
| assert "geometric" in pipelines |
| assert "texture" in pipelines |
|
|
| dummy = np.zeros((100, 100, 3), dtype=np.uint8) + 120 |
| augmented = augment_normal_image(dummy, pipelines, num_augmentations=3) |
| assert len(augmented) == 3 |
| for a in augmented: |
| assert a.shape == (100, 100, 3) |
|
|
|
|
| def test_diversity_clustering_selection(tmp_path): |
| """Verify KMeans feature diversity sampling selects diverse subset.""" |
| img_paths = [] |
| for i in range(10): |
| img = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) |
| p = tmp_path / f"test_{i}.jpg" |
| Image.fromarray(img).save(p) |
| img_paths.append(p) |
|
|
| selected = ensure_diversity(img_paths, target_count=4, num_clusters=3) |
| assert len(selected) == 4 |
| for s in selected: |
| assert Path(s).exists() |
|
|
|
|
| def test_greedy_coreset_minimax_behavior(): |
| """Verify greedy coreset selection reduces sample count and spans maximum distance.""" |
| np.random.seed(42) |
| features = np.random.randn(100, 32).astype(np.float32) |
| coreset = greedy_coreset_subsampling(features, sampling_ratio=0.15) |
|
|
| assert coreset.shape[0] == 15 |
| assert coreset.shape[1] == 32 |
|
|
|
|
| def test_multi_scale_feature_extraction(): |
| """Verify PatchCore multi-scale patch extraction returns correct tensor dimensions.""" |
| detector = PatchCoreAnomalyDetector(backbone_name="resnet18", patch_size=3) |
| import torch |
| dummy_tensor = torch.zeros((1, 3, 224, 224), dtype=torch.float32) |
|
|
| with torch.no_grad(): |
| feats, (h, w) = detector.extract_features(dummy_tensor, patch_size=3) |
|
|
| assert feats.ndim == 2 |
| assert feats.shape[1] == detector.feature_dim |
| assert h > 0 and w > 0 |
|
|
|
|
| def test_patchcore_full_prediction_and_bounding_box(): |
| """Verify full PatchCore inference with calibrated output and localized bounding box.""" |
| detector = PatchCoreAnomalyDetector(backbone_name="resnet18") |
| |
| |
| np.random.seed(42) |
| mock_bank = np.random.randn(50, detector.feature_dim).astype(np.float32) |
| detector.set_memory_bank(mock_bank) |
| detector.calibrator.fit([1.0, 2.0, 3.0, 4.0]) |
|
|
| dummy_frame = np.ones((224, 224, 3), dtype=np.uint8) * 100 |
| signals = detector.predict(dummy_frame) |
|
|
| assert len(signals) == 1 |
| sig = signals[0] |
| assert isinstance(sig, CalibratedSignal) |
| assert sig.signal_type == SignalType.VISUAL_NOVEL |
| assert sig.predicted_class == DefectClass.VISUAL_ANOMALY |
| assert 0.0 <= sig.calibrated_prob <= 1.0 |
|
|
