Update app.py

app.py

@@ -123,49 +123,50 @@ transform_fn = transforms.Compose([
     transforms.ToTensor(),
     transforms.Normalize((0.1307,), (0.3081,))
 ])
-
 def predict_digit(input_image):
     if input_image is None:
         return "Please draw a number!"
-    
+
     try:
-        # 1. Ambil layer coretan pertama langsung seperti kode referensi
-        if isinstance(input_image, dict) and "layers" in input_image and len(input_image["layers"]) > 0:
-            img_array = input_image["layers"][0]
-        elif isinstance(input_image, dict) and "composite" in input_image:
-            img_array = input_image["composite"]
+        # Versi aman: ambil composite jika ada (numpy array HxW atau HxWxC)
+        if isinstance(input_image, dict):
+            # Beberapa versi Gradio meletakkan hasil akhir di 'composite'
+            img_array = input_image.get("composite", input_image["layers"][0])
         else:
             img_array = input_image
 
-        # 2. Pastikan bentuknya 2D Grayscale
-        if len(img_array.shape) == 3:
-            if img_array.shape[-1] == 4: # Jika RGBA, ambil alpha channel (coretannya)
-                grayscale = img_array[:, :, 3]
-            else: # Jika RGB, konversi ke grayscale
-                grayscale = np.dot(img_array[...,:3], [0.2989, 0.5870, 0.1140])
+        # Konversi ke grayscale 2D
+        if isinstance(img_array, np.ndarray):
+            if img_array.ndim == 3:
+                if img_array.shape[-1] == 4:      # RGBA -> ambil alpha
+                    grayscale = img_array[..., 3]
+                else:                             # RGB -> luminance
+                    grayscale = np.dot(img_array[..., :3], [0.2989, 0.5870, 0.1140])
+            else:
+                grayscale = img_array
         else:
-            grayscale = img_array
+            # Kalau ternyata PIL.Image
+            grayscale = np.array(img_array.convert("L"))
 
-        # 3. Cek jika kanvas kosong
-        if np.max(grayscale) == 0:
+        # 3. Cek kanvas kosong (kali ini cek nilai maks > 0, bukan == 0)
+        if grayscale.max() == 0:
             return {str(i): 0.1 for i in range(10)}
 
-        # 4. Konversi ke PIL dan Resize ke 28x28 (Gaya kode referensimu)
-        img = Image.fromarray(grayscale.astype(np.uint8)).convert('L')
-        img = img.resize((28, 28), resample=Image.Resampling.BILINEAR)
-        
-        # 5. Jalankan normalisasi PyTorch sesuai training LookThem V8 kamu
+        # 4. Resize dan 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)
-        
-        # 6. Prediksi dengan Otak Sniper LookThem
+
+        # 5. Prediksi
         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)}
 
     except Exception as e:
-        return {"Error": 0.0}
+        # Untuk debug, kembalikan pesan errornya
+        return {"error": str(e)} 
  
 
 # --- GRADIO INTERFACE CONSTRUCTION ---
@@ -179,12 +180,12 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            # Menggunakan gr.Paint dengan mode L (Luminance/Grayscale)
             input_canvas = gr.Paint(
-                image_mode="L", 
+                image_mode="L",
                 height=280,
                 width=280,
-                brush=gr.components.image_editor.Brush(default_color="rgb(255, 255, 255)") # Kuas warna putih
+                canvas_color="black",   # ⬅️ ini yang wajib ditambahkan
+                brush=gr.components.image_editor.Brush(default_color="rgb(255, 255, 255)")
             )
             submit_btn = gr.Button("Classify Digit 🏎️", variant="primary")