Update app.py

app.py

@@ -1,13 +1,11 @@
 from paddleocr import PaddleOCR
 from gliner import GLiNER
-import json
 from PIL import Image
 import gradio as gr
 import numpy as np
 import cv2
 import logging
 import os
-from pathlib import Path
 import tempfile
 import pandas as pd
 import io
@@ -29,16 +27,8 @@ except Exception as e:
     logger.error("Failed to load GLiNER model")
     raise e
 
-# Helper functions
-# Get a random color (used for drawing bounding boxes, if needed)
 def get_random_color():
-    return tuple(np.random.randint(0, 256, 3).tolist())  # Fixed line
-
-def draw_ocr_bbox(image, boxes, colors):
-    for i in range(len(boxes)):
-        box = np.reshape(np.array(boxes[i]), [-1, 1, 2]).astype(np.int64)
-        image = cv2.polylines(np.array(image), [box], True, colors[i], 2)
-    return image
+    return tuple(np.random.randint(0, 256, 3).tolist()
 
 def scan_qr_code(image):
     try:
@@ -55,32 +45,45 @@ def extract_emails(text):
     return re.findall(email_regex, text)
 
 def extract_websites(text):
-    website_regex = r"(?:https?://)?(?:www\.)?[A-Za-z0-9-]+\.[A-Za-z]{2,}(?:/\S*)?"
+    website_regex = r"\b(?:https?://)?(?:www\.)?([A-Za-z0-9-]+\.[A-Za-z]{2,})(?:/\S*)?\b"
     matches = re.findall(website_regex, text)
     return [m for m in matches if '@' not in m]
 
 def clean_phone_number(phone):
-    return re.sub(r"[^\d+]", "", phone)
+    cleaned = re.sub(r"(?!^\+)[^\d]", "", phone)
+    if len(cleaned) < 9 or (len(cleaned) == 9 and cleaned.startswith("+")):
+        return None
+    return cleaned
+
+def normalize_website(url):
+    url = url.lower().replace("www.", "").split('/')[0]
+    if not re.match(r"^[a-z0-9-]+\.[a-z]{2,}$", url):
+        return None
+    return f"www.{url}"
+
+def extract_address(ocr_texts):
+    address_keywords = ["block", "street", "ave", "area", "industrial", "road"]
+    address_parts = []
+    for text in ocr_texts:
+        if any(kw in text.lower() for kw in address_keywords):
+            address_parts.append(text)
+    return " ".join(address_parts) if address_parts else None
 
-# Main inference function
 def inference(img: Image.Image, confidence):
     try:
-        # Initialize PaddleOCR
         ocr = PaddleOCR(use_angle_cls=True, lang='en', use_gpu=False,
                         det_model_dir='./models/det/en',
                         cls_model_dir='./models/cls/en',
                         rec_model_dir='./models/rec/en')
         
-        # OCR Processing
         img_np = np.array(img)
         result = ocr.ocr(img_np, cls=True)[0]
         ocr_texts = [line[1][0] for line in result]
         ocr_text = " ".join(ocr_texts)
 
-        # Entity Extraction
         labels = ["person name", "company name", "job title", 
-                "phone number", "email address", "physical address", 
-                "website url"]
+                "phone number", "email address", "address", 
+                "website"]
         entities = gliner_model.predict_entities(ocr_text, labels, threshold=confidence, flat_ner=True)
         
         results = {
@@ -89,70 +92,94 @@ def inference(img: Image.Image, confidence):
             "Job Title": [],
             "Phone Number": [],
             "Email Address": [],
-            "Physical Address": [],
-            "Website Url": [],
+            "Address": [],
+            "Website": [],
             "QR Code": []
         }
 
-        # Process GLiNER results
+        # Process entities with validation
         for entity in entities:
-            label = entity["label"].title().replace(" ", "")
-            if label == "PhoneNumber":
-                cleaned = clean_phone_number(entity["text"])
-                if cleaned: results["Phone Number"].append(cleaned)
-            elif label == "EmailAddress":
-                results["Email Address"].append(entity["text"].lower())
-            elif label == "WebsiteUrl":
-                results["Website Url"].append(entity["text"].lower())
-            elif label in results:
-                results[label].append(entity["text"])
+            text = entity["text"].strip()
+            label = entity["label"].lower()
+            
+            if label == "phone number":
+                if (cleaned := clean_phone_number(text)):
+                    results["Phone Number"].append(cleaned)
+            elif label == "email address" and "@" in text:
+                results["Email Address"].append(text.lower())
+            elif label == "website":
+                if (normalized := normalize_website(text)):
+                    results["Website"].append(normalized)
+            elif label == "address":
+                results["Address"].append(text)
+            elif label == "company name":
+                results["Company Name"].append(text)
+            elif label == "person name":
+                results["Person Name"].append(text)
+            elif label == "job title":
+                results["Job Title"].append(text.title())
 
         # Regex fallbacks
-        if not results["Email Address"]:
-            results["Email Address"] = extract_emails(ocr_text)
+        results["Email Address"] += extract_emails(ocr_text)
+        results["Website"] += [normalize_website(w) for w in extract_websites(ocr_text)]
         
-        if not results["Website Url"]:
-            results["Website Url"] = extract_websites(ocr_text)
-
         # Phone number validation
-        phone_numbers = []
-        for text in ocr_texts:
-            numbers = re.findall(r'(?:\+?[0-9]\s?[0-9]+)+', text)
-            phone_numbers.extend([clean_phone_number(n) for n in numbers])
-        results["Phone Number"] = list(set(phone_numbers + results["Phone Number"]))
-
-        # QR Code handling
-        qr_data = scan_qr_code(img)
-        if qr_data:
-            results["QR Code"] = [qr_data]
+        seen_phones = set()
+        for phone in results["Phone Number"] + re.findall(r'\+\d{8,}|\d{9,}', ocr_text):
+            if (cleaned := clean_phone_number(phone)) and cleaned not in seen_phones:
+                results["Phone Number"].append(cleaned)
+                seen_phones.add(cleaned)
+        results["Phone Number"] = list(seen_phones)
+
+        # Address processing
+        if not results["Address"]:
+            if (address := extract_address(ocr_texts)):
+                results["Address"].append(address)
+
+        # Website normalization
+        seen_websites = set()
+        final_websites = []
+        for web in results["Website"]:
+            if web and web not in seen_websites:
+                final_websites.append(web)
+                seen_websites.add(web)
+        results["Website"] = final_websites
+
+        # Company name fallback
+        if not results["Company Name"]:
+            if results["Email Address"]:
+                domain = results["Email Address"][0].split('@')[-1].split('.')[0]
+                results["Company Name"].append(domain.title())
+            elif results["Website"]:
+                domain = results["Website"][0].split('.')[1]
+                results["Company Name"].append(domain.title())
+
+        # Name fallback
+        if not results["Person Name"]:
+            for text in ocr_texts:
+                if re.match(r"^(?:[A-Z][a-z]+\s?){2,}$", text):
+                    results["Person Name"].append(text)
+                    break
+
+        # QR Code
+        if (qr_data := scan_qr_code(img)):
+            results["QR Code"].append(qr_data)
 
         # Create CSV
         csv_data = {k: "; ".join(v) for k, v in results.items() if v}
-        csv_io = io.BytesIO()
-        pd.DataFrame([csv_data]).to_csv(csv_io, index=False)
-        csv_io.seek(0)
-        
-        with tempfile.NamedTemporaryFile(suffix=".csv", delete=False, mode="wb") as tmp_file:
-            tmp_file.write(csv_io.getvalue())
+        with tempfile.NamedTemporaryFile(suffix=".csv", delete=False, mode="w") as tmp_file:
+            pd.DataFrame([csv_data]).to_csv(tmp_file, index=False)
             csv_path = tmp_file.name
 
         return ocr_text, csv_data, csv_path, ""
-    
+
     except Exception as e:
         logger.error(f"Processing failed: {traceback.format_exc()}")
         return "", {}, None, f"Error: {str(e)}\n{traceback.format_exc()}"
 
 # Gradio Interface
 title = 'Enhanced Business Card Parser'
-description = 'Extracts entities with combined AI and regex validation, including QR codes'
-
-examples = [
-    ['example_imgs/example.jpg', 0.4],
-    ['example_imgs/demo003.jpeg', 0.5],
-]
-
-css = """.output_image, .input_image {height: 40rem !important; width: 100% !important;}
-         .gr-interface {max-width: 800px !important;}"""
+description = 'Accurate entity extraction with combined AI and regex validation'
 
 if __name__ == '__main__':
     demo = gr.Interface(
@@ -165,9 +192,6 @@ if __name__ == '__main__':
          gr.Textbox(label="Error Log")],
         title=title,
         description=description,
-        examples=examples,
-        css=css,
-        cache_examples=True
+        css=".gr-interface {max-width: 800px !important;}"
     )
-    demo.queue(max_size=20)
     demo.launch()