supporting presets

app.py

@@ -6,57 +6,56 @@ import matplotlib.pyplot as plt
 import numpy as np
 import torch
 
-from scorer import build_model
+from scorer import build_model, config_presets
 
 
 @cache
-def load_model(device):
-    return build_model(device=device)
+def load_model(preset="edm2-img64-s-fid", device='cpu'):
+    return build_model(preset, device)
 
 @cache
-def load_reference_scores(gmmdir='models'):
-    with np.load(f"{gmmdir}/refscores.npz", "rb") as f:
+def load_reference_scores(model_dir):
+    with np.load(f"{model_dir}/refscores.npz", "rb") as f:
         ref_nll = f["arr_0"]
     return ref_nll
 
-def compute_gmm_likelihood(x_score, gmmdir='models'):
-    with open(f"{gmmdir}/gmm.pkl", "rb") as f:
+def compute_gmm_likelihood(x_score, model_dir):
+    with open(f"{model_dir}/gmm.pkl", "rb") as f:
         clf = load(f)
         nll = -clf.score(x_score)
 
-    ref_nll = load_reference_scores(gmmdir)
+    ref_nll = load_reference_scores(model_dir)
     percentile = (ref_nll < nll).mean() * 100
 
-    return nll, percentile
+    return nll, percentile, ref_nll
 
-def plot_against_reference(nll):
-    ref_nll = load_reference_scores()
-    print(ref_nll.shape)
+def plot_against_reference(nll, ref_nll):
     fig, ax = plt.subplots()
-    ax.hist(ref_nll)
+    ax.hist(ref_nll, label="Reference Scores")
     ax.axvline(nll, label='Image Score', c='red', ls="--")
     plt.legend()
     fig.tight_layout()
     return fig
 
-def run_inference(img, device='cuda'):
+
+def run_inference(img, preset="edm2-img64-s-fid", device="cuda"):
     img = torch.from_numpy(img).permute(2,0,1).unsqueeze(0)
     img = torch.nn.functional.interpolate(img, size=64, mode='bilinear')
-    model = load_model(device=device)
+    model = load_model(preset=preset, device=device)
     x = model(img.cuda())
     x = x.square().sum(dim=(2, 3, 4)) ** 0.5
-    nll, pct = compute_gmm_likelihood(x.cpu())
+    nll, pct, ref_nll = compute_gmm_likelihood(x.cpu(), model_dir=f"models/{preset}")
+
+    plot = plot_against_reference(nll, ref_nll)
 
-    plot = plot_against_reference(nll)
-    print(plot)
-    outstr = f"Anomaly score: {nll:.3f} -> {pct:.2f} percentile"
+    outstr = f"Anomaly score: {nll:.3f} / {pct:.2f} percentile"
     return outstr, plot
 
 
 demo = gr.Interface(
     fn=run_inference,
     inputs=["image"],
-    outputs=["text", gr.Plot()],
+    outputs=["text", gr.Plot(label="Comparing to Imagenette")],
 )
 
 if __name__ == "__main__":

scorer.py

@@ -6,12 +6,21 @@ import numpy as np
 import PIL.Image
 import torch
 from sklearn.mixture import GaussianMixture
+from sklearn.model_selection import GridSearchCV
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import StandardScaler
 from tqdm import tqdm
 
 import dnnlib
 
+model_root = "https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions"
+
+config_presets = {
+    "edm2-img64-s-fid": f"{model_root}/edm2-img64-s-1073741-0.075.pkl",  # fid = 1.58
+    "edm2-img64-m-fid": f"{model_root}/edm2-img64-m-2147483-0.060.pkl",  # fid = 1.43
+    "edm2-img64-l-fid": f"{model_root}/edm2-img64-l-1073741-0.040.pkl",  # fid = 1.33
+}
+
 
 class EDMScorer(torch.nn.Module):
     def __init__(
@@ -34,17 +43,17 @@ class EDMScorer(torch.nn.Module):
         self.net = net.eval()
 
         # Adjust noise levels based on how far we want to accumulate
-        self.sigma_min = sigma_min
+        self.sigma_min = 1e-1
         self.sigma_max = sigma_max * stop_ratio
 
         step_indices = torch.arange(num_steps, dtype=torch.float64, device=device)
         t_steps = (
-            sigma_max ** (1 / rho)
+            self.sigma_max ** (1 / rho)
             + step_indices
             / (num_steps - 1)
-            * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))
+            * (self.sigma_min ** (1 / rho) - self.sigma_max ** (1 / rho))
         ) ** rho
-        print("Using steps:", t_steps)
+        # print("Using steps:", t_steps)
 
         self.register_buffer("sigma_steps", t_steps.to(torch.float64))
 
@@ -61,18 +70,14 @@ class EDMScorer(torch.nn.Module):
             xhat = self.net(x, sigma, force_fp32=force_fp32)
             c_skip = self.net.sigma_data**2 / (sigma**2 + self.net.sigma_data**2)
             score = xhat - (c_skip * x)
-
-            # score_norms = score.mean(1)
-            # score_norms = score.square().sum(dim=(1, 2, 3)) ** 0.5
             batch_scores.append(score)
         batch_scores = torch.stack(batch_scores, axis=1)
 
         return batch_scores
 
 
-def build_model(netpath=f"edm2-img64-s-1073741-0.075.pkl", device="cpu"):
-    model_root = "https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions"
-    netpath = f"{model_root}/{netpath}"
+def build_model(preset="edm2-img64-s-fid", device="cpu"):
+    netpath = config_presets[preset]
     with dnnlib.util.open_url(netpath, verbose=1) as f:
         data = pickle.load(f)
     net = data["ema"]
@@ -80,14 +85,43 @@ def build_model(netpath=f"edm2-img64-s-1073741-0.075.pkl", device="cpu"):
     return model
 
 
-def train_gmm(score_path, outdir="out/msma/"):
+def train_gmm(score_path, outdir):
+    def quantile_scorer(gmm, X, y=None):
+        return np.quantile(gmm.score_samples(X), 0.1)
+
     X = torch.load(score_path)
 
-    gm = GaussianMixture(n_components=5, random_state=42)
+    gm = GaussianMixture(init_params="kmeans", covariance_type="full", max_iter=100000)
     clf = Pipeline([("scaler", StandardScaler()), ("GMM", gm)])
     clf.fit(X)
     inlier_nll = -clf.score_samples(X)
 
+    param_grid = dict(
+        GMM__n_components=range(2, 11, 2),
+    )
+
+    grid = GridSearchCV(
+        estimator=clf,
+        param_grid=param_grid,
+        cv=10,
+        n_jobs=2,
+        verbose=1,
+        scoring=quantile_scorer,
+    )
+
+    grid_result = grid.fit(X)
+
+    print("Best: %f using %s" % (grid_result.best_score_, grid_result.best_params_))
+    print("-----" * 15)
+    means = grid_result.cv_results_["mean_test_score"]
+    stds = grid_result.cv_results_["std_test_score"]
+    params = grid_result.cv_results_["params"]
+    for mean, stdev, param in zip(means, stds, params):
+        print("%f (%f) with: %r" % (mean, stdev, param))
+
+    clf = grid.best_estimator_
+
+    os.makedirs(outdir, exist_ok=True)
     with open(f"{outdir}/refscores.npz", "wb") as f:
         np.savez_compressed(f, inlier_nll)
 
@@ -108,6 +142,7 @@ def compute_gmm_likelihood(x_score, gmmdir):
 
 
 def test_runner(device="cpu"):
+    # f = "doge.jpg"
     f = "goldfish.JPEG"
     image = (PIL.Image.open(f)).resize((64, 64), PIL.Image.Resampling.LANCZOS)
     image = np.array(image)
@@ -118,7 +153,7 @@ def test_runner(device="cpu"):
     return scores
 
 
-def runner(dataset_path, device="cpu"):
+def runner(preset, dataset_path, device="cpu"):
     dsobj = ImageFolderDataset(path=dataset_path, resolution=64)
     refimg, reflabel = dsobj[0]
     print(refimg.shape, refimg.dtype, reflabel)
@@ -126,7 +161,7 @@ def runner(dataset_path, device="cpu"):
         dsobj, batch_size=48, num_workers=4, prefetch_factor=2
     )
 
-    model = build_model(device=device)
+    model = build_model(preset=preset, device=device)
     score_norms = []
 
     for x, _ in tqdm(dsloader):
@@ -137,17 +172,25 @@ def runner(dataset_path, device="cpu"):
     score_norms = torch.cat(score_norms, dim=0)
 
     os.makedirs("out/msma", exist_ok=True)
-    with open("out/msma/imagenette64_score_norms.pt", "wb") as f:
+    with open(f"out/msma/{preset}_imagenette_score_norms.pt", "wb") as f:
         torch.save(score_norms, f)
 
     print(f"Computed score norms for {score_norms.shape[0]} samples")
 
 
 if __name__ == "__main__":
-    # runner("/GROND_STOR/amahmood/datasets/img64/", device="cuda")
-    train_gmm("out/msma/imagenette64_score_norms.pt")
-    s = test_runner(device="cuda")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    preset = "edm2-img64-s-fid"
+    # runner(
+    #     preset=preset,
+    #     dataset_path="/GROND_STOR/amahmood/datasets/img64/",
+    #     device="cuda",
+    # )
+    train_gmm(
+        f"out/msma/{preset}_imagenette_score_norms.pt", outdir=f"out/msma/{preset}"
+    )
+    s = test_runner(device=device)
     s = s.square().sum(dim=(2, 3, 4)) ** 0.5
     s = s.to("cpu").numpy()
-    nll, pct = compute_gmm_likelihood(s, gmmdir="out/msma/")
+    nll, pct = compute_gmm_likelihood(s, gmmdir=f"out/msma/{preset}")
     print(f"Anomaly score for image: {nll[0]:.3f} @ {pct*100:.2f} percentile")