initial commit

app.py

@@ -0,0 +1,66 @@
+import torch, torchvision
+from torchvision import transforms
+import numpy as np
+import gradio as gr
+from PIL import Image
+import os
+import lightning as L
+import torchmetrics
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from model import LightningModel
+
+pytorch_model = torch.hub.load('pytorch/vision', 'resnet18', weights=None)
+pytorch_model.fc = torch.nn.Linear(512, 10)
+model_pth = './epoch=22-step=16169.ckpt'
+lightning_model = LightningModel.load_from_checkpoint(checkpoint_path=model_pth, model=pytorch_model, map_location=torch.device("cpu"))
+
+inv_normalize = transforms.Normalize(
+    mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],
+    std=[1/0.23, 1/0.23, 1/0.23]
+)
+classes = ('plane', 'car', 'bird', 'cat', 'deer',
+           'dog', 'frog', 'horse', 'ship', 'truck')
+
+def inference(input_img, transparency = 0.5, target_layer_number = -1):
+    transform = transforms.ToTensor()
+    org_img = input_img
+    input_img = transform(input_img)
+    input_img = input_img
+    input_img = input_img.unsqueeze(0)
+    lightning_model.eval()
+    with torch.no_grad():
+        outputs = lightning_model(input_img)
+        print(f'outputs, {outputs.shape}')
+    softmax = torch.nn.Softmax(dim=0)
+    print()
+    o = softmax(outputs.flatten())
+    confidences = {classes[i]: float(o[i]) for i in range(10)}
+    _, prediction = torch.max(outputs, 1)
+    target_layers = [pytorch_model.layer2[target_layer_number]]
+    cam = GradCAM(model=lightning_model, target_layers=target_layers, use_cuda=False)
+    grayscale_cam = cam(input_tensor=input_img, targets=None)
+    grayscale_cam = grayscale_cam[0, :]
+    img = input_img.squeeze(0)
+    img = inv_normalize(img)
+    rgb_img = np.transpose(img, (1, 2, 0))
+    rgb_img = rgb_img.numpy()
+    visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency)
+    print(confidences)
+    return confidences, visualization
+
+title = "CIFAR10 trained on ResNet18 Model with GradCAM"
+description = "A simple Gradio interface to infer on ResNet model, and get GradCAM results"
+example1 = './cat.jpg'
+example2 = './dog.jpg'
+examples = [[example1, 0.5, -1], [example2, 0.5, -1]]
+gradio_app = gr.Interface(
+    inference,
+    inputs = [gr.Image(shape=(32, 32), label="Input Image"), gr.Slider(0, 1, value = 0.5, label="Opacity of GradCAM"), gr.Slider(-2, -1, value = -2, step=1,
+                                                                                                                                 label="Which Layer?")],
+    outputs = [gr.Label(num_top_classes=3), gr.Image(shape=(32, 32), label="Output").style(width=128, height=128)],
+    title = title,
+    description = description,
+    examples = examples,
+)
+gradio_app.launch()

model.py

@@ -0,0 +1,70 @@
+import torch
+import lightning as L
+import torchmetrics
+
+
+class LightningModel(L.LightningModule):
+    def __init__(self, model, learning_rate, cosine_t_max, mode):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.cosine_t_max = cosine_t_max
+        self.model = model
+        self.example_input_array = torch.Tensor(1, 3, 32, 32)
+        self.mode = mode
+
+        self.save_hyperparameters(ignore=["model"])
+
+        self.train_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10)
+        self.val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10)
+        self.test_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def _shared_step(self, batch):
+        features, true_labels = batch
+        logits = self(features)
+
+        loss = F.cross_entropy(logits, true_labels)
+        predicted_labels = torch.argmax(logits, dim=1)
+        return loss, true_labels, predicted_labels
+
+    def training_step(self, batch, batch_idx):
+        loss, true_labels, predicted_labels = self._shared_step(batch)
+
+        self.log("train_loss", loss)
+        self.train_acc(predicted_labels, true_labels)
+        self.log(
+            "train_acc", self.train_acc, prog_bar=True, on_epoch=True, on_step=False
+        )
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        loss, true_labels, predicted_labels = self._shared_step(batch)
+
+        self.log("val_loss", loss, prog_bar=True)
+        self.val_acc(predicted_labels, true_labels)
+        self.log("val_acc", self.val_acc, prog_bar=True)
+
+    def test_step(self, batch, batch_idx):
+        loss, true_labels, predicted_labels = self._shared_step(batch)
+        self.test_acc(predicted_labels, true_labels)
+        self.log("test_acc", self.test_acc)
+
+    def configure_optimizers(self):
+        opt = torch.optim.SGD(self.parameters(), lr=self.learning_rate)
+        if self.mode == 'lrfind':
+            return opt
+        else:
+            sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=self.cosine_t_max) # New!
+
+            return {
+                "optimizer": opt,
+                "lr_scheduler": {
+                    "scheduler": sch,
+                    "monitor": "train_loss",
+                    "interval": "step", # step means "batch" here, default: epoch
+                    "frequency": 1, # default
+                },
+            }

requirements.txt

@@ -0,0 +1,7 @@
+torch
+gradio
+grad-cam
+lightning
+torchvision
+pillow
+numpy