Update app.py

app.py

@@ -10,7 +10,6 @@ from utils import page_utils
 
 class BasicBlock(nn.Module):
     """ResNet Basic Block.
-
     Parameters
     ----------
     in_channels : int
@@ -63,7 +62,6 @@ class BasicBlock(nn.Module):
 
 class ResNet18(nn.Module):
     """Construct ResNet-18 Model.
-
     Parameters
     ----------
     input_channels : int
@@ -137,8 +135,8 @@ model = ResNet18(3, 7)
 
 checkpoint = torch.load('ham10000.ckpt', map_location=torch.device('cpu'))
 
-# The state dict will contains net.layer_name
-# Our model doesn't contains `net.` so we have to rename it
+# The state dict will contain net.layer_name
+# Our model doesn't contain `net.` so we have to rename it
 state_dict = checkpoint['state_dict']
 for key in list(state_dict.keys()):
     if 'net.' in key:
@@ -147,7 +145,8 @@ for key in list(state_dict.keys()):
 
 model.load_state_dict(state_dict)
 model.eval()
-"""
+
+# Updated class names
 class_names = {
     'akk': 'Actinic Keratosis',
     'bcc': 'Basal Cell Carcinoma',
@@ -157,13 +156,8 @@ class_names = {
     'nv': 'Melanocytic Nevi',
     'vasc': 'Vascular Lesion'
 }
-"""
-class_names = ['akk', 'bcc', 'bkl', 'df', 'mel', 'nv', 'vasc']
-class_names.sort()
-examples_dir = "sample"
-
-
 
+examples_dir = "sample"
 
 transformation_pipeline = transforms.Compose([
     transforms.ToPILImage(),
@@ -173,43 +167,40 @@ transformation_pipeline = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-
 def preprocess_image(image: np.ndarray):
     """Preprocess the input image.
-
     Note that the input image is in RGB mode.
-
     Parameters
     ----------
     image: np.ndarray
         Input image from callback.
     """
-
     image = transformation_pipeline(image)
     image = torch.unsqueeze(image, 0)
-
     return image
 
-
 def image_classifier(inp):
     """Image Classifier Function.
-
     Parameters
     ----------
     inp: Optional[np.ndarray] = None
         Input image from callback
-
     Returns
     -------
     Dict
         A dictionary class names and its probability
     """
-
-    # If input not valid, return dummy data or raise error
+    # If input not valid, return dummy data or raise an error
     if inp is None:
-        if inp is None:
-            return {'cat': 0.3, 'dog': 0.7}
-            #return {'Actinic Keratosis': 0.0, 'Basal Cell Carcinoma': 0.0, 'Benign Keratosis': 0.0, 'Dermatofibroma': 0.0, 'Melanoma': 0.0, 'Melanocytic Nevi': 0.0, 'Vascular Lesion': 0.0}
+        return {
+            'Actinic Keratosis': 0.0,
+            'Basal Cell Carcinoma': 0.0,
+            'Benign Keratosis': 0.0,
+            'Dermatofibroma': 0.0,
+            'Melanoma': 0.0,
+            'Melanocytic Nevi': 0.0,
+            'Vascular Lesion': 0.0
+        }
     # preprocess
     image = preprocess_image(inp)
     image = image.to(dtype=torch.float32)
@@ -218,21 +209,12 @@ def image_classifier(inp):
     result = model(image)
 
     # postprocess
-    result = torch.nn.functional.softmax(result, dim=1) # apply softmax
-    result = result[0].detach().numpy().tolist() # take the first batch
-    labeled_result = {name:score for name, score in zip(class_names, result)}
+    result = torch.nn.functional.softmax(result, dim=1)  # apply softmax
+    result = result[0].detach().numpy().tolist()  # take the first batch
+    labeled_result = {class_names[name]: score for name, score in zip(class_names, result)}
 
     return labeled_result
 
-# gradio code block for input and output
-with gr.Blocks() as app:
-    gr.Markdown("# Skin Cancer Classification")
-
-with open('index.html', encoding="utf-8") as f:
-    description = f.read()
-
-
-    
 # gradio code block for input and output
 with gr.Blocks(theme=gr.themes.Default(primary_hue=page_utils.KALBE_THEME_COLOR, secondary_hue=page_utils.KALBE_THEME_COLOR).set(
         button_primary_background_fill="*primary_600",
@@ -252,10 +234,10 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue=page_utils.KALBE_THEME_COLOR,
             out_txt = gr.Label(label="Probabilities", num_top_classes=3)
 
     process_btn.click(image_classifier, inputs=inp_img, outputs=out_txt)
-    clear_btn.click(lambda:(
+    clear_btn.click(lambda: (
         gr.update(value=None),
         gr.update(value=None)
-        ),
+    ),
         inputs=None,
         outputs=[inp_img, out_txt])
 
@@ -266,7 +248,7 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue=page_utils.KALBE_THEME_COLOR,
                   os.path.join(examples_dir, "bkl_1.jpeg"),
                   os.path.join(examples_dir, "akk.jpeg"),
                   os.path.join(examples_dir, "mel-_3_.jpeg"),
-                 ],
+                  ],
         inputs=inp_img,
         outputs=out_txt,
         fn=image_classifier,
@@ -275,4 +257,5 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue=page_utils.KALBE_THEME_COLOR,
     gr.Markdown(line_breaks=True, value='Author: M HAIKAL FEBRIAN P ([email protected]) <div class="row"><a href="https://github.com/HAikalfebrianp96?tab=repositories"><img alt="GitHub" src="https://img.shields.io/badge/haikal%20phona-000000?logo=github"> </div>')
 
 # demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
-app.launch(share=True)
+app.launch(share=True)
+