Update app.py

app.py

@@ -27,10 +27,11 @@ except Exception:
     logger.exception("Failed to load GLiNER model")
     raise
 
-# Regex patterns
+# Regex patterns for emails and websites
 EMAIL_REGEX = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b")
-WEBSITE_REGEX = re.compile(r"^(?:https?://)?(?:www\.)?([A-Za-z0-9-]+\.[A-Za-z]{2,})(?:/\S*)?$")
-# Phone number constants and regex
+WEBSITE_REGEX = re.compile(r"(?:https?://)?(?:www\.)?([A-Za-z0-9-]+\.[A-Za-z]{2,})")
+
+# Phone number constants and regex for Saudi/UAE support
 SAUDI_CODE = '+966'
 UAE_CODE = '+971'
 PHONE_REGEX = re.compile(r'^(?:\+9665\d{8}|\+9715\d{8}|05\d{8}|5\d{8})$')
@@ -42,6 +43,10 @@ def extract_emails(text: str) -> list[str]:
 def extract_websites(text: str) -> list[str]:
     return [m.lower() for m in WEBSITE_REGEX.findall(text)]
 
+def normalize_website(url: str) -> str | None:
+    u = url.lower().replace('www.', '').split('/')[0]
+    return f"www.{u}" if re.match(r"^[a-z0-9-]+\.[a-z]{2,}$", u) else None
+
 def clean_phone_number(phone: str) -> str | None:
     cleaned = re.sub(r"[^\d+]", "", phone)
     # International formats
@@ -51,11 +56,12 @@ def clean_phone_number(phone: str) -> str | None:
         return cleaned
     # Local to international
     if cleaned.startswith('05') and len(cleaned) == 10:
-        return f"{UAE_CODE}{cleaned[1:]}"
+        # Determine country by leading digit after 0 (6 Saudi, 5 UAE)
+        return (SAUDI_CODE if cleaned[1]=='5' and cleaned[1:2] == '5' else UAE_CODE) + cleaned[1:]
     if cleaned.startswith('5') and len(cleaned) == 9:
-        return f"{UAE_CODE}{cleaned}"
+        return UAE_CODE + cleaned
     if cleaned.startswith('9665') and len(cleaned) == 12:
-        return f"+{cleaned}"
+        return '+' + cleaned
     return None
 
 def process_phone_numbers(text: str) -> list[str]:
@@ -66,12 +72,6 @@ def process_phone_numbers(text: str) -> list[str]:
             found.append(c)
     return list(set(found))
 
-def normalize_website(url: str) -> str | None:
-    u = url.lower().replace('www.', '').split('/')[0]
-    if re.match(r"^[a-z0-9-]+\.[a-z]{2,}$", u):
-        return f"www.{u}"
-    return None
-
 def extract_address(ocr_texts: list[str]) -> str | None:
     keywords = ["block","street","ave","area","industrial","road"]
     parts = [t for t in ocr_texts if any(kw in t.lower() for kw in keywords)]
@@ -119,9 +119,9 @@ def deduplicate_data(results: dict[str, list[str]]) -> None:
                     seen.add(norm); out.append(norm)
         return out
     # Normalize lists
-    results['Email Address'] = clean_list(results['Email Address'], lambda e: e.lower())
-    results['Website'] = clean_list(results['Website'], normalize_website)
-    results['Phone Number'] = clean_list(results['Phone Number'], clean_phone_number)
+    results['Email Address'] = clean_list(results.get('Email Address', []), lambda e: e.lower())
+    results['Website'] = clean_list(results.get('Website', []), normalize_website)
+    results['Phone Number'] = clean_list(results.get('Phone Number', []), clean_phone_number)
     # Others: simple dedupe
     for key in ['Person Name','Company Name','Job Title','Address','QR Code']:
         seen = set(); out = []
@@ -150,27 +150,35 @@ def inference(img: Image.Image, confidence: float):
         # Entity processing
         for ent in entities:
             txt, lbl = ent['text'].strip(), ent['label'].lower()
-            if lbl == 'person name': results['Person Name'].append(txt)
-            elif lbl == 'company name': results['Company Name'].append(txt)
-            elif lbl == 'job title': results['Job Title'].append(txt.title())
+            if lbl == 'person name':
+                results['Person Name'].append(txt)
+            elif lbl == 'company name':
+                results['Company Name'].append(txt)
+            elif lbl == 'job title':
+                results['Job Title'].append(txt.title())
             elif lbl == 'phone number':
-                if (c:=clean_phone_number(txt)): results['Phone Number'].append(c)
+                if (c:=clean_phone_number(txt)):
+                    results['Phone Number'].append(c)
             elif lbl == 'email address' and EMAIL_REGEX.fullmatch(txt):
                 results['Email Address'].append(txt.lower())
-            elif lbl == 'website' and WEBSITE_REGEX.fullmatch(txt):
-                if (n:=normalize_website(txt)): results['Website'].append(n)
-            elif lbl == 'address': results['Address'].append(txt)
+            elif lbl == 'website' and WEBSITE_REGEX.search(txt):
+                if (n:=normalize_website(txt)):
+                    results['Website'].append(n)
+            elif lbl == 'address':
+                results['Address'].append(txt)
         # Regex fallbacks
         results['Email Address'] += extract_emails(full_text)
         results['Website'] += extract_websites(full_text)
         # Phone regex fallback
         results['Phone Number'] += process_phone_numbers(full_text)
-        # QR
-        if qr := scan_qr_code(img): results['QR Code'].append(qr)
+        # QR code
+        if qr := scan_qr_code(img):
+            results['QR Code'].append(qr)
         # Address fallback
         if not results['Address']:
-            if addr := extract_address(ocr_texts): results['Address'].append(addr)
-        # Dedupe
+            if addr := extract_address(ocr_texts):
+                results['Address'].append(addr)
+        # Deduplicate
         deduplicate_data(results)
         # Company fallback
         if not results['Company Name']:
@@ -184,17 +192,18 @@ def inference(img: Image.Image, confidence: float):
         if not results['Person Name']:
             for t in ocr_texts:
                 if re.match(r'^(?:[A-Z][a-z]+\s?){2,}$', t):
-                    results['Person Name'].append(t); break
-        # CSV
-        csv_map = {k: '; '.join(v) for k,v in results.items() if v}
+                    results['Person Name'].append(t)
+                    break
+        # Build CSV map including all keys
+        csv_map = {k: '; '.join(v) for k,v in results.items()}
         with tempfile.NamedTemporaryFile(suffix='.csv', delete=False, mode='w') as f:
             pd.DataFrame([csv_map]).to_csv(f, index=False)
             csv_path = f.name
-        return full_text, csv_map, csv_path, ''
+        return full_text, results, csv_path, ''
     except Exception:
         err = traceback.format_exc()
         logger.error(f"Processing failed: {err}")
-        return '', {}, None, f"Error:\n{err}"
+        return '', {k: [] for k in ['Person Name','Company Name','Job Title','Phone Number','Email Address','Address','Website','QR Code']}, None, f"Error:\n{err}"
 
 # Gradio Interface
 if __name__ == '__main__':
@@ -211,4 +220,3 @@ if __name__ == '__main__':
         css=".gr-interface {max-width: 800px !important;}"
     )
     demo.launch()
-