Update app.py

app.py

@@ -0,0 +1,44 @@
+import streamlit as st
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from PIL import Image
+import torch
+import io
+
+# Load model and processor once
+@st.cache_resource
+def load_model():
+    model_id = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"
+    processor = AutoProcessor.from_pretrained(model_id)
+    model = AutoModelForVision2Seq.from_pretrained(model_id).to("cuda" if torch.cuda.is_available() else "cpu")
+    return processor, model
+
+processor, model = load_model()
+
+# Streamlit UI
+st.title("Aadhaar Card Information Extractor")
+uploaded_file = st.file_uploader("Upload Aadhaar card image", type=["jpg", "png", "jpeg"])
+
+if uploaded_file is not None:
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded Aadhaar Card", use_column_width=True)
+
+    if st.button("Extract Info"):
+        with st.spinner("Extracting..."):
+            prompt = (
+                "You are an AI system for extracting information from Indian Aadhaar cards. "
+                "From the image, extract and return a structured JSON with:\n"
+                "- Name\n"
+                "- Father's Name\n"
+                "- Date of Birth\n"
+                "- Gender\n"
+                "- Aadhaar Number\n"
+                "- Address (Street, Locality, District, State, PIN)\n"
+                "- QR code data (if visible)\n"
+                "- Bounding box of photograph as [x1, y1, x2, y2]\n"
+                "Respond only with JSON."
+            )
+
+            inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
+            outputs = model.generate(**inputs, max_new_tokens=512)
+            result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+            st.code(result, language="json")