Update app.py

app.py

@@ -20,12 +20,15 @@ config = {
     'max_lr': 0.1,
     'max_lr_epoch': 5,
     'dropout' : 0.01,
+    'LEARNING_RATE' : 1e-5,
+    'WEIGHT_DECAY' : 1e-4,
+    'NUM_EPOCHS' : 100
 }
 
 train_transforms = get_train_transforms()
 test_transforms = get_test_transforms()
 model = CustomResNet(config, config['dropout'], train_transforms, test_transforms)
-model.load_state_dict(torch.load("resnet_model_v2.pth", map_location=torch.device('cpu')), strict=False)
+model.load_state_dict(torch.load("resnet_model_v7.pth", map_location=torch.device('cpu')), strict=False)
 model.setup(stage="test")
 
 inv_normalize = transforms.Normalize(
@@ -98,23 +101,32 @@ examples = [[os.path.join(images_folder, "plane.jpg"), 0.5, -1,10],
             [os.path.join(images_folder, "frog.jpg"), 0.5, -1,2],
             [os.path.join(images_folder, "horse.jpg"), 0.5, -1,10],
             [os.path.join(images_folder, "ship.jpg"), 0.5, -1,10],
-            [os.path.join(images_folder, "truck.jpg"), 0.5, -1,10]]
-
+            [os.path.join(images_folder, "truck.jpeg"), 0.5, -1,10]]
 
+# Create the input interface with the modified template
 input_interface = gr.Interface(
     inference,
-    inputs = [gr.Image(shape=(32, 32), label="Input Image"),
-              gr.Slider(0, 1, value = 0.5, label="Transparency",info = "Set the Opacity of CAM"),
-              gr.Slider(-2, -1, value = -2, step=1, label="Network Layer", info = "GradCAM Network Layer"),
-              gr.Slider(1, 10, step=1, value=10 , label="Top Classes",info = "How many top classes do you want to view")],
-    outputs = [gr.Label(num_top_classes=10), gr.Image(shape=(32, 32), label="Model Prediction").style(width=300, height=300)],
-    description = description,
-    examples=[[f'examples/{k}.jpg'] for k in classes_for_categorize.values()],
-)
+    inputs=[
+        gr.Image(shape=(32, 32), label="Input Image"),
+        gr.Slider(0, 1, value=0.5, label="Transparency", info="Set the Opacity of CAM"),
+        gr.Slider(-2, -1, value=-2, step=1, label="Network Layer", info="GradCAM Network Layer"),
+        gr.Slider(1, 10, step=1, value=10, label="Top Classes", info="How many top classes do you want to view")
+    ],
+    outputs=[
+        gr.Label(num_top_classes=10),
+        gr.Image(shape=(32, 32), label="Model Prediction").style(width=300, height=300)
+    ],
+    description=description,
+    examples=[[f'examples/{k}.jpg'] for k in classes_for_categorize.values()],)
 
 mislclassified_description = "Misclassified Image for Custom Resnet"
-
-# Create a separate interface for the "Misclassified Images" 
+icon_html = '<i class="fas fa-chart-bar"></i>'
+title_with_icon = f"""
+<div style="background-color: #f1f4f0; padding: 10px; display: flex; align-items: center;">
+    {icon_html} <span style="margin-left: 10px;">Custom Resnet on CIFAR10 using PyTorch Lightning and GradCAM</span>
+</div>
+"""
+# Create a separate interface for the "Misclassified Images"
 misclassified_interface = gr.Interface(show_misclassified_images_wrap,
                     inputs=[gr.Number(value=10, label="Misclassified Inputs",info = "Set the Number of Misclassifed Outputs to be Shown"),
                             gr.Radio(["Yes", "No"], value="No"  , label="Enable GradCAM",info = "Do you want to see GradCAM"),
@@ -123,7 +135,6 @@ misclassified_interface = gr.Interface(show_misclassified_images_wrap,
                     outputs=gr.Plot(), description=mislclassified_description)
 
 demo = gr.TabbedInterface([input_interface, misclassified_interface], tab_names=["Top Classes and Prediction", "Misclassified Images"],
-                          title="Custom Resnet on CIFAR10 using pytorch Lightening and GradCAM")
-
+                          title=title_with_icon,)
 demo.launch()