Update app.py

app.py

@@ -17,96 +17,8 @@ if not os.path.exists(MODEL_PATH):
     urllib.request.urlretrieve(HF_URL, MODEL_PATH)
     print("Download complete!")
 
-# --- DEFINE YOUR MODEL ARCHITECTURE (sama seperti sebelumnya) ---
-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)
+# --- DEFINE YOUR MODEL ARCHITECTURE (TETAP SAMA) ---
+# ... (Salin definisi kelas LookThemLayer, LiteResidualBlock, dan LookThemV8MNIST Anda di sini) ...
 
 # --- LOAD WEIGHTS ON CPU/GPU ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -123,40 +35,28 @@ transform_fn = transforms.Compose([
 ])
 
 def predict_digit(input_image):
-    # Selalu kembalikan dictionary 10 digit untuk gr.Label
+    # Default output jika kanvas kosong
     default_output = {str(i): 0.1 for i in range(10)}
     
     if input_image is None:
         return default_output
     
     try:
-        # 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)
+        # gr.Image(source="canvas") mengembalikan numpy array secara langsung
+        img_array = input_image
         
-        # Jika gambar berwarna, ambil channel yang tepat
-        if img_array.ndim == 3:
-            if img_array.shape[-1] == 4:      # RGBA → alpha
+        # Konversi ke grayscale jika perlu (hasil kanvas biasanya sudah grayscale)
+        if isinstance(img_array, np.ndarray) and img_array.ndim == 3:
+            # Ambil channel pertama jika multichannel, atau konversi ke luminance
+            if img_array.shape[-1] == 4:      # RGBA -> alpha
                 grayscale = img_array[..., 3]
-            else:                             # RGB → luminance
+            else:                             # RGB -> luminance
                 grayscale = np.dot(img_array[..., :3], [0.2989, 0.5870, 0.1140])
         else:
             grayscale = img_array
-        
-        # Cek kanvas kosong
-        if grayscale.max() == 0:
+            
+        # Cek apakah kanvas kosong (semua piksel bernilai 0 atau mendekati)
+        if np.max(grayscale) < 5:
             return default_output
         
         # Resize & normalisasi
@@ -166,12 +66,11 @@ def predict_digit(input_image):
         
         with torch.no_grad():
             outputs = model(tensor_img)
-            probs = F.softmax(outputs, dim=1)[0]
-        
-        return {str(i): float(probs[i]) for i in range(10)}
+            probabilities = F.softmax(outputs, dim=1)[0]
+            
+        return {str(i): float(probabilities[i]) for i in range(10)}
     
     except Exception as e:
-        # Kembalikan uniform jika terjadi error tak terduga
         print(f"Prediction error: {e}")
         return default_output
 
@@ -186,12 +85,14 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            # GANTI: gunakan Sketchpad agar latar hitam + pena putih
-            input_canvas = gr.Sketchpad(
+            # Gunakan gr.Image dengan source="canvas"
+            input_canvas = gr.Image(
                 image_mode="L",
                 height=280,
                 width=280,
-                brush=gr.Brush(default_color="rgb(255,255,255)", color_mode="fixed")
+                source="canvas",            # Mengaktifkan mode kanvas untuk menggambar
+                invert_colors=True,         # Membalik warna: latar hitam, coretan putih
+                brush=gr.Brush(default_color="rgb(0,0,0)", color_mode="fixed") # Kuas hitam (akan dibalik jadi putih)
             )
             submit_btn = gr.Button("Classify Digit 🏎️", variant="primary")