Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,123 @@
+
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
+import json
+import string
+
+# --- Recreate Architecture for Inference ---
+# Must match the training notebook architecture
+MAX_SEQ_LEN = 1500 
+
+class CSMTokenizer:
+    def __init__(self):
+        self.chars = list(string.printable) + [chr(i) for i in range(0x0600, 0x06FF + 1)]
+        self.PAD, self.SOS, self.EOS, self.UNK = 0, 1, 2, 3 
+        self.vocab = {c: i+4 for i, c in enumerate(self.chars)}
+        self.inverse_vocab = {i+4: c for i, c in enumerate(self.chars)}
+        self.vocab_size = len(self.vocab) + 4
+        
+    def decode(self, tokens):
+        return "".join([self.inverse_vocab.get(t, "") for t in tokens if t not in [self.PAD, self.SOS, self.EOS]])
+
+class CSMVisionEncoder(nn.Module):
+    def __init__(self, embed_dim=256):
+        super().__init__()
+        self.cnn = nn.Sequential(
+            nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1), nn.ReLU(), nn.BatchNorm2d(32),
+            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), nn.ReLU(), nn.BatchNorm2d(64),
+            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), nn.ReLU(), nn.BatchNorm2d(128),
+            nn.Conv2d(128, embed_dim, kernel_size=3, stride=2, padding=1), nn.ReLU(), nn.BatchNorm2d(embed_dim)
+        )
+        self.pos_embed = nn.Parameter(torch.randn(1, 256, embed_dim))
+
+    def forward(self, x):
+        features = self.cnn(x).flatten(2).permute(0, 2, 1) 
+        return features + self.pos_embed[:, :features.size(1), :]
+
+class CSMJSONDecoder(nn.Module):
+    def __init__(self, vocab_size, embed_dim=256, num_heads=8, num_layers=4):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.pos_encoder = nn.Parameter(torch.randn(1, MAX_SEQ_LEN, embed_dim))
+        decoder_layer = nn.TransformerDecoderLayer(d_model=embed_dim, nhead=num_heads, batch_first=True)
+        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
+        self.fc_out = nn.Linear(embed_dim, vocab_size)
+
+    def forward(self, tgt, memory):
+        tgt_embed = self.embedding(tgt) + self.pos_encoder[:, :tgt.size(1), :]
+        return self.fc_out(self.transformer(tgt_embed, memory))
+
+class CSM_KIE_Universal(nn.Module):
+    def __init__(self, vocab_size):
+        super().__init__()
+        self.encoder = CSMVisionEncoder()
+        self.decoder = CSMJSONDecoder(vocab_size)
+
+# --- Initialization ---
+tokenizer = CSMTokenizer()
+device = torch.device("cpu")
+
+# Load Quantized Model
+print("Downloading trained model...")
+model_path = hf_hub_download(repo_id="Chhagan005/CSM-KIE-Universal", filename="csm_kie_model.pth")
+model = CSM_KIE_Universal(tokenizer.vocab_size)
+model = torch.quantization.quantize_dynamic(model, {nn.Linear, nn.Conv2d}, dtype=torch.qint8)
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
+
+image_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# --- Inference Function ---
+def process_id_card(front_img, back_img):
+    if front_img is None:
+        return '{"error": "Please upload at least the Front side of the ID card."}'
+        
+    # Process Image
+    img_tensor = image_transform(front_img.convert('RGB')).unsqueeze(0)
+    
+    # Autoregressive Generation Logic
+    generated_tokens = [tokenizer.SOS]
+    memory = model.encoder(img_tensor)
+    
+    with torch.no_grad():
+        for _ in range(1000): # Max length
+            tgt_tensor = torch.tensor([generated_tokens], dtype=torch.long)
+            logits = model.decoder(tgt_tensor, memory)
+            next_token = logits[0, -1, :].argmax().item()
+            generated_tokens.append(next_token)
+            if next_token == tokenizer.EOS:
+                break
+                
+    json_string = tokenizer.decode(generated_tokens)
+    
+    # Format and return JSON
+    try:
+        parsed_json = json.loads(json_string)
+        return json.dumps(parsed_json, indent=2, ensure_ascii=False)
+    except:
+        return json_string # Fallback if model generates slight syntax error during early stages
+
+# --- Gradio UI ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🪪 CSM-KIE Universal ID Scanner")
+    gr.Markdown("Upload Front and Back sides of any International ID card (Middle East, Africa, etc.) to extract multilingual structured JSON data using the proprietary CSM-DocVL model.")
+    
+    with gr.Row():
+        with gr.Column():
+            front = gr.Image(type="pil", label="Front Side (Required)")
+            back = gr.Image(type="pil", label="Back Side / MRZ (Optional)")
+            scan_btn = gr.Button("🔍 Scan & Extract JSON", variant="primary")
+        
+        with gr.Column():
+            output_json = gr.Code(language="json", label="Structured JSON Output")
+            
+    scan_btn.click(process_id_card, inputs=[front, back], outputs=output_json)
+
+demo.launch()