|
|
import gradio as gr |
|
|
import base64 |
|
|
import requests |
|
|
from PIL import Image |
|
|
import io |
|
|
|
|
|
API_KEY = "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" |
|
|
IMAGE_MODEL = "OpenGVLab/InternVL3-14B" |
|
|
|
|
|
def extract_passport_info(images, document_type): |
|
|
results = [] |
|
|
|
|
|
for image in images: |
|
|
|
|
|
buffered = io.BytesIO() |
|
|
image.save(buffered, format="JPEG") |
|
|
encoded_image = base64.b64encode(buffered.getvalue()).decode("utf-8") |
|
|
data_url = f"data:image/jpeg;base64,{encoded_image}" |
|
|
|
|
|
|
|
|
prompt = ( |
|
|
f"Extract all passport information from the uploaded {document_type} image. " |
|
|
"Include MRZ (if present), full name, passport number, nationality, gender, " |
|
|
"date of birth, date of issue, expiry date, issuing country, and any other text or labels in other languages. " |
|
|
"Return the result in a JSON format." |
|
|
) |
|
|
|
|
|
|
|
|
payload = { |
|
|
"model": IMAGE_MODEL, |
|
|
"messages": [ |
|
|
{ |
|
|
"role": "user", |
|
|
"content": [ |
|
|
{"type": "text", "text": prompt}, |
|
|
{"type": "image_url", "image_url": {"url": data_url}}, |
|
|
], |
|
|
} |
|
|
], |
|
|
} |
|
|
|
|
|
headers = { |
|
|
"Authorization": f"Bearer {API_KEY}", |
|
|
"Content-Type": "application/json" |
|
|
} |
|
|
|
|
|
try: |
|
|
response = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=payload) |
|
|
result = response.json() |
|
|
print("📡 Status:", response.status_code) |
|
|
print("📡 Raw Result:", result) |
|
|
|
|
|
if "choices" in result: |
|
|
extracted = result["choices"][0]["message"]["content"] |
|
|
results.append({ |
|
|
"document_type": document_type, |
|
|
"extracted_info": extracted |
|
|
}) |
|
|
else: |
|
|
results.append({ |
|
|
"document_type": document_type, |
|
|
"extracted_info": "❌ No data extracted" |
|
|
}) |
|
|
|
|
|
except Exception as e: |
|
|
results.append({ |
|
|
"document_type": document_type, |
|
|
"extracted_info": f"⚠️ Error: {str(e)}" |
|
|
}) |
|
|
|
|
|
return results |
|
|
|
|
|
|
|
|
|
|
|
demo = gr.Interface( |
|
|
fn=extract_passport_info, |
|
|
inputs=[ |
|
|
gr.Image(type="pil", label="Upload Passport/Document Images", multiple=True), |
|
|
gr.Dropdown( |
|
|
choices=["passport_front", "passport_back", "photo", "hotel_reservation"], |
|
|
label="Document Type", |
|
|
value="passport_front", |
|
|
) |
|
|
], |
|
|
outputs="json", |
|
|
title="Passport & Document Info Extractor", |
|
|
description="Upload one or more document images. Extracted information will include MRZ and all available text, structured in JSON format.", |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch() |
|
|
|
