Update app.py

app.py

@@ -17,8 +17,7 @@ if not os.path.exists(MODEL_PATH):
     urllib.request.urlretrieve(HF_URL, MODEL_PATH)
     print("Download complete!")
 
-# --- DEFINE YOUR MODEL ARCHITECTURE ---
-# (Bagian kelas LookThemLayer, LiteResidualBlock, dan LookThemV8MNIST tetap sama)
+# --- DEFINE YOUR MODEL ARCHITECTURE (sama seperti sebelumnya) ---
 class LookThemLayer(nn.Module):
     def __init__(self, num_tokens, in_features, hidden_dim):
         super().__init__()
@@ -108,12 +107,11 @@ class LookThemV8MNIST(nn.Module):
         x = self.compressor(x).flatten(1)
         x = self.res_blocks(self.input_proj(x))
         return self.head(x)
-# ... (Salin definisi model Anda di sini) ...
 
 # --- LOAD WEIGHTS ON CPU/GPU ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = LookThemV8MNIST()
-model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
+model.load_state_dict(torch.load(MODEL_PATH, map_location=device, weights_only=True))
 model.to(device)
 model.eval()
 
@@ -125,34 +123,59 @@ transform_fn = transforms.Compose([
 ])
 
 def predict_digit(input_image):
+    # Selalu kembalikan dictionary 10 digit untuk gr.Label
+    default_output = {str(i): 0.1 for i in range(10)}
+    
     if input_image is None:
-        return "Please draw a number!"
+        return default_output
     
     try:
-        # gr.Sketchpad mengembalikan numpy array secara langsung
-        img_array = input_image
-
-        # Cek apakah kanvas kosong (semua piksel bernilai 0)
-        if np.max(img_array) == 0:
-            return {str(i): 0.1 for i in range(10)}
-
-        # Konversi ke PIL Image dan resize
-        img = Image.fromarray(img_array.astype(np.uint8), mode="L")
-        img = img.resize((28, 28), Image.Resampling.BILINEAR)
+        # Tangani berbagai format input (dict dari Paint, array dari Sketchpad, dll.)
+        if isinstance(input_image, dict):
+            # gr.Paint versi lama -> ambil composite atau layer pertama
+            img_array = input_image.get("composite")
+            if img_array is None and "layers" in input_image:
+                layers = input_image["layers"]
+                img_array = layers[0] if layers else None
+            if img_array is None:
+                return default_output
+        else:
+            img_array = input_image
+        
+        # Konversi ke numpy array jika belum
+        if not isinstance(img_array, np.ndarray):
+            img_array = np.array(img_array)
+        
+        # Jika gambar berwarna, ambil channel yang tepat
+        if img_array.ndim == 3:
+            if img_array.shape[-1] == 4:      # RGBA → alpha
+                grayscale = img_array[..., 3]
+            else:                             # RGB → luminance
+                grayscale = np.dot(img_array[..., :3], [0.2989, 0.5870, 0.1140])
+        else:
+            grayscale = img_array
         
-        # Transformasi dan prediksi
+        # Cek kanvas kosong
+        if grayscale.max() == 0:
+            return default_output
+        
+        # Resize & normalisasi
+        img = Image.fromarray(grayscale.astype(np.uint8), mode="L")
+        img = img.resize((28, 28), Image.Resampling.BILINEAR)
         tensor_img = transform_fn(img).unsqueeze(0).to(device)
         
         with torch.no_grad():
             outputs = model(tensor_img)
-            probabilities = F.softmax(outputs, dim=1)[0]
-            
-        return {str(i): float(probabilities[i]) for i in range(10)}
-
+            probs = F.softmax(outputs, dim=1)[0]
+        
+        return {str(i): float(probs[i]) for i in range(10)}
+    
     except Exception as e:
-        return {"Error": str(e)}
+        # Kembalikan uniform jika terjadi error tak terduga
+        print(f"Prediction error: {e}")
+        return default_output
 
-# --- GRADIO INTERFACE CONSTRUCTION ---
+# --- GRADIO INTERFACE ---
 with gr.Blocks() as demo:
     gr.Markdown(
         """
@@ -163,15 +186,12 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            # Gunakan gr.Sketchpad
+            # GANTI: gunakan Sketchpad agar latar hitam + pena putih
             input_canvas = gr.Sketchpad(
                 image_mode="L",
                 height=280,
                 width=280,
-                brush=gr.Brush(
-                    default_color="rgb(255, 255, 255)", # Kuas putih
-                    color_mode="fixed"
-                )
+                brush=gr.Brush(default_color="rgb(255,255,255)", color_mode="fixed")
             )
             submit_btn = gr.Button("Classify Digit 🏎️", variant="primary")