Update app.py

app.py

@@ -97,24 +97,23 @@ def model_description():
 def prediction():
 
     def load_model(model_name):
+        num_classes = 7  # Pastikan sesuai dengan jumlah kelas yang digunakan saat training
+    
         if model_name == "DenseNet":
             model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
-                                     filename="densenet_finetuned.pth")
-            num_classes = 7
+                                         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")
-            num_classes = 7
+                                         filename="mobileNetV2_finetuned.pth")
             model = models.mobilenet_v2(pretrained=False)
-            model.classifier = torch.nn.Linear(model.classifier[1].in_features, num_classes)
-
+            model.classifier = torch.nn.Linear(model.classifier[0].in_features, num_classes)  # Fix in_features
+    
         elif model_name == "SqueezeNet":
             model_path = hf_hub_download(repo_id="ahmadalfian/mineral-classifier",
-                                     filename="squeezenet1_finetuned.pth")
-            num_classes = 7
+                                         filename="squeezenet1_finetuned.pth")
             model = models.squeezenet1_1(pretrained=False)
             model.classifier = torch.nn.Sequential(
                 torch.nn.Dropout(p=0.5),
@@ -122,47 +121,48 @@ def prediction():
                 torch.nn.ReLU(),
                 torch.nn.AdaptiveAvgPool2d((1, 1))
             )
-
         else:
             raise ValueError("Model not supported.")
-
-    # Load model weights
+    
+        # Load model weights
         model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
         model.eval()
     
         return model
-
+    
+    
     def process_image(image):
+        """Konversi gambar dan lakukan preprocessing sebelum masuk ke model"""
         if image.mode == 'RGBA':
             image = image.convert('RGB')
-            print("Image converted from RGBA to RGB.")
-
+    
         preprocess = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
-
+    
         img_tensor = preprocess(image)
-        print(f"Image tensor shape: {img_tensor.shape}")
-        return img_tensor.unsqueeze(0)
-
+        return img_tensor.unsqueeze(0)  # Tambahkan dimensi batch
+    
+    
     def classify_image(model, image):
-        img_tensor = process_image(image)
-
-        print(f"Image tensor shape after unsqueeze: {img_tensor.shape}")
-
+        """Lakukan prediksi menggunakan model"""
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        img_tensor = process_image(image).to(device)  # Pastikan berada di perangkat yang sesuai
+    
+        model.to(device)
         model.eval()
-
+    
         with torch.no_grad():
             outputs = model(img_tensor)
-
-            print(f"Model output shape: {outputs.shape}")
-
+    
+            # Konversi hasil ke probabilitas
             probabilities = torch.nn.functional.softmax(outputs, dim=1)
-
+    
+            # Ambil prediksi dengan confidence tertinggi
             confidence, predicted = torch.max(probabilities, 1)
-
+    
         return predicted.item(), confidence.item()