Update app.py

app.py

@@ -98,54 +98,54 @@ def prediction():
 
     def load_model(model_name):
 
-    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)
-
-    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)
-
-        # Muat state_dict, tetapi abaikan classifier lama
-        state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-        new_state_dict = {k: v for k, v in state_dict.items() if "classifier.1" not in k}  # Hilangkan classifier lama
-        model.load_state_dict(new_state_dict, strict=False)
-
-        # Definisikan ulang classifier dengan jumlah kelas yang benar
-        model.classifier = torch.nn.Sequential(
-            torch.nn.Dropout(0.2),
-            torch.nn.Linear(model.last_channel, num_classes)
-        )
-
-    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)
-
-        # Muat state_dict, tetapi abaikan classifier lama
-        state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-        new_state_dict = {k: v for k, v in state_dict.items() if "classifier.1" not in k}  # Hilangkan classifier lama
-        model.load_state_dict(new_state_dict, strict=False)
-
-        # Definisikan ulang classifier dengan jumlah kelas yang sesuai
-        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.")
-
-    model.eval()
-    return model
+        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)
+    
+        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)
+    
+            # Muat state_dict, tetapi abaikan classifier lama
+            state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+            new_state_dict = {k: v for k, v in state_dict.items() if "classifier.1" not in k}  # Hilangkan classifier lama
+            model.load_state_dict(new_state_dict, strict=False)
+    
+            # Definisikan ulang classifier dengan jumlah kelas yang benar
+            model.classifier = torch.nn.Sequential(
+                torch.nn.Dropout(0.2),
+                torch.nn.Linear(model.last_channel, num_classes)
+            )
+    
+        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)
+    
+            # Muat state_dict, tetapi abaikan classifier lama
+            state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+            new_state_dict = {k: v for k, v in state_dict.items() if "classifier.1" not in k}  # Hilangkan classifier lama
+            model.load_state_dict(new_state_dict, strict=False)
+    
+            # Definisikan ulang classifier dengan jumlah kelas yang sesuai
+            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.")
+    
+        model.eval()
+        return model