Update app.py

app.py

@@ -18,96 +18,8 @@ if not os.path.exists(MODEL_PATH):
     print("Download complete!")
 
 # --- DEFINE YOUR MODEL ARCHITECTURE ---
-
-class LookThemLayer(nn.Module):
-    def __init__(self, num_tokens, in_features, hidden_dim):
-        super().__init__()
-        self.num_tokens = num_tokens
-        self.mod1_w1 = nn.Parameter(torch.randn(num_tokens, in_features, hidden_dim))
-        self.mod1_b1 = nn.Parameter(torch.zeros(num_tokens, hidden_dim))
-        self.mod1_w2 = nn.Parameter(torch.randn(num_tokens, hidden_dim, 1))
-        self.mod1_b2 = nn.Parameter(torch.zeros(num_tokens, 1))
-        self.mod2_w1 = nn.Parameter(torch.randn(num_tokens, in_features, hidden_dim))
-        self.mod2_b1 = nn.Parameter(torch.zeros(num_tokens, hidden_dim))
-        self.mod2_w2 = nn.Parameter(torch.randn(num_tokens, hidden_dim, 1))
-        self.mod2_b2 = nn.Parameter(torch.zeros(num_tokens, 1))
-        self.trans_w = nn.Parameter(torch.randn(num_tokens, 1, 1))
-        self.trans_b = nn.Parameter(torch.zeros(num_tokens, 1))
-        self._init_weights()
-
-    def _init_weights(self):
-        for w in [self.mod1_w1, self.mod2_w1, self.mod1_w2, self.mod2_w2]:
-            nn.init.xavier_uniform_(w)
-
-    def forward(self, x):
-        N = self.num_tokens
-        h1 = torch.einsum("bti,tij->btj", x, self.mod1_w1) + self.mod1_b1
-        out_m1 = torch.einsum("btj,tjk->btk", F.gelu(h1), self.mod1_w2) + self.mod1_b2
-        h2 = torch.einsum("bti,tij->btj", x, self.mod2_w1) + self.mod2_b1
-        out_m2 = torch.einsum("btj,tjk->btk", F.gelu(h2), self.mod2_w2) + self.mod2_b2
-
-        out_m2_safe = torch.sign(out_m2) * torch.clamp(torch.abs(out_m2), min=1e-6)
-        compare = torch.tanh(out_m1.unsqueeze(2) / out_m2_safe.unsqueeze(1))
-        compare2 = torch.tanh(out_m1.unsqueeze(1) / out_m2_safe.unsqueeze(2))
-
-        trans_compare = torch.einsum("bije,jef->bijf", compare, self.trans_w) + self.trans_b.view(1, 1, N, 1)
-        trans_compare2 = torch.einsum("bije,jef->bijf", compare2, self.trans_w) + self.trans_b.view(1, 1, N, 1)
-
-        interaksi = (trans_compare * x.unsqueeze(2) + trans_compare2 * x.unsqueeze(1)) / 2
-        mask = (1.0 - torch.eye(N, device=x.device)).view(1, N, N, 1)
-        return (interaksi * mask).sum(dim=2) / (N - 1.0)
-
-class LiteResidualBlock(nn.Module):
-    def __init__(self, dim, dropout=0.05):
-        super().__init__()
-        self.block = nn.Sequential(nn.Linear(dim, dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim, dim))
-        self.norm = nn.LayerNorm(dim)
-    def forward(self, x):
-        return self.norm(x + self.block(x))
-
-class LookThemV8MNIST(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.stream_a = nn.Sequential(
-                nn.Conv2d(1, 4, 3, 2, 1),
-                nn.BatchNorm2d(4), nn.GELU(),
-                nn.Conv2d(4, 8, 3, 2, 1),
-                nn.BatchNorm2d(8), nn.GELU(),
-                nn.AdaptiveMaxPool2d((8, 8)))
-        self.stream_b = nn.Sequential(
-                nn.Conv2d(1, 4, 3, 1, 1),
-                nn.BatchNorm2d(4), nn.GELU(),
-                nn.Conv2d(4, 8, 3, 1, 1),
-                nn.BatchNorm2d(8), nn.GELU(),
-                nn.AdaptiveMaxPool2d((8, 8)))
-
-        self.lookthemA = LookThemLayer(64, 8, 32)
-        self.lookthemB = LookThemLayer(64, 8, 32)
-        self.lookthem_comb = LookThemLayer(64, 16, 32)
-        self.comb_norm = nn.LayerNorm(16)
-
-        self.FFN1 = nn.Conv1d(16, 8, 1)
-        self.lookthem2 = LookThemLayer(64, 8, 32)
-        self.FFN2 = nn.Conv1d(8, 8, 1)
-
-        self.compressor = nn.Conv1d(8, 4, 1)
-        self.input_proj = nn.Linear(64 * 4, 128)
-        self.res_blocks = nn.Sequential(LiteResidualBlock(128), LiteResidualBlock(128))
-        self.head = nn.Sequential(nn.Linear(128, 128), nn.GELU(), nn.Linear(128, 10))
-
-    def forward(self, x):
-        b = x.size(0)
-        fa = self.lookthemA(self.stream_a(x).view(b, 8, 64).transpose(1, 2))
-        fb = self.lookthemB(self.stream_b(x).view(b, 8, 64).transpose(1, 2))
-        x = self.comb_norm(self.lookthem_comb(torch.cat([fa, fb], dim=2)))
-        x = x.transpose(1, 2)
-        x = self.FFN1(x).transpose(1, 2)
-        res = x
-        x = self.lookthem2(x).transpose(1, 2)
-        x = self.FFN2(x) + res.transpose(1, 2)
-        x = self.compressor(x).flatten(1)
-        x = self.res_blocks(self.input_proj(x))
-        return self.head(x)
+# (Bagian kelas LookThemLayer, LiteResidualBlock, dan LookThemV8MNIST tetap sama)
+# ... (Salin definisi model Anda di sini) ...
 
 # --- LOAD WEIGHTS ON CPU/GPU ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -117,57 +29,39 @@ model.to(device)
 model.eval()
 
 # --- PREPROCESSING MATCHING TRAINING PIPELINE ---
-# Using the exact MNIST normalization values from your training code
 transform_fn = transforms.Compose([
     transforms.Resize((28, 28)),
     transforms.ToTensor(),
     transforms.Normalize((0.1307,), (0.3081,))
 ])
+
 def predict_digit(input_image):
     if input_image is None:
         return "Please draw a number!"
-
+    
     try:
-        # 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
+        # gr.Sketchpad mengembalikan numpy array secara langsung
+        img_array = input_image
 
-        # 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:
-            # Kalau ternyata PIL.Image
-            grayscale = np.array(img_array.convert("L"))
-
-        # 3. Cek kanvas kosong (kali ini cek nilai maks > 0, bukan == 0)
-        if grayscale.max() == 0:
+        # Cek apakah kanvas kosong (semua piksel bernilai 0)
+        if np.max(img_array) == 0:
             return {str(i): 0.1 for i in range(10)}
 
-        # 4. Resize dan normalisasi
-        img = Image.fromarray(grayscale.astype(np.uint8), mode="L")
+        # Konversi ke PIL Image dan resize
+        img = Image.fromarray(img_array.astype(np.uint8), mode="L")
         img = img.resize((28, 28), Image.Resampling.BILINEAR)
+        
+        # Transformasi dan prediksi
         tensor_img = transform_fn(img).unsqueeze(0).to(device)
-
-        # 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:
-        # Untuk debug, kembalikan pesan errornya
-        return {"error": str(e)} 
- 
+        return {"Error": str(e)}
 
 # --- GRADIO INTERFACE CONSTRUCTION ---
 with gr.Blocks() as demo:
@@ -180,12 +74,15 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            input_canvas = gr.Paint(
+            # Gunakan gr.Sketchpad
+            input_canvas = gr.Sketchpad(
                 image_mode="L",
                 height=280,
                 width=280,
-                canvas_color="black",   # ⬅️ ini yang wajib ditambahkan
-                brush=gr.components.image_editor.Brush(default_color="rgb(255, 255, 255)")
+                brush=gr.Brush(
+                    default_color="rgb(255, 255, 255)", # Kuas putih
+                    color_mode="fixed"
+                )
             )
             submit_btn = gr.Button("Classify Digit 🏎️", variant="primary")
             
@@ -194,6 +91,5 @@ with gr.Blocks() as demo:
 
     submit_btn.click(fn=predict_digit, inputs=input_canvas, outputs=output_label)
 
-
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Soft())
+    demo.launch()