Update app.py

app.py

@@ -97,17 +97,18 @@ def model_description():
 def prediction():
 
     def load_model(model_name):
-    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")
+        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
@@ -124,6 +125,7 @@ def prediction():
     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