update images

app.py

@@ -3,22 +3,21 @@ Gradio demo of image classification with OOD detection.
 
 If the image example is probably OOD, the model will abstain from the prediction.
 """
-import os
-import pickle
 import json
+import logging
+import pickle
 from glob import glob
 
 import gradio as gr
-from gradio.components import Image, Label, JSON
 import numpy as np
-import torch
 import timm
+import torch
+import torch.nn.functional as F
+from gradio.components import JSON, Image, Label
 from timm.data import resolve_data_config
 from timm.data.transforms_factory import create_transform
 from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names
 
-import logging
-
 _logger = logging.getLogger(__name__)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -26,7 +25,7 @@ TOPK = 3
 
 # load model
 print("Loading model...")
-model = timm.create_model("resnet50", pretrained=True)
+model = timm.create_model("resnet50.tv2_in1k", pretrained=True, checkpoint_path="resnet50.tv2_in1k.bin")
 model.to(device)
 model.eval()
 
@@ -34,28 +33,25 @@ model.eval()
 idx2label = json.loads(open("ilsvrc2012.json").read())
 idx2label = {int(k): v for k, v in idx2label.items()}
 print(idx2label)
+print(idx2label.values())
 
 # transformation
 config = resolve_data_config({}, model=model)
 config["is_training"] = False
 transform = create_transform(**config)
 
-# print features names
-print(get_graph_node_names(model)[0])
-
-# load train scores
+# create feature extractor
 penultimate_features_key = "global_pool.flatten"
 logits_key = "fc"
 features_names = [penultimate_features_key, logits_key]
 
-# create feature extractor
 feature_extractor = create_feature_extractor(model, features_names)
 
-# OOD dtector thresholds
+centroids = torch.from_numpy(pickle.load(open("centroids_resnet50.tv2_in1k_igeood_logits.pkl", "rb"))).to(device)
+# OOD detector thresholds
 msp_threshold = 0.3796
 energy_threshold = 0.3781
-
-## unpickle detectors
+igeood_threshold = 2.4984
 
 
 def mahalanobis_penult(features):
@@ -72,9 +68,18 @@ def energy(logits):
     return torch.logsumexp(logits, dim=1).item()
 
 
+def igeoodlogits_vec(logits, temperature, centroids, epsilon=1e-12):
+    logits = torch.sqrt(F.softmax(logits / temperature, dim=1))
+    centroids = torch.sqrt(F.softmax(centroids / temperature, dim=1))
+    mult = logits @ centroids.T
+    stack = 2 * torch.acos(torch.clamp(mult, -1 + epsilon, 1 - epsilon))
+    return stack.mean(dim=1).item()
+
+
 def predict(image):
     # forward pass
     inputs = transform(image).unsqueeze(0)
+    inputs = inputs.to(device)
     with torch.no_grad():
         features = feature_extractor(inputs)
 
@@ -86,13 +91,16 @@ def predict(image):
 
     result = {idx2label[i.item()]: v.item() for i, v in zip(class_idxs.squeeze(), softmax.squeeze())}
     # OOD
-    msp_score = msp(features[logits_key])
-    energy_score = energy(features[logits_key])
+    msp_score = round(msp(features[logits_key]), 4)
+    energy_score = round(energy(features[logits_key]), 4)
+    igeood_scores = round(igeoodlogits_vec(features[logits_key], 1, centroids), 4)
     ood_scores = {
-        "msp": msp_score,
-        "msp_is_ood": msp_score < msp_threshold,
-        "energy": energy_score,
-        "energy_is_ood": energy_score < energy_threshold,
+        "MSP": msp_score,
+        "MSP, is the input OOD?": msp_score < msp_threshold,
+        "Energy": energy_score,
+        "Energy, is the input OOD?": energy_score < energy_threshold,
+        "Igeood": igeood_scores,
+        "Igeood, is the input OOD?": igeood_scores < igeood_threshold,
     }
     _logger.info(ood_scores)
     return result, ood_scores
@@ -100,9 +108,9 @@ def predict(image):
 
 def main():
     # image examples for demo shuffled
-    examples = glob("images/imagenet/*.jpg") + glob("images/ood/*.jpg")
+    examples = glob("images/imagenet/*") + glob("images/ood/*")
     np.random.seed(42)
-    np.random.shuffle(examples)
+    # np.random.shuffle(examples)
 
     # gradio interface
     interface = gr.Interface(
@@ -117,11 +125,24 @@ def main():
         allow_flagging="never",
         theme="default",
         title="OOD Detection 🧐",
-        description="Out-of-distribution (OOD) detection is an essential safety measure for machine learning models. This app demonstrates how these methods can be useful. They try to determine wether we can trust the predictions of a ResNet-50 model trained on ImageNet-1K. Enjoy the demo!",
-    )
-    interface.launch(
-        server_port=7860,
+        description=(
+            "Out-of-distribution (OOD) detection is an essential safety measure for machine learning models. "
+            "The objective of an OOD detector is to determine wether the input sample comes from the distribution known by the AI model. "
+            "For instance, an input that does not belong to any of the known classes or is from a different domain should be flagged by the detector.\n"
+            "In this demo we will display the decision of three OOD detectors on a ResNet-50 model trained to classify on the ImageNet-1K dataset (top-1 accuracy 80%)."
+            "This model can classify among 1000 classes from several categories, including `animals`, `vehicles`, `clothing`, `instruments`, `plants`, etc. "
+            "For the complete hierarchy of classes, please check the website https://observablehq.com/@mbostock/imagenet-hierarchy. "
+            "\n\n"
+            "## Instructions:\n"
+            "1. Upload an image of your choice or select one from the examples bar.\n"
+            "2. The model will predict the top 3 most likely classes for the image.\n"
+            "3. The OOD detectors will output their scores and decision on the image. The smaller the score, the least confident the detector is on the sample being in-distribution.\n"
+            "4. If the image is OOD, the model will abstain from the prediction and flag it to the practicioner.\n"
+            "\n\n\nEnjoy the demo!"
+        ),
+        cache_examples=True,
     )
+    interface.launch(server_port=7860)
     interface.close()