Add application file

app.py

@@ -6,20 +6,45 @@ 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 ---
-# 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_KEY = os.environ.get("OPENROUTER_API_KEY")
+IMAGE_MODEL = "opengvlab/internvl3-14b:free"
 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
+# 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.
 
-# --- Helper Functions ---
+# --- 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."}
@@ -42,68 +67,56 @@ def extract_json_from_text(text):
                 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}
+                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", "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.
+- "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.  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.
+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"
+        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,
+            "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", 
+            "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)
@@ -114,139 +127,250 @@ def call_openrouter_ocr(image_filepath):
             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.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}"
+        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:
-        return {"error": f"An unexpected error occurred during OCR: {str(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 "extracted_fields" not in ocr_json or not isinstance(ocr_json.get("extracted_fields"), dict):
+    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): # ocr_json itself might be an error dict
+        if isinstance(ocr_json, 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}
+        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")
-    name_keys = ["full name", "name", "account holder name", "guest name"]
+
+    # 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"]
-    passport_keys = ["document number", "passport number"]
+    # 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():
-                extracted_name = str(value) if value else None
-                break
-        if extracted_name: break
+                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():
-                extracted_dob = str(value) if value else None
+            if key == field_key.lower() and value and isinstance(value, str):
+                extracted_dob = value.strip()
                 break
         if extracted_dob: break
-    extracted_passport_no = None
-    for key in passport_keys:
+    
+    extracted_main_id = None
+    for key in id_keys:
         for field_key, value in fields.items():
-            if key == field_key.lower():
-                extracted_passport_no = str(value).replace(" ", "").upper() if value else None
+            if key == field_key.lower() and value and isinstance(value, str):
+                extracted_main_id = value.replace(" ", "").upper().strip() # Normalize
                 break
-        if extracted_passport_no: break
+        if extracted_main_id: break
+        
     return {
         "name": extracted_name,
         "dob": extracted_dob,
-        "passport_no": extracted_passport_no,
-        "doc_type": doc_type
+        "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 get_person_id_and_update_profiles(doc_id, entities, current_persons_data):
-    passport_no = entities.get("passport_no")
+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")
-    if passport_no:
+
+    # Tier 1: Match by Main ID (Passport, National ID, etc.)
+    if main_id:
         for p_key, p_data in current_persons_data.items():
-            if passport_no in p_data.get("passport_numbers", set()):
+            if main_id in p_data.get("ids", 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
+                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_key = f"person_{passport_no}"
+        # New person based on this main_id
+        new_person_key = f"person_id_{main_id}"
         current_persons_data[new_person_key] = {
-            "canonical_name": name, "canonical_dob": dob,
+            "display_name": name or f"Person (ID: {main_id})",
             "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})"
+            "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)
-        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)
+                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_key = f"person_{composite_key_nd}_{str(uuid.uuid4())[:4]}"
+        # 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] = {
-            "canonical_name": name, "canonical_dob": dob,
-            "names": {norm_name}, "dobs": {dob},
-            "passport_numbers": set(), "doc_ids": {doc_id},
-            "display_name": name
+            "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)
-        new_person_key = f"person_{norm_name}_{str(uuid.uuid4())[:4]}"
+        # 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] = {
-            "canonical_name": name, "canonical_dob": None,
-            "names": {norm_name}, "dobs": set(), "passport_numbers": set(),
-            "doc_ids": {doc_id}, "display_name": name
+            "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] = {
-        "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]})"
+        "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 {} # CORRECTED LINE HERE
+        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("passport_no", "N/A"),
-            f_data.get("assigned_person_key", "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."
+    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():
+    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("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'])}")
+        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 = p_data.get("doc_ids", set())
+        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)
@@ -254,11 +378,10 @@ def format_persons_markdown(current_persons_data, current_files_data):
                     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.")
+                    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)
 
@@ -267,125 +390,146 @@ def process_uploaded_files(files_list, progress=gr.Progress(track_tqdm=True)):
     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."
-        )
+        # 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
-    for i, file_obj in enumerate(files_list):
+
+    # 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.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
+            "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": ""
         })
-    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.")
+    
+    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"]
-        if not file_data_item["filepath"]: # Check if filepath is valid
-            file_data_item["status"] = "Error: Invalid file path"
+        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 with file {current_filename}")
+            yield(df_data, persons_md, "{}", f"({i+1}/{len(processed_files_data)}) Error for {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}")
+        # 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 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}")
+            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 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}")
+        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 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["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"({i+1}/{len(processed_files_data)}) Classified {current_filename} -> {person_key}")
+        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")
+    gr.Markdown("# 📄 Intelligent Document Processor & Classifier v2 (with Face ID)")
     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."
+        "**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("<h3 style='color:red;'>⚠️ ERROR: `OPENROUTER_API_KEY` is not set in Space Secrets! OCR will fail.</h3>")
+    if not OPENROUTER_API_KEY: gr.Markdown("<h3 style='color:red;'>⚠️ ERROR: `OPENROUTER_API_KEY` Secret missing! OCR will fail.</h3>")
+    if not DEEPFACE_AVAILABLE: gr.Markdown("<h3 style='color:orange;'>⚠️ WARNING: `deepface` library not installed. Facial recognition features are disabled.</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
+            files_input = gr.Files(label="Upload Document Images (Bulk)", file_count="multiple", type="filepath")
             process_button = gr.Button("🚀 Process Uploaded Documents", variant="primary")
-            overall_status_textbox = gr.Textbox(label="Overall Progress", interactive=False, lines=1)
+        with gr.Column(scale=2):
+            overall_status_textbox = gr.Textbox(label="Current Task & Overall Progress", interactive=False, lines=2)
+    
     gr.Markdown("---")
     gr.Markdown("## Document Processing Details")
-    dataframe_headers = ["Doc ID (short)", "Filename", "Status", "Detected Type", "Name", "DOB", "Passport No.", "Assigned Person Key"]
+    dataframe_headers = ["Doc ID", "Filename", "Status", "Type", "Face?", "Name", "DOB", "Main ID", "Person Key", "Linking Method"]
     document_status_df = gr.Dataframe(
-        headers=dataframe_headers,
-        datatype=["str"] * len(dataframe_headers),
+        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
+        row_count=(1, "dynamic"), col_count=(len(dataframe_headers), "fixed"), wrap=True, height=400
     )
-    ocr_json_output = gr.Code(label="Selected Document OCR JSON", language="json", interactive=False)
+    
+    with gr.Accordion("Selected Document Full OCR JSON", open=False):
+        ocr_json_output = gr.Code(label="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
-        ]
+        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 "{}"
+        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.
+        # Access global state. Be cautious with globals in complex apps.
         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)
+        return json.dumps({"message": "No OCR data or selection out of bounds."}, indent=2)
 
 if __name__ == "__main__":
     demo.queue().launch(debug=True, share=os.environ.get("GRADIO_SHARE", "true").lower() == "true")

old_app.py

@@ -1,70 +1,91 @@
 import gradio as gr
-import requests
 import base64
-import os
+import requests
 import json
-import mimetypes
+import re
+import os
+import uuid
+from datetime import datetime
 
 # --- Configuration ---
-OPENROUTER_API_KEY = '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'
-IMAGE_MODEL = "opengvlab/internvl3-14b:free"
+# 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"
 
-# --- Application State ---
-current_batch = []
+# --- 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 generate_extraction_prompt(doc_type_provided_by_user):
-    prompt = f"""You are an advanced OCR and information extraction AI.
-The user has provided an image and identified it as a '{doc_type_provided_by_user}'.
+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_provided": (string) The type provided by the user: "{doc_type_provided_by_user}".
-- "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", "Boarding Pass", "Photograph of a person").
+- "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", "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", "Room Type".
-    - For photos: "Description" (e.g., "Portrait of a person", "Image contains text: [text if any]").
+    - 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 (e.g., "passport_type", "issuing_country", "surname", "given_names", "passport_number", "nationality", "dob", "sex", "expiry_date", "personal_number").
+    - "parsed_mrz": (object) Key-value pairs of parsed MRZ fields.
     If no MRZ, this field should be null.
-- "multilingual_info": (array of objects or null) For any text segments not in English:
-    - Each object: {{"language_detected": "ISO 639-1 code", "original_text": "...", "english_translation_or_transliteration": "..."}}
-    If no non-English text, this field can be null or an empty array.
 - "full_text_ocr": (string) Concatenation of all text found on the document.
 
 Extraction Guidelines:
-1.  Prioritize accuracy. If unsure about a character or word, indicate uncertainty if possible, or extract the most likely interpretation.
-2.  Extract all visible text, including small print, stamps, and handwritten annotations if legible.
+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.
-4.  If the image is a photo of a person without much text, the "extracted_fields" might contain a description, and "full_text_ocr" might be minimal.
-5.  If the document is multi-page and only one page is provided, note this if apparent.
 
 Ensure the entire output strictly adheres to the JSON format.
 """
-    return prompt
 
-def process_single_image_with_openrouter(image_path, doc_type):
+def call_openrouter_ocr(image_filepath):
     if not OPENROUTER_API_KEY:
-        return {"error": "OpenRouter API key not set.", "document_type_provided": doc_type}
+        return {"error": "OpenRouter API Key not configured."}
     try:
-        with open(image_path, "rb") as f:
-            encoded_image_bytes = f.read()
-            encoded_image_string = base64.b64encode(encoded_image_bytes).decode("utf-8")
-        mime_type, _ = mimetypes.guess_type(image_path)
-        if not mime_type:
-            ext = os.path.splitext(image_path)[1].lower()
-            if ext == ".png": mime_type = "image/png"
-            elif ext in [".jpg", ".jpeg"]: mime_type = "image/jpeg"
-            elif ext == ".webp": mime_type = "image/webp"
-            else: mime_type = "image/jpeg"
-        data_url = f"data:{mime_type};base64,{encoded_image_string}"
-        prompt_text = generate_extraction_prompt(doc_type)
+        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": [
@@ -76,192 +97,295 @@ def process_single_image_with_openrouter(image_path, doc_type):
                     ]
                 }
             ],
-            "max_tokens": 3000,
+            "max_tokens": 3500,
             "temperature": 0.1,
         }
         headers = {
             "Authorization": f"Bearer {OPENROUTER_API_KEY}",
             "Content-Type": "application/json",
-            "HTTP-Referer": "https://huggingface.co/spaces/Passport_Extractor",
-            "X-Title": "Document Classifier"
+            "HTTP-Referer": "https://huggingface.co/spaces/YOUR_SPACE", 
+            "X-Title": "Gradio Document Processor"
         }
-        print(f"Sending request to OpenRouter for image: {os.path.basename(image_path)}, type: {doc_type}")
-        response = requests.post(OPENROUTER_API_URL, headers=headers, json=payload, timeout=120)
+        response = requests.post(OPENROUTER_API_URL, headers=headers, json=payload, timeout=180)
         response.raise_for_status()
         result = response.json()
-        print(f"Received response from OpenRouter. Status: {response.status_code}")
         if "choices" in result and result["choices"]:
-            content_text = result["choices"][0]["message"]["content"]
-            clean_content = content_text.strip()
-            if clean_content.startswith("```json"):
-                clean_content = clean_content[7:]
-                if clean_content.endswith("```"):
-                    clean_content = clean_content[:-3]
-            elif clean_content.startswith("`") and clean_content.endswith("`"):
-                 clean_content = clean_content[1:-1]
-            try:
-                parsed_json = json.loads(clean_content)
-                if "document_type_provided" not in parsed_json:
-                    parsed_json["document_type_provided"] = doc_type
-                return parsed_json
-            except json.JSONDecodeError as e:
-                print(f"JSONDecodeError: {e}. Raw content was:\n{content_text}")
-                return {
-                    "error": "Failed to parse LLM output as JSON.",
-                    "raw_content_from_llm": content_text,
-                    "document_type_provided": doc_type
-                }
+            raw_content = result["choices"][0]["message"]["content"]
+            return extract_json_from_text(raw_content)
         else:
-            print(f"No 'choices' in API response: {result}")
-            return {"error": "No choices in API response.", "details": result, "document_type_provided": doc_type}
+            return {"error": "No 'choices' in API response from OpenRouter.", "details": result}
     except requests.exceptions.Timeout:
-        print(f"API Request Timeout for {os.path.basename(image_path)}")
-        return {"error": "API request timed out.", "document_type_provided": doc_type}
+        return {"error": "API request timed out."}
     except requests.exceptions.RequestException as e:
         error_message = f"API Request Error: {str(e)}"
-        if e.response is not None:
+        if hasattr(e, 'response') and e.response is not None:
             error_message += f" Status: {e.response.status_code}, Response: {e.response.text}"
-        print(error_message)
-        return {"error": error_message, "document_type_provided": doc_type}
+        return {"error": error_message}
     except Exception as e:
-        print(f"An unexpected error occurred during processing {os.path.basename(image_path)}: {str(e)}")
-        return {"error": f"An unexpected error: {str(e)}", "document_type_provided": doc_type}
+        return {"error": f"An unexpected error occurred during OCR: {str(e)}"}
 
-def add_document_to_batch_ui(image_filepath, doc_type_selection):
-    global current_batch
-    if image_filepath and doc_type_selection:
-        filename = os.path.basename(image_filepath)
-        current_batch.append({"path": image_filepath, "type": doc_type_selection, "filename": filename})
-        batch_display_data = [[item["filename"], item["type"]] for item in current_batch]
-        return batch_display_data, f"Added '{filename}' as '{doc_type_selection}'."
-    batch_display_data = [[item["filename"], item["type"]] for item in current_batch]
-    return batch_display_data, "Failed to add: Image or document type missing."
+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}
 
-def process_batch_ui():
-    global current_batch
-    if not OPENROUTER_API_KEY:
-        return {"error": "OPENROUTER_API_KEY is not set. Please configure it."}, "API Key Missing."
-    if not current_batch:
-        return {"message": "Batch is empty. Add documents first."}, "Batch is empty."
-    all_results = []
-    status_updates = []
-    for i, item_to_process in enumerate(current_batch):
-        status_msg = f"Processing document {i+1}/{len(current_batch)}: {item_to_process['filename']} ({item_to_process['type']})..."
-        print(status_msg)
-        extracted_data = process_single_image_with_openrouter(item_to_process["path"], item_to_process["type"])
-        all_results.append(extracted_data)
-        if "error" in extracted_data:
-            status_updates.append(f"Error processing {item_to_process['filename']}: {extracted_data['error']}")
-        else:
-            status_updates.append(f"Successfully processed {item_to_process['filename']}.")
-    grouped_by_person = {}
-    unidentified_docs = []
-    for result_item in all_results:
-        doc_id = None
-        if isinstance(result_item, dict) and "extracted_fields" in result_item and isinstance(result_item["extracted_fields"], dict):
-            fields = result_item["extracted_fields"]
-            passport_no = fields.get("Document Number") or fields.get("Passport Number") or fields.get("passport_number")
-            name = fields.get("Given Names") or fields.get("Given Name") or fields.get("Name")
-            surname = fields.get("Surname") or fields.get("Family Name")
-            dob = fields.get("Date of Birth") or fields.get("DOB")
-            if passport_no:
-                doc_id = f"passport_{str(passport_no).replace(' ', '').lower()}"
-            elif name and surname and dob:
-                doc_id = f"{str(name).replace(' ', '').lower()}_{str(surname).replace(' ', '').lower()}_{str(dob).replace(' ', '')}"
-            elif name and surname:
-                 doc_id = f"{str(name).replace(' ', '').lower()}_{str(surname).replace(' ', '').lower()}"
-        if doc_id:
-            if doc_id not in grouped_by_person:
-                grouped_by_person[doc_id] = {"person_identifier": doc_id, "documents": []}
-            grouped_by_person[doc_id]["documents"].append(result_item)
-        else:
-            unidentified_docs.append(result_item)
-    final_structured_output = {
-        "summary": f"Processed {len(current_batch)} documents.",
-        "grouped_by_person": list(grouped_by_person.values()) if grouped_by_person else [],
-        "unidentified_documents_or_errors": unidentified_docs
+    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
     }
-    final_status = "Batch processing complete. " + " | ".join(status_updates)
-    print(final_status)
-    return final_structured_output, final_status
 
-def clear_batch_ui():
-    global current_batch
-    current_batch = []
-    return [], "Batch cleared successfully."
+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("# 📄 Document Information Extractor (OpenGVLab/InternVL3-14B via OpenRouter)")
+    gr.Markdown("# 📄 Intelligent Document Processor & Classifier")
     gr.Markdown(
-        "**Instructions:**\n"
-        "1. Upload a document image (e.g., passport front/back, photo, hotel reservation).\n"
-        "2. Select the correct document type.\n"
-        "3. Click 'Add Document to Current Batch'. Repeat for all documents of a person or a related set.\n"
-        "4. Review the batch. Click 'Clear Entire Batch' to start over.\n"
-        "5. Click 'Process Batch and Extract Information' to send documents to the AI.\n"
-        "6. View the extracted information in JSON format below."
+        "**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;'>⚠️ Warning: `OPENROUTER_API_KEY` environment variable is not detected. "
-            "API calls will fail. Please set it and restart this application.</h3>"
-        )
+        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):
-            gr.Markdown("### Step 1: Add Document")
-            image_input = gr.Image(
-                label="Upload Document Image",
-                type="filepath",
-                sources=["upload"],
-                height=300
-            )
-            doc_type_choices = [
-                'passport_front', 'passport_back', 'national_id_front', 'national_id_back',
-                'drivers_license_front', 'drivers_license_back', 'visa_sticker',
-                'photo', 'hotel_reservation', 'boarding_pass', 'utility_bill', 'other_document'
-            ]
-            doc_type_input = gr.Dropdown(
-                label="Select Document Type",
-                choices=doc_type_choices,
-                value='passport_front',
-                filterable=True
-            )
-            add_button = gr.Button("➕ Add Document to Current Batch", variant="secondary")
-        with gr.Column(scale=2):
-            gr.Markdown("### Step 2: Review Current Batch")
-            batch_dataframe = gr.Dataframe(
-                headers=["Filename", "Document Type"],
-                datatype=["str", "str"],
-                row_count=1,  # Changed: Start with 1 row, should grow dynamically
-                col_count=2,  # Changed: Simpler integer for fixed columns
-                wrap=True
-            )
-            clear_batch_button = gr.Button("🗑️ Clear Entire Batch", variant="stop")
-    gr.Markdown("### Step 3: Process Batch")
-    process_button = gr.Button("🚀 Process Batch and Extract Information", variant="primary")
-    status_message_textbox = gr.Textbox(label="Processing Status", interactive=False, lines=2)
-    gr.Markdown("### Step 4: View Results")
-    output_json_display = gr.JSON(label="Extracted Information (JSON Format)")
-    add_button.click(
-        fn=add_document_to_batch_ui,
-        inputs=[image_input, doc_type_input],
-        outputs=[batch_dataframe, status_message_textbox]
-    ).then(lambda: None, outputs=image_input)
-    clear_batch_button.click(
-        fn=clear_batch_ui,
-        inputs=[],
-        outputs=[batch_dataframe, status_message_textbox]
+            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_batch_ui,
-        inputs=[],
-        outputs=[output_json_display, status_message_textbox]
+        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__":
-    if not OPENROUTER_API_KEY:
-        print("ERROR: The OPENROUTER_API_KEY environment variable is not set.")
-        print("Please set it before running the application, e.g.:")
-        print("  export OPENROUTER_API_KEY='your_openrouter_key_here'")
-        print("The application will launch, but API calls will fail.")
-    demo.launch(share=True) # Added share=True
+    demo.queue().launch(debug=True, share=os.environ.get("GRADIO_SHARE", "true").lower() == "true")

requirements.txt

@@ -1,4 +1,8 @@
-gradio>=4.26.0
-requests>=2.25.0,<3.0.0
-# pillow might be needed explicitly if not pulled by gradio for image handling
-Pillow>=9.0.0
+gradio>=4.0.0
+requests>=2.25.0
+Pillow>=9.0.0
+deepface>=0.0.79
+tensorflow-cpu>=2.10.0 # Or tensorflow-cpu if GPU is not available/needed
+opencv-python-headless>=4.5.0
+# retina-face Pypi package for the detector if deepface doesn't pull it correctly
+retina-face>=0.0.12