Update app.py

app.py

@@ -97,38 +97,54 @@ def model_description():
 def prediction():
 
     def load_model(model_name):
-        num_classes = 7  # Sesuaikan dengan jumlah kelas mineral
-        
-        if model_name == "DenseNet":
-            model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
-                                         filename="densenet_finetuned.pth")
-            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")
-            model = models.mobilenet_v2(pretrained=False)
-            model.classifier = torch.nn.Linear(model.last_channel, num_classes)  # Perbaikan classifier
-    
-        elif model_name == "SqueezeNet":
-            model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
-                                         filename="squeezenet1_finetuned.pth")
-            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()
-    
-        return model
+    num_classes = 7  # Sesuaikan dengan jumlah kelas mineral yang digunakan
+
+    if model_name == "DenseNet":
+        model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
+                                     filename="densenet_finetuned.pth")
+        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")
+        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")
+        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
+
     
     
     def process_image(image):