Upload 6 files

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+ex2.jpg filter=lfs diff=lfs merge=lfs -text

gradio.py

@@ -0,0 +1,169 @@
+import warnings
+warnings.filterwarnings('ignore')
+import torch, numpy as np, os
+from torch import nn
+from transformers import AutoModelForImageClassification, AutoConfig, AutoImageProcessor
+import matplotlib.pyplot as plt
+from PIL import Image
+import saliency.core as saliency
+import io
+import gradio as gr
+import PIL
+
+model_choice = 0
+model_names = ["nvidia/mit-b0",'facebook/convnext-base-224', 'microsoft/resnet-18', 'microsoft/swin-tiny-patch4-window7-224']
+model_name = model_names[model_choice]
+device     = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+class Model(nn.Module):
+    def __init__(self, MODEL_NAME=model_name):
+        super().__init__()
+        self.config = AutoConfig.from_pretrained(MODEL_NAME, finetuning_task="image-classification")
+        self.model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
+        self.class_len = self.config.num_labels
+        self.id2label = self.config.id2label
+        self.label2id = self.config.label2id
+
+    def forward(self, x):
+        if isinstance(x, np.ndarray):
+            x = torch.from_numpy(x)
+        if len(x.shape) == 3:
+            x = x.unsqueeze(0)
+        if x.shape[-1] == 3:
+            x = x.permute(0, 3, 1, 2)
+        x = x.to(device)
+        x = self.model(x)
+        return x.logits
+
+def conv_layer_forward_hook(module, input, output):
+    """Method from Examples_pytorch.ipynb for the gradcam library https://github.com/PAIR-code/saliency."""
+    global last_conv_layer_outputs
+    last_conv_layer_outputs[saliency.base.CONVOLUTION_LAYER_VALUES] = torch.movedim(output, 3, 1).detach().cpu().numpy()
+def conv_layer_backward_hook(module, grad_input, grad_output):
+    """Method from Examples_pytorch.ipynb for the gradcam library https://github.com/PAIR-code/saliency."""
+    global last_conv_layer_outputs
+    last_conv_layer_outputs[saliency.base.CONVOLUTION_OUTPUT_GRADIENTS] = torch.movedim(grad_output[0], 3, 1).detach().cpu().numpy()
+
+auto_transformer, class_to_id, id_to_class, last_conv_layer, last_conv_layer_outputs = None, None, None, None, None
+
+
+def swap_models(name):
+    global model, auto_transformer, class_to_id, id_to_class, last_conv_layer, last_conv_layer_outputs
+    auto_transformer = AutoImageProcessor.from_pretrained(name)
+    model = Model(MODEL_NAME=name)
+    model = model.to(device).eval()
+    # register the hooks for the last convolution layer for Grad-Cam
+    named_modules = dict(model.model.named_modules())
+    last_conv_layer_name = None
+    for name, module in named_modules.items():
+        if isinstance(module, torch.nn.Conv2d):
+            last_conv_layer_name = name
+
+    last_conv_layer = named_modules[last_conv_layer_name]
+    last_conv_layer_outputs = {}
+
+    last_conv_layer.register_forward_hook(conv_layer_forward_hook)
+    last_conv_layer.register_backward_hook(conv_layer_backward_hook)
+    class_to_id = {v:k for k,v in model.model.config.id2label.items()}
+    id_to_class = {k:v for k,v in model.model.config.id2label.items()}
+
+swap_models(model_name)
+
+def saliency_graph(img1, steps=120):
+    img1 = auto_transformer(img1)
+    img1 = np.squeeze(np.array(img1.pixel_values))
+    if img1.shape[0] < img1.shape[1]:
+        img1 = np.moveaxis(img1, 0, -1)
+    img1 = (img1 - np.min(img1)) / (np.max(img1) - np.min(img1))
+
+    class_idx_str = 'class_idx_str'
+    def gradcam_call(images, call_model_args=None, expected_keys=None):
+        if not isinstance(images, np.ndarray) and not isinstance(images, torch.Tensor) and not isinstance(images, PIL.Image.Image):
+            # return two blank images
+            im1 = np.zeros((224, 224, 3))
+            im2 = np.zeros((224, 224, 3))
+            return im1, im2
+        
+        if len(images.shape) == 3:
+            images = np.expand_dims(images, 0)
+        images = torch.tensor(images, dtype=torch.float32)
+        images = images.requires_grad_(True)
+        target_class_idx = call_model_args[class_idx_str]
+        y_pred = model(images)
+        
+        if saliency.base.INPUT_OUTPUT_GRADIENTS in expected_keys:
+            out =  y_pred[:, target_class_idx]
+            # move actual color channel to the 1st dimension
+            #images = torch.movedim(images, 3, 1)
+            grads = torch.autograd.grad(out, images, grad_outputs=torch.ones_like(out))
+            grads = grads[0].detach().cpu().numpy()
+            return {saliency.base.INPUT_OUTPUT_GRADIENTS: grads}
+        else:
+            hot = torch.zeroes_like(y_pred)
+            hot[:, target_class_idx] = 1
+            model.zero_grad()
+            y_pred.backward(gradient=hot, retain_graph=True)
+            return last_conv_layer_outputs
+
+    im = img1.astype(np.float32)
+    base = np.zeros(img1.shape)
+
+    pred = model(torch.from_numpy(im))
+    class_pred = pred.argmax(dim=1).item()
+    call_model_args = {class_idx_str: class_pred}
+    gradients = saliency.IntegratedGradients()
+
+    s = gradients.GetSmoothedMask(im, gradcam_call, call_model_args, x_steps=steps, x_baseline=base, batch_size=25)
+
+    smoothgrad_mask_grayscale = saliency.VisualizeImageGrayscale(s)
+
+    with torch.no_grad():
+        output = model.forward(img1)
+        output = torch.nn.functional.softmax(output, dim=1)
+        output = output.cpu().numpy()
+    top_5 = [(id_to_class[int(i)], output[0][i]) for i in np.argsort(output)[0][-5:][::-1]]
+
+
+    # Render the saliency masks.
+    fig, ax = plt.subplots(1, 1, figsize=(10, 10))
+    ax.barh([x[0] for x in top_5], [x[1] for x in top_5])
+    ax.set_title('Top 5 Predictions')
+    buf = io.BytesIO()
+    fig.savefig(buf, format='jpg')
+    buf.seek(0)
+    fig_img = Image.open(buf)
+    plt.close(fig)
+    return smoothgrad_mask_grayscale, fig_img
+
+# gradio Interface
+def gradio_interface(img):
+    smoothgrad_mask_grayscale, fig_img = saliency_graph(img, steps=25)
+    return smoothgrad_mask_grayscale, fig_img
+
+with gr.Blocks(title='Looking at the pixels models attend to', description="This function finds the most critical pixels in an image for predicting a class. The best models will ideally make predictions by highlighting the expected object. Poorly generalizable models will often rely on environmental cues instead and forego looking at the most important pixels. Highlighting the most important pixels helps explain/build trust about whether a given model uses the correct features to make its prediction.", live=True) as iface:
+    #examples = gr.Examples(examples=["ex1.jpg", "ex2.jpg", "ex3.jpg", "ex4.jpg"], label="Examples", inputs="image", examples_per_page=4)
+    gr.Markdown("Choose a model to use for classying images:")
+    with gr.Row():
+        with gr.Column():
+            test_image = gr.Image(label="Input Image", live=True)
+            input_btn = gr.Button(label="Classify image")
+            model_select_dropdown = gr.Radio(model_names, label="Model to test", interactive=True, default=0)
+        with gr.Column():
+            output = gr.Image(label="Pixels used for classification")
+            output2 = gr.Image(label="Top 5 Predictions")
+
+    input_btn.click(gradio_interface, test_image, outputs=[output, output2])
+    model_select_dropdown.change(swap_models, inputs=[model_select_dropdown])
+    examples = gr.Examples(
+        examples = [os.path.join('./', x) for x in os.listdir('./') if x.endswith('.jpg')],
+        inputs=gr.Image(),
+        label="Examples",
+        fn=gradio_interface,
+        cache_examples=False,
+        run_on_click=True,
+        postprocess=True,
+        preprocess=True,
+        outputs=[output, output2])
+
+
+iface.launch()

requirements.txt

@@ -0,0 +1,8 @@
+torch
+numpy
+matplotlib
+pillow
+matplotlib
+saliency
+gradio
+transformers