Update app.py

app.py

@@ -1,6 +1,5 @@
 import torch
 import torch.nn as nn
-import torch.optim as optim
 import torchvision.transforms as transforms
 from PIL import Image
 import gradio as gr
@@ -10,16 +9,19 @@ class SimpleCNN(nn.Module):
     def __init__(self):
         super(SimpleCNN, self).__init__()
         self.conv_layers = nn.Sequential(
-            nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),
+            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
             nn.ReLU(),
             nn.MaxPool2d(kernel_size=2, stride=2),
-            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
+            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
             nn.ReLU(),
             nn.MaxPool2d(kernel_size=2, stride=2)
         )
         self.fc_layers = nn.Sequential(
             nn.Flatten(),
-            nn.Linear(32 * 32 * 32, 128),
+            nn.Linear(128 * 16 * 16, 128),  # Pastikan ukuran input sesuai
             nn.ReLU(),
             nn.Linear(128, 6)  # Output layer untuk 6 kelas
         )
@@ -32,9 +34,13 @@ class SimpleCNN(nn.Module):
 # Inisialisasi model
 model = SimpleCNN()
 
-# Load model
-model.load_state_dict(torch.load("model_deri.pth", map_location=torch.device("cpu")))
-model.eval()
+# Load model dengan error handling
+try:
+    model.load_state_dict(torch.load("model_deri.pth", map_location=torch.device("cpu")), strict=False)
+    model.eval()
+    print("✅ Model berhasil dimuat!")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
 
 # Kelas mapping
 class_mapping = {0: 'Bu dian', 1: 'Deri', 2: 'Putra', 3: 'Unknown', 4: 'Uqi', 5: 'Uwa'}
@@ -50,11 +56,17 @@ def predict(image):
     image = transform(image).unsqueeze(0)  # Tambah batch dimension
     with torch.no_grad():
         output = model(image)
-        predicted_class = torch.argmax(output, dim=1).item()
-    return f"Predicted: {class_mapping[predicted_class]}"
+        probabilities = torch.nn.functional.softmax(output, dim=1)
+        predicted_class = torch.argmax(probabilities, dim=1).item()
+        confidence = probabilities[0, predicted_class].item() * 100  # Konversi ke persen
+    return f"Predicted: {class_mapping[predicted_class]} (Confidence: {confidence:.2f}%)"
 
 # Buat UI Gradio
 iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs="text")
 
 # Jalankan aplikasi
-iface.launch()
+if __name__ == "__main__":
+    try:
+        iface.launch()
+    except Exception as e:
+        print(f"❌ Gradio error: {e}")