Update app.py

app.py

@@ -15,7 +15,6 @@ 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 models import custom_resnet
 import gradio as gr
 from pytorch_lightning import LightningModule, Trainer, seed_everything
 from pytorch_lightning.callbacks import LearningRateMonitor
@@ -23,60 +22,9 @@ from pytorch_lightning.callbacks.progress import TQDMProgressBar
 from pytorch_lightning.loggers import CSVLogger
 from pytorch_lightning.loggers import TensorBoardLogger
 from torchmetrics import Accuracy
+from models import custom_resnet
+from network import LitResnet
 
-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")
 
@@ -112,10 +60,10 @@ def inference(input_img, see_misclassified=False,num_misclassified_imgs=0,see_gr
       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()
+      # 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)
     else:
       visualization = org_img
@@ -126,7 +74,6 @@ def inference(input_img, see_misclassified=False,num_misclassified_imgs=0,see_gr
 
 title = "CIFAR10 trained on ResNet18 Model with GradCAM"
 description = "A simple Gradio interface to infer on ResNet model, and get GradCAM results"
-#examples = [["img_eg_0.jpg", False,0,False,0.5, -1,3], ["img_eg_1.jpg", False,0,False,0.5, -1,3],["img_eg_2.jpg", False,0,False,0.5, -1,3],["img_eg_3.jpg", False,0,False,0.5, -1,3],["img_eg_4.jpg", False,0,False,0.5, -1,3],["img_eg_5.jpg", False,0,False,0.5, -1,3],["img_eg_6.jpg", False,0,False,0.5, -1,3],["img_eg_7.jpg", False,0,False,0.5, -1,3],["img_eg_8.jpg", False,0,False,0.5, -1,3]]
 examples = [["img_eg_0.jpg"], ["img_eg_1.jpg"],["img_eg_2.jpg"],["img_eg_3.jpg"],["img_eg_4.jpg"],["img_eg_5.jpg"],["img_eg_6.jpg"],["img_eg_7.jpg"],["img_eg_8.jpg"],["img_eg_9.jpg"]]