import os os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # Suppress TensorFlow INFO and WARNING messages os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0' import gradio as gr import base64 import requests import json import re import os import uuid from datetime import datetime import time # For potential sleeps if needed, or timing # Attempt to import deepface and handle import error gracefully try: from deepface import DeepFace from deepface.commons import functions as deepface_functions DEEPFACE_AVAILABLE = True except ImportError: DEEPFACE_AVAILABLE = False print("Warning: deepface library not found. Facial recognition features will be disabled.") # Mock DeepFace object if not available to prevent NameErrors, though functions won't work class DeepFaceMock: def represent(self, *args, **kwargs): return [] def verify(self, *args, **kwargs): return {'verified': False, 'distance': float('inf')} def detectFace(self, *args, **kwargs): raise NotImplementedError("DeepFace not installed") DeepFace = DeepFaceMock() # --- Configuration --- OPENROUTER_API_KEY = "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" IMAGE_MODEL = "opengvlab/internvl3-14b:free" OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions" # Facial Recognition Configuration FACE_DETECTOR_BACKEND = 'retinaface' # common and effective FACE_RECOGNITION_MODEL_NAME = 'VGG-Face' # good balance # Threshold for deepface.verify (model-specific, VGG-Face with cosine is often around 0.40 for verification) # Lower threshold means stricter match for verify. For similarity search, we might use raw distance. # DeepFace.verify uses model-specific thresholds internally. Let's rely on its 'verified' flag. FACE_SIMILARITY_THRESHOLD = 0.60 # For cosine distance, lower is more similar. For similarity, higher is better. # Deepface verify returns 'distance'. For cosine, lower distance = more similar. # Let's use a distance threshold. For VGG-Face with cosine, this might be < 0.4 for a match. # We will use deepface.verify which handles this internally. # --- Global State --- processed_files_data = [] person_profiles = {} # --- 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 (no outer braces found): {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(): # Enhanced 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 Front", "Passport Back", "National ID Card", "Photo of a person", "Hotel Reservation", "Bank Statement"). - "extracted_fields": (object) A key-value map of all extracted information. Be comprehensive. - For ALL document types, if a primary person is the subject, try to include: "Primary Person Name", "Full Name". - List other names found under specific keys like "Guest Name", "Account Holder Name", "Mother's Name", "Spouse's Name". - Extract critical identifiers like "Passport Number", "Document Number", "ID Number", "Account Number", "Reservation Number" FROM ANY PART OF THE DOCUMENT where they appear. Use consistent key names for these if possible. - For passports/IDs: "Surname", "Given Names", "Nationality", "Date of Birth", "Sex", "Place of Birth", "Date of Issue", "Date of Expiry". - For photos: "Description" (e.g., "Portrait of John Doe", "User's profile photo"), "People Present" (array of names if discernible). - "mrz_data": (object or null) If a Machine Readable Zone (MRZ) is present. - "full_text_ocr": (string) Concatenation of all text found on the document. Extraction Guidelines: 1. Extract "Passport Number" or "Document Number" even from back sides or less prominent areas. 2. Identify and list all prominent names. If one person is clearly the main subject, label their name as "Primary Person Name" or "Full Name". 3. For dates, aim for YYYY-MM-DD. Ensure the entire output strictly adheres to the JSON format. """ def call_openrouter_ocr(image_filepath): # (User's existing function - kept mostly as is, ensure YOUR_SPACE is updated if needed) 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": os.environ.get("GRADIO_ROOT_PATH", "http://localhost:7860"), # Better placeholder "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 get_facial_embeddings_with_deepface(image_filepath): if not DEEPFACE_AVAILABLE: return {"error": "DeepFace library not installed.", "embeddings": []} try: # Use represent to get embeddings. It can find multiple faces. # Setting align=True, detector_backend for robustness. # enforce_detection=False will return empty list if no face, rather than error. embedding_objs = DeepFace.represent( img_path=image_filepath, model_name=FACE_RECOGNITION_MODEL_NAME, detector_backend=FACE_DETECTOR_BACKEND, enforce_detection=False, # Don't raise error if no face align=True ) # DeepFace.represent returns a list of dictionaries, each with an 'embedding' key embeddings = [obj['embedding'] for obj in embedding_objs if 'embedding' in obj] if not embeddings: return {"message": "No face detected or embedding failed.", "embeddings": []} return {"embeddings": embeddings, "count": len(embeddings)} except Exception as e: # Catch errors from DeepFace if enforce_detection was True or other issues # Like "Face detector ... could not find anyıs face" if "could not find any face" in str(e).lower(): return {"message": "No face detected.", "embeddings": []} return {"error": f"Facial embedding extraction failed: {str(e)}", "embeddings": []} def extract_entities_from_ocr(ocr_json): if not ocr_json or not isinstance(ocr_json, dict) 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): doc_type_from_ocr = ocr_json.get("document_type_detected", "Unknown (error in OCR)") return {"name": None, "dob": None, "main_id": None, "doc_type": doc_type_from_ocr, "all_names_roles": []} fields = ocr_json["extracted_fields"] doc_type = ocr_json.get("document_type_detected", "Unknown") # Expanded and prioritized name keys # Order matters: more specific or primary names first name_keys = [ "primary person name", "full name", "name", "account holder name", "guest name", "cardholder name", "policy holder name", "applicant name", "beneficiary name", "student name", "employee name", "sender name", "receiver name", "patient name", "traveler name", "customer name", "member name", "user name" ] dob_keys = ["date of birth", "dob"] # Expanded ID keys (passport, national ID, etc.) id_keys = ["passport number", "document number", "id number", "personal no", "member id", "customer id", "account number", "reservation number"] extracted_name = None all_names_roles = [] # To store all names found with their original JSON key for key in name_keys: for field_key, value in fields.items(): if key == field_key.lower(): if value and isinstance(value, str) and value.strip(): if not extracted_name: # Take the first one found as primary for now extracted_name = value.strip() all_names_roles.append({"name_text": value.strip(), "source_key": field_key}) # If "People Present" exists (e.g., for photos), add them if "people present" in (k.lower() for k in fields.keys()): people = fields.get([k for k in fields if k.lower() == "people present"][0]) if isinstance(people, list): for person_name in people: if isinstance(person_name, str) and person_name.strip(): all_names_roles.append({"name_text": person_name.strip(), "source_key": "People Present"}) if not extracted_name: extracted_name = person_name.strip() # Prioritize if no other name found extracted_dob = None for key in dob_keys: for field_key, value in fields.items(): if key == field_key.lower() and value and isinstance(value, str): extracted_dob = value.strip() break if extracted_dob: break extracted_main_id = None for key in id_keys: for field_key, value in fields.items(): if key == field_key.lower() and value and isinstance(value, str): extracted_main_id = value.replace(" ", "").upper().strip() # Normalize break if extracted_main_id: break return { "name": extracted_name, "dob": extracted_dob, "main_id": extracted_main_id, # This will be used as the primary linking ID "doc_type": doc_type, "all_names_roles": list({tuple(d.items()): d for d in all_names_roles}.values()) # Deduplicate } def normalize_name(name): if not name: return "" return "".join(filter(str.isalnum, name)).lower() def are_faces_similar(emb1_list, emb2_gallery_list): if not DEEPFACE_AVAILABLE or not emb1_list or not emb2_gallery_list: return False # Compare each embedding from emb1_list against each in emb2_gallery_list for emb1 in emb1_list: for emb2 in emb2_gallery_list: try: # DeepFace.verify expects embeddings directly if not paths # It uses built-in thresholds per model. result = DeepFace.verify( img1_path=emb1, # Pass embedding directly img2_path=emb2, # Pass embedding directly model_name=FACE_RECOGNITION_MODEL_NAME, detector_backend=FACE_DETECTOR_BACKEND, # Though not used for verify with embeddings distance_metric='cosine' # Or 'euclidean', 'euclidean_l2' ) if result.get("verified", False): # print(f"Face match found: distance {result.get('distance')}") return True except Exception as e: print(f"DeepFace verify error: {e}") # e.g. if embeddings are not in expected format return False def get_person_id_and_update_profiles(doc_id, entities, facial_embeddings, current_persons_data, linking_method_log): main_id = entities.get("main_id") # Passport No, Document No, Account No etc. name = entities.get("name") dob = entities.get("dob") # Tier 1: Match by Main ID (Passport, National ID, etc.) if main_id: for p_key, p_data in current_persons_data.items(): if main_id in p_data.get("ids", set()): p_data["doc_ids"].add(doc_id) if name and normalize_name(name) not in p_data["names"]: p_data["names"].add(normalize_name(name)) if dob and dob not in p_data["dobs"]: p_data["dobs"].add(dob) if facial_embeddings: p_data["face_gallery"].extend(facial_embeddings) # Add new faces linking_method_log.append(f"Linked by Main ID ({main_id}) to {p_key}") return p_key # New person based on this main_id new_person_key = f"person_id_{main_id}" current_persons_data[new_person_key] = { "display_name": name or f"Person (ID: {main_id})", "names": {normalize_name(name)} if name else set(), "dobs": {dob} if dob else set(), "ids": {main_id}, "face_gallery": list(facial_embeddings or []), # Initialize gallery "doc_ids": {doc_id} } linking_method_log.append(f"New person by Main ID ({main_id}): {new_person_key}") return new_person_key # Tier 2: Match by Facial Recognition if facial_embeddings: for p_key, p_data in current_persons_data.items(): if are_faces_similar(facial_embeddings, p_data.get("face_gallery", [])): p_data["doc_ids"].add(doc_id) if name and normalize_name(name) not in p_data["names"]: p_data["names"].add(normalize_name(name)) if dob and dob not in p_data["dobs"]: p_data["dobs"].add(dob) p_data["face_gallery"].extend(facial_embeddings) # Freshen gallery linking_method_log.append(f"Linked by Facial Match to {p_key}") return p_key # If no facial match to existing, but we have a face and name/dob, it will be used for new profile below # Tier 3: Match by Normalized Name + DOB if name and dob: norm_name = normalize_name(name) 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) if facial_embeddings: p_data["face_gallery"].extend(facial_embeddings) linking_method_log.append(f"Linked by Name+DOB to {p_key}") return p_key # New person based on name and DOB new_person_key = f"person_{norm_name}_{dob}_{str(uuid.uuid4())[:4]}" current_persons_data[new_person_key] = { "display_name": name, "names": {norm_name}, "dobs": {dob}, "ids": set(), "face_gallery": list(facial_embeddings or []), "doc_ids": {doc_id} } linking_method_log.append(f"New person by Name+DOB: {new_person_key}") return new_person_key # Tier 4: Match by Normalized Name only (creates a more tentative profile) if name: norm_name = normalize_name(name) # Check if any existing profile primarily matches this name AND has no stronger identifiers yet (e.g. no DOB, no ID, no face) # This logic could be refined to prevent overly aggressive merging or splitting. # For now, we'll create a new profile if not matched above. new_person_key = f"person_name_{norm_name}_{str(uuid.uuid4())[:4]}" current_persons_data[new_person_key] = { "display_name": name, "names": {norm_name}, "dobs": set(), "ids": set(), "face_gallery": list(facial_embeddings or []), "doc_ids": {doc_id} } linking_method_log.append(f"New person by Name only: {new_person_key}") return new_person_key # Tier 5: Unclassifiable by PII, but might have a face generic_person_key = f"unidentified_person_{str(uuid.uuid4())[:6]}" current_persons_data[generic_person_key] = { "display_name": f"Unknown Person ({doc_id[:6]})", "names": set(), "dobs": set(), "ids": set(), "face_gallery": list(facial_embeddings or []), "doc_ids": {doc_id} } linking_method_log.append(f"New Unidentified Person: {generic_person_key}") 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 {} face_info = f_data.get("face_analysis_result", {}) face_detected_status = "Y" if face_info.get("count", 0) > 0 else "N" if "error" in face_info : face_detected_status = "Error" elif "message" in face_info and "No face detected" in face_info["message"]: face_detected_status = "N" 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"), face_detected_status, entities.get("name", "N/A"), entities.get("dob", "N/A"), entities.get("main_id", "N/A"), # Changed from passport_no to main_id f_data.get("assigned_person_key", "N/A"), f_data.get("linking_method", "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 sorted(current_persons_data.items()): # Sort for consistent display display_name = p_data.get('display_name', p_key) md_parts.append(f"### Person: {display_name} (Profile Key: {p_key})") if p_data.get("dobs"): md_parts.append(f"* Known DOB(s): {', '.join(p_data['dobs'])}") if p_data.get("ids"): md_parts.append(f"* Known ID(s): {', '.join(p_data['ids'])}") if p_data.get("face_gallery") and len(p_data.get("face_gallery")) > 0: md_parts.append(f"* Facial Signatures Stored: {len(p_data.get('face_gallery'))}") md_parts.append("* Documents:") doc_ids_for_person = sorted(list(p_data.get("doc_ids", set()))) # Sort for consistency 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") linking_method = doc_detail.get("linking_method", "") md_parts.append(f" - {filename} (`{doc_type}`) {linking_method}") 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: # Expected number of output components: df_data, persons_md, ocr_json_output, status_textbox yield ([["N/A", "ERROR", "API Key Missing", "N/A","N/A", "N/A", "N/A", "N/A","N/A", "N/A"]], "API Key Missing.", "{}", "Error: API Key not set.") return if not files_list: yield ([], "No files uploaded.", "{}", "Upload files to begin.") return # Initialize file data structures for i, file_obj_path in enumerate(files_list): # gr.Files with type="filepath" returns list of path strings doc_uid = str(uuid.uuid4()) processed_files_data.append({ "doc_id": doc_uid, "filename": os.path.basename(file_obj_path), "filepath": file_obj_path, "status": "Queued", "ocr_json": None, "entities": None, "face_analysis_result": None, "facial_embeddings": None, "assigned_person_key": None, "linking_method": "" }) df_data = format_dataframe_data(processed_files_data) persons_md = format_persons_markdown(person_profiles, processed_files_data) yield (df_data, persons_md, "{}", f"Initialized {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"] linking_method_log_for_doc = [] # To store how this doc was linked if not file_data_item["filepath"] or not os.path.exists(file_data_item["filepath"]): file_data_item["status"] = "Error: Invalid file" linking_method_log_for_doc.append("File path error.") file_data_item["linking_method"] = " ".join(linking_method_log_for_doc) 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 for {current_filename}") continue # 1. OCR file_data_item["status"] = "OCR..." df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, file_data_item.get("ocr_json_str","{}"), f"OCR: {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 Err: {str(ocr_result['error'])[:30]}.." linking_method_log_for_doc.append("OCR Failed.") file_data_item["linking_method"] = " ".join(linking_method_log_for_doc) df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"OCR Err: {current_filename}") continue file_data_item["status"] = "OCR OK. Entities..." df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"Entities: {current_filename}") # 2. Entity Extraction entities = extract_entities_from_ocr(ocr_result) file_data_item["entities"] = entities file_data_item["status"] = "Entities OK. Face..." df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"Face Detect: {current_filename}") # 3. Facial Feature Extraction doc_type_lower = (entities.get("doc_type") or "").lower() # Attempt face detection on photos, passports, IDs. if DEEPFACE_AVAILABLE and ("photo" in doc_type_lower or "passport" in doc_type_lower or "id card" in doc_type_lower or "selfie" in doc_type_lower): face_result = get_facial_embeddings_with_deepface(file_data_item["filepath"]) file_data_item["face_analysis_result"] = face_result if "embeddings" in face_result and face_result["embeddings"]: file_data_item["facial_embeddings"] = face_result["embeddings"] file_data_item["status"] = f"Face OK ({face_result.get('count',0)}). Classify..." linking_method_log_for_doc.append(f"{face_result.get('count',0)} face(s).") elif "error" in face_result: file_data_item["status"] = f"Face Err: {face_result['error'][:20]}.." linking_method_log_for_doc.append("Face Ext. Error.") else: # No error, but no embeddings (e.g. no face detected) file_data_item["status"] = "No Face. Classify..." linking_method_log_for_doc.append("No face det.") else: file_data_item["status"] = "No Face Ext. Classify..." linking_method_log_for_doc.append("Face Ext. Skipped.") df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"Classifying: {current_filename}") # 4. Person Classification person_key = get_person_id_and_update_profiles(current_doc_id, entities, file_data_item.get("facial_embeddings"), person_profiles, linking_method_log_for_doc) file_data_item["assigned_person_key"] = person_key file_data_item["status"] = "Classified" file_data_item["linking_method"] = " ".join(linking_method_log_for_doc) 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"Done: {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 v2 (with Face ID)") gr.Markdown( "**Upload multiple documents. The system will OCR, extract entities & faces, and classify documents by person.**\n" "Ensure `OPENROUTER_API_KEY` is set as a Secret. Facial recognition uses `deepface` ('VGG-Face' model, 'retinaface' detector)." ) if not OPENROUTER_API_KEY: gr.Markdown("