Masked Image Modeling

"""Masked image modeling (MIM) mini ala MAE pada dataset digits.

Setiap gambar digit 8x8 dipecah menjadi 16 patch 2x2. Saat pelatihan,
8 dari 16 patch (50%) disembunyikan secara acak. Encoder mini hanya
memproses patch yang terlihat; slot patch tersembunyi diisi satu token
topeng yang dipelajari. Decoder kecil lalu merekonstruksi seluruh
64 piksel, tetapi loss hanya dihitung pada patch yang disembunyikan.
Pembanding: menebak piksel topeng dengan rata-rata dataset.
Evaluasi ulang pada rasio topeng 75% (angka yang dipakai MAE asli).
Diuji di CPU dengan PyTorch 2.14, seed tetap, deterministik.
"""

import numpy as np
import torch
import torch.nn as nn
from sklearn.datasets import load_digits

torch.manual_seed(0)

# ---------- 1. Data: gambar digits 8x8 dinormalkan ke [0, 1] ----------
digits = load_digits()
X = digits.data.astype(np.float32) / 16.0            # (1797, 64)
acak = np.random.default_rng(42).permutation(len(X))
X = X[acak]
target = digits.target[acak]                      # hanya untuk memilih contoh
X_latih = torch.tensor(X[:1500])
X_uji = torch.tensor(X[1500:])
target_uji = target[1500:]
SISI, PATCH, N_PATCH = 8, 2, 16                      # 16 patch 2x2 per gambar
N_VIS, N_RUSAK = 8, 8                                # 8 terlihat, 8 disembunyikan
print(f"data     : {X_latih.shape[0]} gambar latih, {X_uji.shape[0]} gambar uji")
print(f"patch    : {N_PATCH} patch {PATCH}x{PATCH} piksel per gambar {SISI}x{SISI}")
print(f"rasio    : {N_RUSAK}/{N_PATCH} = {N_RUSAK / N_PATCH:.0%} patch disembunyikan")
print(f"baseline : tebak piksel topeng dengan rata-rata dataset "
      f"({X_latih.mean():.3f})")
print()


def patchify(img):
    """(B, 64) -> (B, 16, 4): belah gambar jadi grid patch 2x2."""
    b = img.reshape(-1, SISI, SISI)
    p = b.reshape(-1, SISI // PATCH, PATCH, SISI // PATCH, PATCH)
    return p.permute(0, 1, 3, 2, 4).reshape(-1, N_PATCH, PATCH * PATCH)


def pilih_topeng(B, n_rusak, generator=None):
    """Pilih n_rusak indeks patch acak per gambar, tanpa duplikat."""
    if generator is not None:
        r = torch.rand(B, N_PATCH, generator=generator)
    else:
        r = torch.rand(B, N_PATCH)
    idx = r.topk(n_rusak, dim=1).indices             # n_rusak indeks acak
    topeng = torch.zeros(B, N_PATCH, dtype=torch.bool)
    return topeng.scatter(1, idx, True)              # True = disembunyikan


class MAEMini(nn.Module):
    """Encoder per-patch + token topeng + decoder rekonstruksi."""

    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(4, 32), nn.ReLU())
        self.token_topeng = nn.Parameter(torch.zeros(32))
        self.pos = nn.Parameter(torch.zeros(N_PATCH, 8))
        self.decoder = nn.Sequential(
            nn.Linear(40 * N_PATCH, 128), nn.ReLU(), nn.Linear(128, 64)
        )

    def forward(self, img, topeng):
        B = img.shape[0]
        pat = patchify(img)                           # (B, 16, 4)
        # urutkan indeks: patch terlihat (False) duluan -> ambil N_VIS
        idx_vis = torch.argsort(topeng.int(), dim=1)[:, :N_VIS]
        pat_vis = pat.gather(
            1, idx_vis[..., None].expand(-1, -1, PATCH * PATCH))
        z_vis = self.encoder(pat_vis)                 # encode patch terlihat
        z = self.token_topeng.expand(B, N_PATCH, 32).clone()
        z.scatter_(1, idx_vis[..., None].expand(-1, -1, 32), z_vis)
        fitur = torch.cat([z, self.pos.expand(B, -1, -1)], dim=2)
        return self.decoder(fitur.flatten(1))         # (B, 64) rekonstruksi


model = MAEMini()
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
mse = nn.MSELoss()

# ---------- 2. Pelatihan: rekonstruksi patch yang disembunyikan ----------
print("=== Pelatihan MAE-mini (loss pada patch tersembunyi saja) ===")
for ep in range(300):
    i = torch.randint(0, X_latih.shape[0], (128,))
    img = X_latih[i]
    topeng = pilih_topeng(img.shape[0], N_RUSAK)
    pred = model(img, topeng)
    target_patch = patchify(img)
    pred_p = patchify(pred)
    loss = mse(pred_p[topeng], target_patch[topeng])
    opt.zero_grad(); loss.backward(); opt.step()
    if ep in (0, 299):
        print(f"  epoch {ep + 1:3d} : loss {loss.item():.4f}")
print()


# ---------- 3. Evaluasi: galat rekonstruksi vs baseline ----------
def galat(n_rusak, seed):
    """MSE pada piksel topeng uji + MSE baseline rata-rata dataset."""
    g = torch.Generator().manual_seed(seed)
    total_model, total_base, n = 0.0, 0.0, 0
    with torch.no_grad():
        for i in range(0, X_uji.shape[0], 100):
            img = X_uji[i:i + 100]
            topeng = pilih_topeng(img.shape[0], n_rusak, g)
            pred = patchify(model(img, topeng))
            asli = patchify(img)
            total_model += mse(pred[topeng], asli[topeng]).item() * img.shape[0]
            rata = X_latih.mean().expand_as(asli)
            total_base += mse(rata[topeng], asli[topeng]).item() * img.shape[0]
            n += img.shape[0]
    return total_model / n, total_base / n


print("=== Galat rekonstruksi pada data uji (piksel topeng) ===")
for n_rusak, nama in ((8, "topeng 50% (8/16 patch)"), (12, "topeng 75% (12/16)")):
    m, b = galat(n_rusak, seed=1234)
    print(f"  {nama:24s} : model {m:.4f} vs baseline {b:.4f} "
          f"({b / m:.1f}x lebih baik)")
print()


# ---------- 4. Visualisasi satu gambar: asli, masukan, rekonstruksi ----------
KARAKTER = " .:-=+*#%@"


def gambar_ascii(piksel, topeng=None, judul=""):
    """Cetak gambar 8x8 sebagai karakter.

    Bila topeng diberikan: piksel pada patch tersembunyi ditandai '?'
    sehingga pembaca melihat patch mana saja yang dibaca model.
    """
    baris = []
    g = piksel.reshape(SISI, SISI)
    t = topeng.reshape(-1).tolist() if topeng is not None else None
    for r in range(SISI):
        sel = []
        for c in range(SISI):
            if t is not None and t[(r // PATCH) * PATCH + c // PATCH]:
                sel.append("?")
            else:
                sel.append(KARAKTER[min(int(g[r, c] * 9.99), 9)])
        baris.append("".join(sel))
    return f"{judul}\n" + "\n".join(baris)


print("=== Contoh rekonstruksi digit 0 (tanda ? = patch disembunyikan) ===")
i = int(np.flatnonzero(target_uji == 0)[0])   # digit 0: bentuk cincin
img = X_uji[i:i + 1]
g = torch.Generator().manual_seed(7)
topeng = pilih_topeng(1, N_RUSAK, g)
with torch.no_grad():
    rekon = model(img, topeng)[0]
print(gambar_ascii(img[0], None, "gambar asli         :"))
print()
print(gambar_ascii(img[0], topeng, "masukan model       :"))
print()
print(gambar_ascii(rekon, None, "rekonstruksi penuh  :"))
with torch.no_grad():
    mse_satu = mse(patchify(rekon)[topeng], patchify(img[0])[topeng]).item()
print(f"MSE piksel topeng gambar ini : {mse_satu:.4f}")

data     : 1500 gambar latih, 297 gambar uji
patch    : 16 patch 2x2 piksel per gambar 8x8
rasio    : 8/16 = 50% patch disembunyikan
baseline : tebak piksel topeng dengan rata-rata dataset (0.305)

=== Pelatihan MAE-mini (loss pada patch tersembunyi saja) ===
  epoch   1 : loss 0.2216
  epoch 300 : loss 0.0361

=== Galat rekonstruksi pada data uji (piksel topeng) ===
  topeng 50% (8/16 patch)  : model 0.0360 vs baseline 0.1436 (4.0x lebih baik)
  topeng 75% (12/16)       : model 0.0353 vs baseline 0.1416 (4.0x lebih baik)

=== Contoh rekonstruksi digit 0 (tanda ? = patch disembunyikan) ===
gambar asli         :
  *@%.  
 .@+=%  
 -@..@- 
 -@. #* 
 =%  #* 
 =@ .@- 
 :@=#*  
  *@%.  

masukan model       :
??*@????
??@+????
????.@- 
???? #* 
 =%  #??
 =@ .@??
 :??#*??
  ??%.??

rekonstruksi penuh  :
  +%#.  
 .@*+*  
 -%..*. 
 -#  =- 
 -*  =- 
 :# .#- 
 .@*+%. 
  +%#.  
MSE piksel topeng gambar ini : 0.0196