import os import torch import torch.nn as nn from PIL import Image import torchvision.transforms as transforms import numpy as np import gradio as gr # ── Professional "Cyber-Dark" CSS ───────────────────────────── custom_css = """ .gradio-container {background-color: #050505;} .feedback-card {border: 1px solid #6366f1; padding: 20px; border-radius: 12px; background: #0f172a; margin: 10px 0;} #header-text {text-align: center; background: linear-gradient(to right, #818cf8, #c084fc); -webkit-background-clip: text; -webkit-text-fill-color: transparent;} button.primary {background: linear-gradient(90deg, #6366f1, #a855f7) !important; color: white !important; font-weight: bold !important; border: none !important;} .tabs {border: none !important;} footer {visibility: hidden} """ # [Model classes: PatchifyAndMask, TransformerBlock, ViT_Encoder, ViT_Decoder remain identical] class PatchifyAndMask(nn.Module): def __init__(self, img_size=224, patch_size=16, in_channels=3): super().__init__() self.img_size = img_size self.patch_size = patch_size self.num_patches = (img_size // patch_size) ** 2 self.patch_dim = in_channels * patch_size * patch_size def patchify(self, imgs): p = self.patch_size h = w = self.img_size // p x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p)) x = torch.einsum('nchpwq->nhwpqc', x) x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3)) return x def random_masking(self, x, mask_ratio=0.75): N, L, D = x.shape len_keep = int(L * (1 - mask_ratio)) noise = torch.rand(N, L, device=x.device) ids_shuffle = torch.argsort(noise, dim=1) ids_restore = torch.argsort(ids_shuffle, dim=1) ids_keep = ids_shuffle[:, :len_keep] x_kept = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D)) mask = torch.ones([N, L], device=x.device) mask[:, :len_keep] = 0 mask = torch.gather(mask, dim=1, index=ids_restore) return x_kept, mask, ids_restore class TransformerBlock(nn.Module): def __init__(self, embed_dim, num_heads): super().__init__() self.norm1 = nn.LayerNorm(embed_dim) self.norm2 = nn.LayerNorm(embed_dim) self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) self.mlp = nn.Sequential( nn.Linear(embed_dim, embed_dim * 4), nn.GELU(), nn.Linear(embed_dim * 4, embed_dim) ) def forward(self, x): attn_output, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x)) x = x + attn_output x = x + self.mlp(self.norm2(x)) return x class ViT_Encoder(nn.Module): def __init__(self, patch_dim=768, embed_dim=768, depth=12, num_heads=12, num_patches=196): super().__init__() self.patch_embed = nn.Linear(patch_dim, embed_dim) self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) self.blocks = nn.ModuleList([ TransformerBlock(embed_dim, num_heads) for _ in range(depth) ]) self.norm = nn.LayerNorm(embed_dim) def forward(self, x, patcher_module): x = self.patch_embed(x) x = x + self.pos_embed x_visible, mask, ids_restore = patcher_module.random_masking(x, mask_ratio=0.75) for block in self.blocks: x_visible = block(x_visible) x_visible = self.norm(x_visible) return x_visible, mask, ids_restore class ViT_Decoder(nn.Module): def __init__(self, encoder_dim=768, decoder_dim=384, depth=12, num_heads=6, num_patches=196, patch_dim=768): super().__init__() self.decoder_embed = nn.Linear(encoder_dim, decoder_dim) self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim)) self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches, decoder_dim)) self.blocks = nn.ModuleList([ TransformerBlock(embed_dim=decoder_dim, num_heads=num_heads) for _ in range(depth) ]) self.norm = nn.LayerNorm(decoder_dim) self.pred = nn.Linear(decoder_dim, patch_dim) def forward(self, x, ids_restore): x = self.decoder_embed(x) B = x.shape[0] num_visible = x.shape[1] num_total = ids_restore.shape[1] num_masks = num_total - num_visible mask_tokens = self.mask_token.repeat(B, num_masks, 1) x_full = torch.cat([x, mask_tokens], dim=1) ids_restore_expanded = ids_restore.unsqueeze(-1).repeat(1, 1, x_full.shape[2]) x_unshuffled = torch.gather(x_full, dim=1, index=ids_restore_expanded) x_unshuffled = x_unshuffled + self.decoder_pos_embed for block in self.blocks: x_unshuffled = block(x_unshuffled) x_unshuffled = self.norm(x_unshuffled) predictions = self.pred(x_unshuffled) return predictions def unpatchify(x, patch_size=16, img_size=224): p = patch_size h = w = img_size // p x = x.reshape(shape=(x.shape[0], h, w, p, p, 3)) x = torch.einsum('nhwpqc->nchpwq', x) imgs = x.reshape(shape=(x.shape[0], 3, h * p, w * p)) return imgs class MaskedAutoencoderApp(nn.Module): def __init__(self): super().__init__() self.patcher = PatchifyAndMask(img_size=224, patch_size=16, in_channels=3) self.encoder = ViT_Encoder(patch_dim=768, embed_dim=768, depth=12, num_heads=12) self.decoder = ViT_Decoder(encoder_dim=768, decoder_dim=384, depth=12, num_heads=6) def forward(self, imgs, mask_ratio): target_patches = self.patcher.patchify(imgs) x = self.encoder.patch_embed(target_patches) x = x + self.encoder.pos_embed x_visible, mask, ids_restore = self.patcher.random_masking(x, mask_ratio=mask_ratio) for block in self.encoder.blocks: x_visible = block(x_visible) x_visible = self.encoder.norm(x_visible) predictions = self.decoder(x_visible, ids_restore) return predictions, target_patches, mask # ── Load Model ──────────────────────────────────────────────── model = MaskedAutoencoderApp() try: model.load_state_dict(torch.load("best_mae_weights.pth", map_location=torch.device("cpu"))) except Exception as e: print(f"Error loading weights: {e}") model.eval() # ── Inference ───────────────────────────────────────────────── def process_image(input_image, mask_ratio_percent): if input_image is None: return None, None, None mask_ratio = mask_ratio_percent / 100.0 image = input_image.convert('RGB') transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()]) input_tensor = transform(image).unsqueeze(0) with torch.no_grad(): preds, targets, mask = model(input_tensor, mask_ratio=mask_ratio) mask_expanded = mask.unsqueeze(-1).repeat(1, 1, targets.shape[2]) masked_input = targets * (1 - mask_expanded) masked_img_np = unpatchify(masked_input).squeeze().numpy() final_recon = preds * mask_expanded + targets * (1 - mask_expanded) recon_img_np = unpatchify(final_recon).squeeze().numpy() orig_img_np = unpatchify(targets).squeeze().numpy() return (np.clip(np.transpose(masked_img_np, (1, 2, 0)), 0, 1), np.clip(np.transpose(recon_img_np, (1, 2, 0)), 0, 1), np.clip(np.transpose(orig_img_np, (1, 2, 0)), 0, 1)) # ── Enhanced UI Construction ────────────────────────────────── with gr.Blocks(css=custom_css, theme=gr.themes.Monochrome()) as demo: with gr.Column(): gr.Markdown("# 💎 PIXEL REVIVE: AI RECONSTRUCTION", elem_id="header-text") gr.Markdown("Self-supervised Masked Autoencoder (MAE) designed to reconstruct images by learning global context from visible patches.") with gr.Tabs(elem_classes="tabs"): with gr.TabItem("🚀 Reconstruction Studio"): with gr.Row(): with gr.Column(scale=1, variant="panel"): gr.Markdown("### Processing Settings") img_input = gr.Image(type="pil", label="Source Image") mask_slider = gr.Slider(10, 90, value=75, step=5, label="Masking Intensity (%)") submit_btn = gr.Button("✨ RECONSTRUCT", variant="primary") gr.Markdown("*(MAE models are often pre-trained with 75% masking)*") with gr.Column(scale=2): with gr.Row(): out_masked = gr.Image(label="Masked Input") out_recon = gr.Image(label="AI Output") out_orig = gr.Image(label="Reference Image", height=200) with gr.TabItem("📊 Technical Architecture"): with gr.Row(): with gr.Column(elem_classes="feedback-card"): gr.Markdown("### Vision Transformer (ViT) Encoder") gr.Markdown("Only processes visible patches. By masking 75% of the image, the encoder is forced to learn highly descriptive structural representations.") with gr.Column(elem_classes="feedback-card"): gr.Markdown("### Lightweight Decoder") gr.Markdown("Reconstructs the original pixels by combining encoded visible patches with learnable 'mask tokens' to fill in the gaps.") submit_btn.click(process_image, [img_input, mask_slider], [out_masked, out_recon, out_orig]) if __name__ == "__main__": demo.launch()