import gradio as gr
import base64
import requests
import json
import re
import os
import uuid
from datetime import datetime
# --- Configuration ---
# IMPORTANT: Set your OPENROUTER_API_KEY as a Hugging Face Space Secret
OPENROUTER_API_KEY = "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"
IMAGE_MODEL = "opengvlab/internvl3-14b:free" # Using the free tier model as specified
OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
# --- Global State (managed within Gradio's session if possible, or module-level for simplicity here) ---
# This will be reset each time the processing function is called.
processed_files_data = [] # Stores dicts for each file's details and status
person_profiles = {} # Stores dicts for each identified person and their documents
# --- Helper Functions ---
def extract_json_from_text(text):
if not text:
return {"error": "Empty text provided for JSON extraction."}
match_block = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL | re.IGNORECASE)
if match_block:
json_str = match_block.group(1)
else:
text_stripped = text.strip()
if text_stripped.startswith("`") and text_stripped.endswith("`"):
json_str = text_stripped[1:-1]
else:
json_str = text_stripped
try:
return json.loads(json_str)
except json.JSONDecodeError as e:
try:
first_brace = json_str.find('{')
last_brace = json_str.rfind('}')
if first_brace != -1 and last_brace != -1 and last_brace > first_brace:
potential_json_str = json_str[first_brace : last_brace+1]
return json.loads(potential_json_str)
else:
return {"error": f"Invalid JSON structure: {str(e)}", "original_text": text}
except json.JSONDecodeError as e2:
return {"error": f"Invalid JSON structure after attempting substring: {str(e2)}", "original_text": text}
def get_ocr_prompt():
return f"""You are an advanced OCR and information extraction AI.
Your task is to meticulously analyze this image and extract all relevant information.
Output Format Instructions:
Provide your response as a SINGLE, VALID JSON OBJECT. Do not include any explanatory text before or after the JSON.
The JSON object should have the following top-level keys:
- "document_type_detected": (string) Your best guess of the specific document type (e.g., "Passport", "National ID Card", "Driver's License", "Visa Sticker", "Hotel Confirmation Voucher", "Bank Statement", "Photo of a person").
- "extracted_fields": (object) A key-value map of all extracted information. Be comprehensive. Examples:
- For passports/IDs: "Surname", "Given Names", "Full Name", "Document Number", "Nationality", "Date of Birth", "Sex", "Place of Birth", "Date of Issue", "Date of Expiry", "Issuing Authority", "Country Code".
- For hotel reservations: "Guest Name", "Hotel Name", "Booking Reference", "Check-in Date", "Check-out Date".
- For bank statements: "Account Holder Name", "Account Number", "Bank Name", "Statement Period", "Ending Balance".
- For photos: "Description" (e.g., "Portrait of a person", "Group photo at a location"), "People Present" (array of strings if multiple).
- "mrz_data": (object or null) If a Machine Readable Zone (MRZ) is present:
- "raw_mrz_lines": (array of strings) Each line of the MRZ.
- "parsed_mrz": (object) Key-value pairs of parsed MRZ fields.
If no MRZ, this field should be null.
- "full_text_ocr": (string) Concatenation of all text found on the document.
Extraction Guidelines:
1. Prioritize accuracy.
2. Extract all visible text. Include "Full Name" by combining given and surnames if possible.
3. For dates, try to use ISO 8601 format (YYYY-MM-DD) if possible, but retain original format if conversion is ambiguous.
Ensure the entire output strictly adheres to the JSON format.
"""
def call_openrouter_ocr(image_filepath):
if not OPENROUTER_API_KEY:
return {"error": "OpenRouter API Key not configured."}
try:
with open(image_filepath, "rb") as f:
encoded_image = base64.b64encode(f.read()).decode("utf-8")
mime_type = "image/jpeg"
if image_filepath.lower().endswith(".png"):
mime_type = "image/png"
elif image_filepath.lower().endswith(".webp"):
mime_type = "image/webp"
data_url = f"data:{mime_type};base64,{encoded_image}"
prompt_text = get_ocr_prompt()
payload = {
"model": IMAGE_MODEL,
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": prompt_text},
{"type": "image_url", "image_url": {"url": data_url}}
]
}
],
"max_tokens": 3500,
"temperature": 0.1,
}
headers = {
"Authorization": f"Bearer {OPENROUTER_API_KEY}",
"Content-Type": "application/json",
"HTTP-Referer": "https://huggingface.co/spaces/YOUR_SPACE",
"X-Title": "Gradio Document Processor"
}
response = requests.post(OPENROUTER_API_URL, headers=headers, json=payload, timeout=180)
response.raise_for_status()
result = response.json()
if "choices" in result and result["choices"]:
raw_content = result["choices"][0]["message"]["content"]
return extract_json_from_text(raw_content)
else:
return {"error": "No 'choices' in API response from OpenRouter.", "details": result}
except requests.exceptions.Timeout:
return {"error": "API request timed out."}
except requests.exceptions.RequestException as e:
error_message = f"API Request Error: {str(e)}"
if hasattr(e, 'response') and e.response is not None:
error_message += f" Status: {e.response.status_code}, Response: {e.response.text}"
return {"error": error_message}
except Exception as e:
return {"error": f"An unexpected error occurred during OCR: {str(e)}"}
def extract_entities_from_ocr(ocr_json):
if not ocr_json or "extracted_fields" not in ocr_json or not isinstance(ocr_json.get("extracted_fields"), dict):
doc_type_from_ocr = "Unknown"
if isinstance(ocr_json, dict): # ocr_json itself might be an error dict
doc_type_from_ocr = ocr_json.get("document_type_detected", "Unknown (error in OCR)")
return {"name": None, "dob": None, "passport_no": None, "doc_type": doc_type_from_ocr}
fields = ocr_json["extracted_fields"]
doc_type = ocr_json.get("document_type_detected", "Unknown")
name_keys = ["full name", "name", "account holder name", "guest name"]
dob_keys = ["date of birth", "dob"]
passport_keys = ["document number", "passport number"]
extracted_name = None
for key in name_keys:
for field_key, value in fields.items():
if key == field_key.lower():
extracted_name = str(value) if value else None
break
if extracted_name: break
extracted_dob = None
for key in dob_keys:
for field_key, value in fields.items():
if key == field_key.lower():
extracted_dob = str(value) if value else None
break
if extracted_dob: break
extracted_passport_no = None
for key in passport_keys:
for field_key, value in fields.items():
if key == field_key.lower():
extracted_passport_no = str(value).replace(" ", "").upper() if value else None
break
if extracted_passport_no: break
return {
"name": extracted_name,
"dob": extracted_dob,
"passport_no": extracted_passport_no,
"doc_type": doc_type
}
def normalize_name(name):
if not name: return ""
return "".join(filter(str.isalnum, name)).lower()
def get_person_id_and_update_profiles(doc_id, entities, current_persons_data):
passport_no = entities.get("passport_no")
name = entities.get("name")
dob = entities.get("dob")
if passport_no:
for p_key, p_data in current_persons_data.items():
if passport_no in p_data.get("passport_numbers", set()):
p_data["doc_ids"].add(doc_id)
if name and not p_data.get("canonical_name"): p_data["canonical_name"] = name
if dob and not p_data.get("canonical_dob"): p_data["canonical_dob"] = dob
return p_key
new_person_key = f"person_{passport_no}"
current_persons_data[new_person_key] = {
"canonical_name": name, "canonical_dob": dob,
"names": {normalize_name(name)} if name else set(),
"dobs": {dob} if dob else set(),
"passport_numbers": {passport_no}, "doc_ids": {doc_id},
"display_name": name or f"Person (ID: {passport_no})"
}
return new_person_key
if name and dob:
norm_name = normalize_name(name)
composite_key_nd = f"{norm_name}_{dob}"
for p_key, p_data in current_persons_data.items():
if norm_name in p_data.get("names", set()) and dob in p_data.get("dobs", set()):
p_data["doc_ids"].add(doc_id)
return p_key
new_person_key = f"person_{composite_key_nd}_{str(uuid.uuid4())[:4]}"
current_persons_data[new_person_key] = {
"canonical_name": name, "canonical_dob": dob,
"names": {norm_name}, "dobs": {dob},
"passport_numbers": set(), "doc_ids": {doc_id},
"display_name": name
}
return new_person_key
if name:
norm_name = normalize_name(name)
new_person_key = f"person_{norm_name}_{str(uuid.uuid4())[:4]}"
current_persons_data[new_person_key] = {
"canonical_name": name, "canonical_dob": None,
"names": {norm_name}, "dobs": set(), "passport_numbers": set(),
"doc_ids": {doc_id}, "display_name": name
}
return new_person_key
generic_person_key = f"unidentified_person_{str(uuid.uuid4())[:6]}"
current_persons_data[generic_person_key] = {
"canonical_name": "Unknown", "canonical_dob": None,
"names": set(), "dobs": set(), "passport_numbers": set(),
"doc_ids": {doc_id}, "display_name": f"Unknown Person ({doc_id[:6]})"
}
return generic_person_key
def format_dataframe_data(current_files_data):
df_rows = []
for f_data in current_files_data:
entities = f_data.get("entities") or {} # CORRECTED LINE HERE
df_rows.append([
f_data.get("doc_id", "N/A")[:8],
f_data.get("filename", "N/A"),
f_data.get("status", "N/A"),
entities.get("doc_type", "N/A"),
entities.get("name", "N/A"),
entities.get("dob", "N/A"),
entities.get("passport_no", "N/A"),
f_data.get("assigned_person_key", "N/A")
])
return df_rows
def format_persons_markdown(current_persons_data, current_files_data):
if not current_persons_data:
return "No persons identified yet."
md_parts = ["## Classified Persons & Documents\n"]
for p_key, p_data in current_persons_data.items():
display_name = p_data.get('display_name', p_key)
md_parts.append(f"### Person: {display_name} (Profile Key: {p_key})")
if p_data.get("canonical_dob"): md_parts.append(f"* DOB: {p_data['canonical_dob']}")
if p_data.get("passport_numbers"): md_parts.append(f"* Passport(s): {', '.join(p_data['passport_numbers'])}")
md_parts.append("* Documents:")
doc_ids_for_person = p_data.get("doc_ids", set())
if doc_ids_for_person:
for doc_id in doc_ids_for_person:
doc_detail = next((f for f in current_files_data if f["doc_id"] == doc_id), None)
if doc_detail:
filename = doc_detail.get("filename", "Unknown File")
doc_entities = doc_detail.get("entities") or {}
doc_type = doc_entities.get("doc_type", "Unknown Type")
md_parts.append(f" - {filename} (`{doc_type}`)")
else:
md_parts.append(f" - Document ID: {doc_id[:8]} (details error)")
else:
md_parts.append(" - No documents currently assigned.")
md_parts.append("\n---\n")
return "\n".join(md_parts)
def process_uploaded_files(files_list, progress=gr.Progress(track_tqdm=True)):
global processed_files_data, person_profiles
processed_files_data = []
person_profiles = {}
if not OPENROUTER_API_KEY:
yield (
[["N/A", "ERROR", "OpenRouter API Key not configured.", "N/A", "N/A", "N/A", "N/A", "N/A"]],
"Error: OpenRouter API Key not configured. Please set it in Space Secrets.",
"{}", "API Key Missing. Processing halted."
)
return
if not files_list:
yield ([], "No files uploaded.", "{}", "Upload files to begin.")
return
for i, file_obj in enumerate(files_list):
doc_uid = str(uuid.uuid4())
processed_files_data.append({
"doc_id": doc_uid,
"filename": os.path.basename(file_obj.name if hasattr(file_obj, 'name') else f"file_{i+1}.unknown"),
"filepath": file_obj.name if hasattr(file_obj, 'name') else None, # file_obj itself is filepath if from gr.Files type="filepath"
"status": "Queued",
"ocr_json": None,
"entities": None,
"assigned_person_key": None
})
initial_df_data = format_dataframe_data(processed_files_data)
initial_persons_md = format_persons_markdown(person_profiles, processed_files_data)
yield (initial_df_data, initial_persons_md, "{}", f"Initialized. Found {len(files_list)} files.")
for i, file_data_item in enumerate(progress.tqdm(processed_files_data, desc="Processing Documents")):
current_doc_id = file_data_item["doc_id"]
current_filename = file_data_item["filename"]
if not file_data_item["filepath"]: # Check if filepath is valid
file_data_item["status"] = "Error: Invalid file path"
df_data = format_dataframe_data(processed_files_data)
persons_md = format_persons_markdown(person_profiles, processed_files_data)
yield(df_data, persons_md, "{}", f"({i+1}/{len(processed_files_data)}) Error with file {current_filename}")
continue
file_data_item["status"] = "OCR in Progress..."
df_data = format_dataframe_data(processed_files_data)
persons_md = format_persons_markdown(person_profiles, processed_files_data)
yield (df_data, persons_md, "{}", f"({i+1}/{len(processed_files_data)}) OCR for: {current_filename}")
ocr_result = call_openrouter_ocr(file_data_item["filepath"])
file_data_item["ocr_json"] = ocr_result
if "error" in ocr_result:
file_data_item["status"] = f"OCR Error: {str(ocr_result['error'])[:50]}..."
df_data = format_dataframe_data(processed_files_data)
yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"({i+1}/{len(processed_files_data)}) OCR Error on {current_filename}")
continue
file_data_item["status"] = "OCR Done. Extracting Entities..."
df_data = format_dataframe_data(processed_files_data)
yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"({i+1}/{len(processed_files_data)}) OCR Done for {current_filename}")
entities = extract_entities_from_ocr(ocr_result)
file_data_item["entities"] = entities
file_data_item["status"] = "Entities Extracted. Classifying..."
df_data = format_dataframe_data(processed_files_data)
yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"({i+1}/{len(processed_files_data)}) Entities for {current_filename}")
person_key = get_person_id_and_update_profiles(current_doc_id, entities, person_profiles)
file_data_item["assigned_person_key"] = person_key
file_data_item["status"] = "Classified"
df_data = format_dataframe_data(processed_files_data)
persons_md = format_persons_markdown(person_profiles, processed_files_data)
yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"({i+1}/{len(processed_files_data)}) Classified {current_filename} -> {person_key}")
final_df_data = format_dataframe_data(processed_files_data)
final_persons_md = format_persons_markdown(person_profiles, processed_files_data)
yield (final_df_data, final_persons_md, "{}", f"All {len(processed_files_data)} documents processed.")
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("# 📄 Intelligent Document Processor & Classifier")
gr.Markdown(
"**Upload multiple documents (images of passports, bank statements, hotel reservations, photos, etc.). "
"The system will perform OCR, attempt to extract key entities, and classify documents by the person they belong to.**\n"
"Ensure `OPENROUTER_API_KEY` is set as a Secret in your Hugging Face Space."
)
if not OPENROUTER_API_KEY:
gr.Markdown("<h3 style='color:red;'>⚠️ ERROR: `OPENROUTER_API_KEY` is not set in Space Secrets! OCR will fail.</h3>")
with gr.Row():
with gr.Column(scale=1):
files_input = gr.Files(label="Upload Document Images (Bulk)", file_count="multiple", type="filepath") # Using filepath
process_button = gr.Button("🚀 Process Uploaded Documents", variant="primary")
overall_status_textbox = gr.Textbox(label="Overall Progress", interactive=False, lines=1)
gr.Markdown("---")
gr.Markdown("## Document Processing Details")
dataframe_headers = ["Doc ID (short)", "Filename", "Status", "Detected Type", "Name", "DOB", "Passport No.", "Assigned Person Key"]
document_status_df = gr.Dataframe(
headers=dataframe_headers,
datatype=["str"] * len(dataframe_headers),
label="Individual Document Status & Extracted Entities",
row_count=(1, "dynamic"), # Start with 1 row, dynamically grows
col_count=(len(dataframe_headers), "fixed"),
wrap=True
)
ocr_json_output = gr.Code(label="Selected Document OCR JSON", language="json", interactive=False)
gr.Markdown("---")
person_classification_output_md = gr.Markdown("## Classified Persons & Documents\nNo persons identified yet.")
process_button.click(
fn=process_uploaded_files,
inputs=[files_input],
outputs=[
document_status_df,
person_classification_output_md,
ocr_json_output,
overall_status_textbox
]
)
@document_status_df.select(inputs=None, outputs=ocr_json_output, show_progress="hidden")
def display_selected_ocr(evt: gr.SelectData):
if evt.index is None or evt.index[0] is None:
return "{}"
selected_row_index = evt.index[0]
# Ensure processed_files_data is accessible here. If it's truly global, it should be.
# For safety, one might pass it or make it part of a class if this were more complex.
if 0 <= selected_row_index < len(processed_files_data):
selected_doc_data = processed_files_data[selected_row_index]
if selected_doc_data and selected_doc_data.get("ocr_json"):
# Check if ocr_json is already a dict, if not, try to parse (though it should be)
ocr_data_to_display = selected_doc_data["ocr_json"]
if isinstance(ocr_data_to_display, str): # Should not happen if stored correctly
try:
ocr_data_to_display = json.loads(ocr_data_to_display)
except json.JSONDecodeError:
return json.dumps({"error": "Stored OCR data is not valid JSON string."}, indent=2)
return json.dumps(ocr_data_to_display, indent=2, ensure_ascii=False)
return json.dumps({ "message": "No OCR data found for selected row or selection out of bounds (check if processing is complete). Current rows: " + str(len(processed_files_data))}, indent=2)
if __name__ == "__main__":
demo.queue().launch(debug=True, share=os.environ.get("GRADIO_SHARE", "true").lower() == "true")