Create app.py

app.py

@@ -0,0 +1,120 @@
+import torch, torchvision
+from torchvision import transforms
+import numpy as np
+import gradio as gr
+from PIL import Image
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from resnet import ResNet18
+import gradio as gr
+
+class LitResnet(LightningModule):
+    def __init__(self, num_classes=10, lr=0.05):
+        super().__init__()
+
+        self.save_hyperparameters()
+        self.model = custom_resnet.Net()
+        self.criterion = nn.CrossEntropyLoss()
+        self.BATCH_SIZE = 512
+        self.torchmetrics_accuracy = Accuracy(task="multiclass", num_classes= self.hparams.num_classes)
+
+    def forward(self, x):
+        out = self.model(x)
+        return out
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_pred = self(x)
+        loss = self.criterion(y_pred, y)
+        acc  = self.torchmetrics_accuracy(y_pred, y)
+
+        self.log('train_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
+        self.log('train_acc', acc, prog_bar=True, on_step=False, on_epoch=True)
+        return loss
+
+
+    def evaluate(self, batch, stage=None):
+        x, y = batch
+        y_test_pred = self(x)
+        loss = self.criterion(y_test_pred, y)
+        acc  = self.torchmetrics_accuracy(y_test_pred, y)
+
+        if stage:
+            self.log(f"{stage}_loss", loss, prog_bar=True)
+            self.log(f"{stage}_acc", acc, prog_bar=True)
+
+    def test_step(self, batch, batch_idx):
+        self.evaluate(batch, "test")
+
+    def validation_step(self, batch, batch_idx):
+        self.evaluate(batch, "val")
+
+    def configure_optimizers(self):
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr, weight_decay=1e-4)
+        scheduler = OneCycleLR(
+                optimizer,
+                max_lr= 5.38E-02, #self.hparams.lr,
+                pct_start = 5/self.trainer.max_epochs,
+                epochs=self.trainer.max_epochs,
+                steps_per_epoch=len(train_loader),
+                div_factor=100,verbose=False,
+                three_phase=False
+            )
+        return ([optimizer],[scheduler])
+
+inference_model = LitResnet.load_from_checkpoint("cifar10_customresnet_20_epoch.ckpt")
+
+def inference(input_img, see_misclassified,num_misclassified_imgs,see_gradcam,num_gradcam_imgs,transparency = 0.85, target_layer_number = -1,top_classes=3):
+
+    if see_misclassified: # show misclassified images
+      org_img   =  cv2.imread('/content/drive/MyDrive/AI/ERA_course/session12/example_images/img_eg_0.jpg')
+      input_img = org_img
+
+    elif num_gradcam_imgs > 0: # show gradcam on example images
+      org_img   =  cv2.imread('/content/drive/MyDrive/AI/ERA_course/session12/example_images/img_eg_0.jpg')
+      input_img = org_img
+
+    else: # nothing chosen - misclassified or gradcam
+      org_img = input_img
+
+    # model inference
+    transform = transforms.ToTensor()
+    input_img = transform(input_img)
+    input_img = input_img.unsqueeze(0)
+    outputs = inference_model.model(input_img)
+    softmax = torch.nn.Softmax(dim=0)
+    o = softmax(outputs.flatten())
+    confidences = {classes[i]: float(o[i]) for i in range(10)}
+    _, prediction = torch.max(outputs, 1)
+
+    # gradcam
+    if see_gradcam:
+      target_layers = [inference_model.model.layer2[target_layer_number]]
+      cam = GradCAM(model=inference_model.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.0, grayscale_cam, use_rgb=True, image_weight=transparency)
+      plt.imshow(visualization)
+    else:
+      plt.imshow(org_img)
+      visualization = org_img
+
+    # top n classes only
+    confidences = {k: confidences[k] for k in list(confidences)[:top_classes]}
+    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"
+
+demo = gr.Interface(
+    inference,
+    inputs =  [gr.Image(shape=(32, 32), label="Input Image"), gr.Checkbox(label="Misclassified"),gr.Slider(0, 10, value = 0, step=1,label="Total Misclassified Images"),gr.Checkbox(label="Gradcam"),gr.Slider(0, 10, value = 0, step=1,label="Total GradCam Images"),gr.Slider(0, 1, value = 0.5, label="Opacity of GradCAM"), gr.Slider(-2, -1, value = -1, step=1, label="Which Layer?"), gr.Slider(1, 10, value=3, step=1, label="How many top classes?")],
+    outputs = [gr.Label(), gr.Image(shape=(32, 32), label="Output").style(width=128, height=128)],
+    title = title,
+    description = description,)
+
+demo.launch()