Update app.py

app.py

@@ -97,41 +97,40 @@ def model_description():
 def prediction():
 
     def load_model(model_name):
-
-    if model_name == "DenseNet":
-        model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
+        if model_name == "DenseNet":
+            model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
                                      filename="densenet_finetuned.pth")
-        num_classes = 7
-        model = models.densenet121(pretrained=False)
-        model.classifier = torch.nn.Linear(model.classifier.in_features, num_classes)
+            num_classes = 7
+            model = models.densenet121(pretrained=False)
+            model.classifier = torch.nn.Linear(model.classifier.in_features, num_classes)
 
-    elif model_name == "MobileNet":
-        model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
+        elif model_name == "MobileNet":
+            model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
                                      filename="mobileNetV2_finetuned.pth")
-        num_classes = 7
-        model = models.mobilenet_v2(pretrained=False)
-        model.classifier = torch.nn.Linear(model.classifier[1].in_features, num_classes)
+            num_classes = 7
+            model = models.mobilenet_v2(pretrained=False)
+            model.classifier = torch.nn.Linear(model.classifier[1].in_features, num_classes)
 
-    elif model_name == "SqueezeNet":
-        model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
+        elif model_name == "SqueezeNet":
+            model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
                                      filename="squeezenet1_finetuned.pth")
-        num_classes = 7
-        model = models.squeezenet1_1(pretrained=False)
-        model.classifier = torch.nn.Sequential(
-            torch.nn.Dropout(p=0.5),
-            torch.nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1)),
-            torch.nn.ReLU(),
-            torch.nn.AdaptiveAvgPool2d((1, 1))
-        )
-
-    else:
-        raise ValueError("Model not supported.")
+            num_classes = 7
+            model = models.squeezenet1_1(pretrained=False)
+            model.classifier = torch.nn.Sequential(
+                torch.nn.Dropout(p=0.5),
+                torch.nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1)),
+                torch.nn.ReLU(),
+                torch.nn.AdaptiveAvgPool2d((1, 1))
+            )
+
+        else:
+            raise ValueError("Model not supported.")
 
     # Load model weights
-    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
-    model.eval()
+        model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+        model.eval()
     
-    return model
+        return model
 
     def process_image(image):
         if image.mode == 'RGBA':