Update app.py

app.py

@@ -1,36 +1,38 @@
-# app.py
-import gradio as gr
+# final_app.py
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import torchvision.transforms as T
+from transformers import AutoTokenizer, AutoModel
+from torchvision import transforms
 from PIL import Image
+import gradio as gr
 
-# --- Sacred Geometry Symbolic Classifier Core ---
+# --- SriYantra Custom Layer ---
 class SriYantraLayer(nn.Module):
-    def __init__(self, in_features, out_features, num_heads=3):
-        super(SriYantraLayer, self).__init__()
+    def __init__(self, in_dim, out_dim, num_heads):
+        super().__init__()
         self.triangle_heads = nn.ModuleList([
-            nn.Linear(in_features, out_features) for _ in range(num_heads)
+            nn.Linear(in_dim, out_dim) for _ in range(num_heads)
         ])
-        self.norm = nn.LayerNorm(out_features)
+        self.norm = nn.LayerNorm(out_dim)
 
     def forward(self, x):
-        head_outputs = [F.relu(head(x)) for head in self.triangle_heads]
-        merged = sum(head_outputs) / len(head_outputs)
-        return self.norm(merged)
+        outputs = [F.relu(head(x)) for head in self.triangle_heads]
+        combined = sum(outputs) / len(outputs)
+        return self.norm(combined)
 
+# --- Full Model ---
 class SriYantraNet(nn.Module):
-    def __init__(self, input_dim, hidden_dim, latent_dim, output_dim):
-        super(SriYantraNet, self).__init__()
-        self.outer1 = SriYantraLayer(input_dim, hidden_dim, num_heads=4)
-        self.outer2 = SriYantraLayer(hidden_dim, hidden_dim, num_heads=4)
-        self.inner1 = SriYantraLayer(hidden_dim, hidden_dim, num_heads=3)
-        self.inner2 = SriYantraLayer(hidden_dim, latent_dim, num_heads=3)
-        self.center = nn.Linear(latent_dim, latent_dim)
-        self.decoder1 = SriYantraLayer(latent_dim, hidden_dim, num_heads=3)
-        self.decoder2 = SriYantraLayer(hidden_dim, hidden_dim, num_heads=4)
-        self.final = nn.Linear(hidden_dim, output_dim)
+    def __init__(self):
+        super().__init__()
+        self.outer1 = SriYantraLayer(896, 256, 4)
+        self.outer2 = SriYantraLayer(256, 256, 4)
+        self.inner1 = SriYantraLayer(256, 256, 3)
+        self.inner2 = SriYantraLayer(256, 64, 3)
+        self.center = nn.Linear(64, 64)
+        self.decoder1 = SriYantraLayer(64, 256, 3)
+        self.decoder2 = SriYantraLayer(256, 256, 4)
+        self.final = nn.Linear(256, 10)
 
     def forward(self, x):
         x = self.outer1(x)
@@ -42,72 +44,59 @@ class SriYantraNet(nn.Module):
         x = self.decoder2(x)
         return self.final(x)
 
-# --- Image Feature Extractor ---
-class ImageToVector(nn.Module):
-    def __init__(self, out_dim=128):
-        super(ImageToVector, self).__init__()
-        self.encoder = nn.Sequential(
-            nn.Conv2d(1, 16, 3, stride=2),
-            nn.ReLU(),
-            nn.Conv2d(16, 32, 3, stride=2),
-            nn.ReLU(),
-            nn.Conv2d(32, 64, 3, stride=2),
-            nn.ReLU(),
-            nn.Flatten(),
-            nn.Linear(64*6*6, out_dim),
-            nn.ReLU()
-        )
-
-    def forward(self, x):
-        return self.encoder(x)
-
-# --- Complete Classifier ---
-class SacredSymbolClassifier(nn.Module):
-    def __init__(self):
-        super(SacredSymbolClassifier, self).__init__()
-        self.visual = ImageToVector(out_dim=128)
-        self.symbolic = SriYantraNet(128, 256, 64, 10)
-
-    def forward(self, x):
-        v = self.visual(x)
-        return self.symbolic(v)
+# Load tokenizer and text model
+print("Loading IndicBERTv2...")
+tokenizer = AutoTokenizer.from_pretrained("ai4bharat/IndicBERTv2-MLM-only")
+sanskrit_model = AutoModel.from_pretrained("ai4bharat/IndicBERTv2-MLM-only")
 
-# Instantiate model
-model = SacredSymbolClassifier()
+# Load symbol classifier
+model = SriYantraNet()
 model.eval()
 
-# Dummy weights for demo (replace with trained model)
+# Dummy weights (replace with trained weights if available)
 with torch.no_grad():
     for param in model.parameters():
         param.uniform_(-0.1, 0.1)
 
-# --- Preprocessing ---
-transform = T.Compose([
-    T.Grayscale(),
-    T.Resize((64, 64)),
-    T.ToTensor()
+# Image preprocessing
+image_transform = transforms.Compose([
+    transforms.Resize((64, 64)),
+    transforms.Grayscale(),
+    transforms.ToTensor()
 ])
 
 # --- Inference Function ---
-def classify_symbol(img: Image.Image):
+def predict(image, sanskrit_text):
     try:
-        img_tensor = transform(img).unsqueeze(0)  # [1, 1, 64, 64]
+        image = image.convert("RGB")
+        img_tensor = image_transform(image).view(1, -1)[:, :128]  # shape: [1, 128]
+
+        tokens = tokenizer(sanskrit_text, return_tensors="pt", truncation=True, padding=True)
         with torch.no_grad():
-            output = model(img_tensor)
+            text_emb = sanskrit_model(**tokens).last_hidden_state.mean(dim=1)  # shape: [1, 768]
+
+        fused = torch.cat([img_tensor, text_emb], dim=1)  # shape: [1, 896]
+
+        with torch.no_grad():
+            output = model(fused)
             pred = torch.argmax(output, dim=1).item()
+
         return f"🔮 Predicted Symbolic Pattern Class: {pred}"
     except Exception as e:
         return f"❌ Error during prediction: {str(e)}"
 
 # --- Gradio Interface ---
-demo = gr.Interface(
-    fn=classify_symbol,
-    inputs=gr.Image(type="pil", label="Upload Sacred Image (Yantra, Mandala, etc.)"),
-    outputs=gr.Textbox(label="Prediction Result"),
-    title="🛸 Sacred Geometry Symbol Classifier",
-    description="A Sri-Yantra inspired deep learning model to classify symbolic sacred patterns.",
-    theme="compact"
+iface = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil", label="Upload Symbol Image"),
+        gr.Textbox(label="Enter Sanskrit Text")
+    ],
+    outputs=gr.Textbox(label="Prediction"),
+    title="🔺 SriYantra-Net: Symbolic Pattern Classifier",
+    description="Upload a sacred symbol image and Sanskrit phrase to classify symbolic pattern using a fused image-text deep network.",
+    theme="default"
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()