Update app.py

app.py

@@ -1,161 +1,153 @@
-import numpy as np
-import matplotlib.pyplot as plt
-from threading import Thread
-from matplotlib.colors import ListedColormap
-from sklearn.datasets import make_moons, make_circles, make_classification
-from sklearn.datasets import make_blobs, make_circles, make_moons
 import gradio as gr
-import math
-from functools import partial
-import time
-
-import matplotlib
-
-from sklearn import svm
-from sklearn.datasets import make_moons, make_blobs
-from sklearn.covariance import EllipticEnvelope
-from sklearn.ensemble import IsolationForest
-from sklearn.neighbors import LocalOutlierFactor
-from sklearn.linear_model import SGDOneClassSVM
-from sklearn.kernel_approximation import Nystroem
-from sklearn.pipeline import make_pipeline
-
-def get_groundtruth_model(X, labels):
-    # dummy model to show true label distribution
-    class Dummy:
-        def __init__(self, y):
-            self.labels_ = labels
-
-    return Dummy(labels)
-
-#### PLOT
-FIGSIZE = 10,10
-figure = plt.figure(figsize=(25, 10))
-
-
-def train_models(input_data,  outliers_fraction, n_samples, clf_name):
-    n_outliers = int(outliers_fraction * n_samples)
-    n_inliers = n_samples - n_outliers
-    blobs_params = dict(random_state=0, n_samples=n_inliers, n_features=2)
-    NAME_CLF_MAPPING = {"Robust covariance": EllipticEnvelope(contamination=outliers_fraction),
-    "One-Class SVM": svm.OneClassSVM(nu=outliers_fraction, kernel="rbf", gamma=0.1),
-    "One-Class SVM (SGD)":make_pipeline(
-            Nystroem(gamma=0.1, random_state=42, n_components=150),
-            SGDOneClassSVM(
-                nu=outliers_fraction,
-                shuffle=True,
-                fit_intercept=True,
-                random_state=42,
-                tol=1e-6,
-            ),
-        ),
-        "Isolation Forest": IsolationForest(contamination=outliers_fraction, random_state=42),
-        "Local Outlier Factor": LocalOutlierFactor(n_neighbors=35, contamination=outliers_fraction),
-    }   
-    DATA_MAPPING = {
-    "Central Blob":make_blobs(centers=[[0, 0], [0, 0]], cluster_std=0.5, **blobs_params)[0],
-    "Two Blobs": make_blobs(centers=[[2, 2], [-2, -2]], cluster_std=[0.5, 0.5], **blobs_params)[0],
-    "Blob with Noise": make_blobs(centers=[[2, 2], [-2, -2]], cluster_std=[1.5, 0.3], **blobs_params)[0],
-    "Moons": 4.0
-    * (
-        make_moons(n_samples=n_samples, noise=0.05, random_state=0)[0]
-        - np.array([0.5, 0.25])
-    ),
-    "Noise": 14.0 * (np.random.RandomState(42).rand(n_samples, 2) - 0.5),
-    }
-    DATASETS = [
-    make_blobs(centers=[[0, 0], [0, 0]], cluster_std=0.5, **blobs_params)[0],
-    make_blobs(centers=[[2, 2], [-2, -2]], cluster_std=[0.5, 0.5], **blobs_params)[0],
-    make_blobs(centers=[[2, 2], [-2, -2]], cluster_std=[1.5, 0.3], **blobs_params)[0],
-    4.0
-    * (
-        make_moons(n_samples=n_samples, noise=0.05, random_state=0)[0]
-        - np.array([0.5, 0.25])
-    ),
-    14.0 * (np.random.RandomState(42).rand(n_samples, 2) - 0.5),
-    ]
+import os
+import sys
+import yaml
+import torch
+from PIL import Image
+import numpy as np
+import torchvision.transforms as transforms # Added this import
+
+# Add the project root to sys.path to allow imports from sibling directories
+# Assuming app.py is in the root of the space, and visual-quality-inspection is a subdirectory
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'visual-quality-inspection')))
+
+# Import your core anomaly detection functions
+# Make sure these imports work relative to the sys.path adjustments
+from visual_quality_inspection.anomaly_detection import load_custom_model, prepare_torchscript_model, inference_score, get_PCA_kernel, get_partial_model, get_train_features # Added get_partial_model, get_train_features
+from visual_quality_inspection.dataset import Mvtec # Your custom dataset class
+
+# --- Configuration Loading ---
+# Define the path to your eval.yaml within the Space
+CONFIG_FILE_PATH = 'visual-quality-inspection/configs/eval.yaml'
+# Define the path where your model is located within the Space
+MODEL_OUTPUT_PATH = 'visual-quality-inspection/models' # This should point to the 'models' directory you create
+
+# Load config once at startup
+with open(CONFIG_FILE_PATH, "r") as f:
+    config = yaml.safe_load(f)
+
+# --- Global Model and PCA Kernel Loading (run once when the app starts) ---
+# This ensures the model is loaded only once, not on every inference call.
+print("Loading model and preparing PCA kernel...")
+
+# Ensure the correct feature_extractor and category_type are set in config
+# This assumes you've pre-modified eval.yaml or you set them here programmatically
+# For this example, let's assume eval.yaml is already set to 'simsiam' and 'bottle' or 'all'
+# If you need to override:
+# config['model']['feature_extractor'] = 'simsiam'
+# config['dataset']['category_type'] = 'bottle' # Or 'all' if you want to iterate
+
+# Load the pre-trained model
+model = load_custom_model(MODEL_OUTPUT_PATH, config)
+if model is None:
+    raise RuntimeError("Failed to load the custom model. Check model path and file integrity.")
+
+# Prepare a dummy dataset for feature shape inference and PCA training
+current_category = config['dataset']['category_type']
+if current_category == 'all':
+    print("Config category is 'all'. Using 'bottle' for initial PCA training for demo purposes.")
+    pca_train_category = 'bottle'
+else:
+    pca_train_category = current_category
+
+trainset = Mvtec(
+    root_dir=config['dataset']['root_dir'],
+    object_type=pca_train_category,
+    split='train',
+    im_size=config['dataset']['image_size']
+)
+
+partial_model, feature_shape = get_partial_model(model, trainset, config['model'])
+model_ts = prepare_torchscript_model(partial_model, config)
+train_features, _ = get_train_features(model_ts, trainset, feature_shape, config)
+pca_kernel = get_PCA_kernel(train_features, config)
+
+print("Model and PCA kernel loaded successfully.")
+
+# --- Anomaly Detection Function for Gradio ---
+def predict_anomaly(input_image: Image.Image, current_category_choice: str):
+    \"\"\"
+    Performs anomaly detection on a single input image for a chosen category.
+    \"\"\"
+    # Ensure the model is in evaluation mode
+    model.eval()
+
+    # Apply the same transformations as defined in Mvtec
+    im_size = config['dataset']['image_size']
+    transform = transforms.Compose([
+        transforms.Resize((im_size, im_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
     
-    xx, yy = np.meshgrid(np.linspace(-7, 7, 150), np.linspace(-7, 7, 150))
-    clf = NAME_CLF_MAPPING[clf_name]
-    plt.figure(figsize=(len(NAME_CLF_MAPPING) * 2 + 4, 12.5))
+    transformed_image = transform(input_image.convert('RGB')).unsqueeze(0) # Add batch dimension
+
+    # Dynamically update category in config for inference if 'all' is chosen or new category
+    # Note: This config change is local to this function call and won't affect global `config`
+    # for subsequent calls, which is fine for Gradio's stateless nature per call.
+    original_category_config = config['dataset']['category_type'] # Store original
+    config['dataset']['category_type'] = current_category_choice # Use user's choice for this inference
+
+    with torch.cpu.amp.autocast(enabled=config['precision']=='bfloat16'):
+        inputs = transformed_image.contiguous(memory_format=torch.channels_last)
+        if config['precision'] == 'bfloat16':
+            inputs = inputs.to(torch.bfloat16)
+
+        features = partial_model(inputs)[config['model']['layer']]
+        pool_out = torch.nn.functional.avg_pool2d(features, config['model']['pool']) if config['model']['pool'] > 1 else features
+        outputs = pool_out.contiguous().view(pool_out.size(0), -1)
+
+        oi = outputs
+        oi_or = oi
+        oi_j = pca_kernel.transform(oi)
+        oi_reconstructed = pca_kernel.inverse_transform(oi_j)
+        fre = torch.square(oi_or - oi_reconstructed).reshape(outputs.shape)
+        fre_score = torch.sum(fre, dim=1)
+        score = -fre_score.item() # Get the single scalar score
+
+    # Revert category_type in config if it was changed (good practice, though not strictly needed for Gradio)
+    config['dataset']['category_type'] = original_category_config
+
+    # Simple anomaly threshold for display
+    # You might want to get a threshold from your eval.yaml or a pre-computed one
+    # For now, a simple rule: if score is very low (highly negative), it's anomalous.
+    # This threshold is illustrative and should be determined from training/validation.
+    ANOMALY_THRESHOLD = -100.0 # Example threshold, adjust based on your model's score range
+    
+    status = "Anomaly Detected!" if score < ANOMALY_THRESHOLD else "Normal"
 
+    return f"Status: {status} | Anomaly Score: {score:.4f}", input_image
 
-    plot_num = 1
-    rng = np.random.RandomState(42)
-    X = DATA_MAPPING[input_data]
-    X = np.concatenate([X, rng.uniform(low=-6, high=6, size=(n_outliers, 2))], axis=0)
-    
-    t0 = time.time()
-    clf.fit(X)
-    t1 = time.time()
-    # fit the data and tag outliers
-    if clf_name == "Local Outlier Factor":
-        y_pred = clf.fit_predict(X)
-    else:
-        y_pred = clf.fit(X).predict(X)
-
-    # plot the levels lines and the points
-    if clf_name != "Local Outlier Factor":  
-        Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])
-        Z = Z.reshape(xx.shape)
-        plt.contour(xx, yy, Z, levels=[0], linewidths=10, colors="black")
-
-    colors = np.array(["#377eb8", "#ff7f00"])
-    plt.scatter(X[:, 0], X[:, 1], s=100, color=colors[(y_pred + 1) // 2])
-
-    plt.xlim(-7, 7)
-    plt.ylim(-7, 7)
-    plt.xticks(())
-    plt.yticks(())
-    plt.text(
-        0.99,
-        0.01,
-        ("%.2fs" % (t1 - t0)).lstrip("0"),
-        transform=plt.gca().transAxes,
-        size=60,
-        horizontalalignment="right",
-    )
-    plot_num += 1
-
-    return plt
-
-description = "Learn how different anomaly detection algorithms perform in different datasets."
-
-def iter_grid(n_rows, n_cols):
-    # create a grid using gradio Block
-    for _ in range(n_rows):
-        with gr.Row():
-            for _ in range(n_cols):
-                with gr.Column():
-                    yield
-
-title = "🕵️‍♀️ compare anomaly detection algorithms 🕵️‍♂️"
-with gr.Blocks() as demo:
-    gr.Markdown(f"## {title}")
-    gr.Markdown(description)
-
-    input_models = ["Robust covariance","One-Class SVM","One-Class SVM (SGD)","Isolation Forest",
-    "Local Outlier Factor"]
-    input_data = gr.Radio(
-        choices=["Central Blob", "Two Blobs", "Blob with Noise", "Moons", "Noise"],
-        value="Moons"
-    )
-    n_samples = gr.Slider(minimum=100, maximum=500, step=25, label="Number of Samples")
-    outliers_fraction = gr.Slider(minimum=0.1, maximum=0.9, step=0.1, label="Fraction of Outliers")
-    counter = 0
-
-
-    for _ in iter_grid(5, 5):
-        if counter >= len(input_models):
-            break
-
-        input_model = input_models[counter]
-        plot = gr.Plot(label=input_model)
-        fn = partial(train_models, clf_name=input_model)
-        input_data.change(fn=fn, inputs=[input_data, outliers_fraction, n_samples], outputs=plot)
-        n_samples.change(fn=fn, inputs=[input_data, outliers_fraction, n_samples], outputs=plot)
-        outliers_fraction.change(fn=fn, inputs=[input_data, outliers_fraction, n_samples], outputs=plot)
-        counter += 1
-
-demo.launch(enable_queue=True, debug=True)
 
+# Get available categories from the data directory
+DATA_ROOT_DIR = config['dataset']['root_dir']
+# Ensure DATA_ROOT_DIR exists before listing
+if not os.path.isdir(DATA_ROOT_DIR):
+    print(f"Warning: Data root directory '{DATA_ROOT_DIR}' not found. Falling back to default categories.")
+    available_categories = ["bottle", "cable", "capsule", "carpet", "grid", "hazelnut", "leather", "metal_nut", "pill", "screw", "tile", "toothbrush", "transistor", "wood", "zipper"]
+else:
+    available_categories = [
+        os.path.basename(d) for d in os.listdir(DATA_ROOT_DIR)
+        if os.path.isdir(os.path.join(DATA_ROOT_DIR, d)) and d not in ['ground_truth'] # Exclude ground_truth if it's a top-level dir
+    ]
+    available_categories.sort()
+
+if not available_categories:
+    available_categories = ["bottle"] # Final fallback if no categories found
+
+# --- Gradio Interface ---
+iface = gr.Interface(
+    fn=predict_anomaly,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image for Anomaly Detection"),
+        gr.Dropdown(choices=available_categories, label="Select Category", value=available_categories[0] if available_categories else "bottle")
+    ],
+    outputs=[
+        gr.Textbox(label="Anomaly Detection Result"),
+        gr.Image(type="pil", label="Input Image")
+    ],
+    title="Visual Anomaly Detection (SimSiam + PCA)",
+    description="Upload an image and select its category to detect anomalies using a pre-trained SimSiam model with PCA-based anomaly scoring. Note: The anomaly threshold is illustrative and may need tuning."
+)
+
+iface.launch()