Upload folder using huggingface_hub

app.py

@@ -1,7 +1,7 @@
-###############################################
-# Bajaj Datathon - FINAL PATCHED BILL EXTRACTOR
-# High Accuracy | Robust OCR | Gemini Refinement
-###############################################
+# app_bill_extractor_final_v2.py
+# Humanized, high-accuracy bill extraction API.
+# Robust OCR preprocessing, TSV layout parsing, numeric-column inference,
+# header prefiltering, deterministic Gemini refinement (if configured).
 
 import os
 import re
@@ -19,42 +19,37 @@ from pytesseract import Output
 import numpy as np
 import cv2
 
-# Optional Gemini SDK
+# Optional: Google Gemini SDK (if available)
 try:
     import google.generativeai as genai
-except:
+except Exception:
     genai = None
 
 # ---------------- LLM CONFIG ----------------
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 GEMINI_MODEL_NAME = os.getenv("GEMINI_MODEL_NAME", "gemini-2.5-flash")
-
 if GEMINI_API_KEY and genai is not None:
     try:
         genai.configure(api_key=GEMINI_API_KEY)
-    except:
+    except Exception:
         pass
 
-
-# ---------------- FASTAPI APP ----------------
-app = FastAPI(title="Bajaj Datathon - Bill Extractor (patched v3)")
+# ---------------- FastAPI app ----------------
+app = FastAPI(title="Bajaj Datathon - Bill Extractor (final, humanized)")
 
 class BillRequest(BaseModel):
     document: str
 
-
-###############################################
-#      COMMON REGEX AND UTILITY FUNCTIONS
-###############################################
-
+# ---------------- Regex and keywords ----------------
 NUM_RE = re.compile(r"[-+]?\d{1,3}(?:[,0-9]*)(?:\.\d+)?")
+TOTAL_KEYWORDS = re.compile(
+    r"(grand\s*total|net\s*payable|total\s*amount|amount\s*payable|bill\s*amount|final\s*amount|balance\s*due|sub\s*total|subtotal|round\s*off)",
+    re.I,
+)
+FOOTER_KEYWORDS = re.compile(r"(page|printed on|printed:|date:|time:|am|pm)", re.I)
 
-HEADER_KEYWORDS = [
-    "description", "qty", "hrs", "rate",
-    "discount", "net", "amt", "amount",
-    "qty/hrs", "qty / hrs"
-]
-
+# generalized header-related tokens & exact header phrase blacklist (common variants)
+HEADER_KEYWORDS = ["description", "qty", "hrs", "rate", "discount", "net", "amt", "amount", "consultation", "qty/hrs", "qty / hrs"]
 HEADER_PHRASES = [
     "description qty / hrs consultation rate discount net amt",
     "description qty / hrs rate discount net amt",
@@ -64,14 +59,7 @@ HEADER_PHRASES = [
 ]
 HEADER_PHRASES = [h.lower() for h in HEADER_PHRASES]
 
-TOTAL_KEYWORDS = re.compile(
-    r"(grand\s*total|net\s*payable|total\s*amount|amount\s*payable|bill\s*amount|final\s*amount|balance\s*due|sub\s*total|subtotal|round\s*off)",
-    re.I,
-)
-
-FOOTER_KEYWORDS = re.compile(r"(page|printed on|printed:|date:|time:|am|pm)", re.I)
-
-
+# ---------------- small utilities ----------------
 def sanitize_ocr_text(s: str) -> str:
     if not s:
         return ""
@@ -80,718 +68,647 @@ def sanitize_ocr_text(s: str) -> str:
     s = s.replace("\r\n", "\n").replace("\r", "\n")
     s = re.sub(r"[ \t]+", " ", s)
     s = s.strip()
-    return s[:5000]
-
+    return s[:4000]
 
 def normalize_num_str(s: Optional[str]) -> Optional[float]:
     if s is None:
         return None
     s = str(s).strip()
-    s = re.sub(r"[^\d\-\+\,\.\(\)]", "", s)
     if s == "":
         return None
+    s = re.sub(r"[^\d\-\+\,\.\(\)]", "", s)
     negative = False
     if s.startswith("(") and s.endswith(")"):
         negative = True
         s = s[1:-1]
     s = s.replace(",", "")
-    try:
-        v = float(s)
-        return -v if negative else v
-    except:
+    if s in ("", "-", "+"):
         return None
-
+    try:
+        return -float(s) if negative else float(s)
+    except Exception:
+        try:
+            return float(s.replace(" ", ""))
+        except Exception:
+            return None
 
 def is_numeric_token(t: Optional[str]) -> bool:
     return bool(t and NUM_RE.search(str(t)))
 
-
 def clean_name_text(s: str) -> str:
     s = s.replace("—", "-")
     s = re.sub(r"\s+", " ", s)
     s = s.strip(" -:,.")
-    # Fix doctor prefix only if followed by name
-    s = re.sub(r"\bOR (?=[A-Z][a-z])", "DR ", s)
+    s = re.sub(r"\bSG0?(\d+)\b", r"SG\1", s, flags=re.I)
+    s = re.sub(r"\b(RR)[\s\-]*2\b", r"RR-2", s, flags=re.I)
+    # fix common OCR mistakes for doctor prefixes
+    s = re.sub(r"\bOR\b", "DR", s)  # sometimes OCR turns 'DR' -> 'OR'
     return s.strip()
 
-
-###############################################
-#          IMAGE PREPROCESSING
-###############################################
-
-def pil_to_cv2(img: Image.Image):
+# ---------------- image preprocessing ----------------
+def pil_to_cv2(img: Image.Image) -> Any:
     arr = np.array(img)
     if arr.ndim == 2:
         return arr
     return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
 
-
-def preprocess_image(pil_img: Image.Image):
+def preprocess_image(pil_img: Image.Image) -> Any:
     pil_img = pil_img.convert("RGB")
     w, h = pil_img.size
-
-    if w < 1500:
-        scale = 1500 / float(w)
+    target_w = 1500
+    if w < target_w:
+        scale = target_w / float(w)
         pil_img = pil_img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
-
-    img = pil_to_cv2(pil_img)
-    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-
+    cv_img = pil_to_cv2(pil_img)
+    gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
     gray = cv2.fastNlMeansDenoising(gray, h=10)
-
     try:
-        bw = cv2.adaptiveThreshold(gray, 255,
-                                   cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
-                                   cv2.THRESH_BINARY, 41, 15)
-    except:
-        _, bw = cv2.threshold(gray, 127, 255,
-                              cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-
-    bw = cv2.morphologyEx(bw, cv2.MORPH_OPEN, np.ones((1, 1), np.uint8))
+        bw = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 41, 15)
+    except Exception:
+        _, bw = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
+    kernel = np.ones((1,1), np.uint8)
+    bw = cv2.morphologyEx(bw, cv2.MORPH_OPEN, kernel)
     return bw
 
-
-###############################################
-#                OCR TSV EXTRACTION
-###############################################
-
-def image_to_tsv_cells(cv_img):
+# ---------------- OCR TSV ----------------
+def image_to_tsv_cells(cv_img: Any) -> List[Dict[str, Any]]:
     try:
-        ocr = pytesseract.image_to_data(
-            cv_img,
-            output_type=Output.DICT,
-            config="--psm 6"
-        )
-    except:
-        ocr = pytesseract.image_to_data(cv_img, output_type=Output.DICT)
-
+        o = pytesseract.image_to_data(cv_img, output_type=Output.DICT, config="--psm 6")
+    except Exception:
+        o = pytesseract.image_to_data(cv_img, output_type=Output.DICT)
     cells = []
-    n = len(ocr.get("text", []))
-
+    n = len(o.get("text", []))
     for i in range(n):
-        t = (ocr["text"][i] or "").strip()
-        if not t:
+        raw = o["text"][i]
+        if raw is None:
+            continue
+        txt = str(raw).strip()
+        if not txt:
             continue
         try:
-            conf = float(ocr["conf"][i])
-        except:
+            conf = float(o["conf"][i]) if o["conf"][i] not in (None, "", "-1") else -1.0
+        except Exception:
             conf = -1.0
-
-        left = int(ocr.get("left", [0])[i])
-        top = int(ocr.get("top", [0])[i])
-        width = int(ocr.get("width", [0])[i])
-        height = int(ocr.get("height", [0])[i])
-
-        cells.append({
-            "text": t,
-            "conf": conf,
-            "left": left,
-            "top": top,
-            "width": width,
-            "height": height,
-            "center_x": left + width / 2,
-            "center_y": top + height / 2,
-        })
+        left = int(o.get("left", [0])[i])
+        top = int(o.get("top", [0])[i])
+        width = int(o.get("width", [0])[i])
+        height = int(o.get("height", [0])[i])
+        center_y = top + height / 2.0
+        center_x = left + width / 2.0
+        cells.append({"text": txt, "conf": conf, "left": left, "top": top, "width": width, "height": height, "center_y": center_y, "center_x": center_x})
     return cells
 
-
-###############################################
-#          GROUPING INTO TEXT LINES
-###############################################
-
-def group_cells_into_rows(cells, y_tol=12):
+# ---------------- grouping & merge helpers ----------------
+def group_cells_into_rows(cells: List[Dict[str, Any]], y_tolerance: int = 12) -> List[List[Dict[str, Any]]]:
     if not cells:
         return []
-    cells = sorted(cells, key=lambda c: (c["center_y"], c["center_x"]))
-
+    sorted_cells = sorted(cells, key=lambda c: (c["center_y"], c["center_x"]))
     rows = []
-    current = [cells[0]]
-    last_y = cells[0]["center_y"]
-
-    for c in cells[1:]:
-        if abs(c["center_y"] - last_y) <= y_tol:
+    current = [sorted_cells[0]]
+    last_y = sorted_cells[0]["center_y"]
+    for c in sorted_cells[1:]:
+        if abs(c["center_y"] - last_y) <= y_tolerance:
             current.append(c)
             last_y = (last_y * (len(current) - 1) + c["center_y"]) / len(current)
         else:
             rows.append(sorted(current, key=lambda cc: cc["left"]))
             current = [c]
             last_y = c["center_y"]
-
     if current:
         rows.append(sorted(current, key=lambda cc: cc["left"]))
-
     return rows
 
-
-###############################################
-#          DOCTOR-NAME MERGING (PATCH)
-###############################################
-
-def merge_multiline_names(rows):
+def merge_multiline_names(rows: List[List[Dict[str, Any]]]) -> List[List[Dict[str, Any]]]:
     if not rows:
         return rows
-
     merged = []
     i = 0
     while i < len(rows):
         row = rows[i]
         tokens = [c["text"] for c in row]
-        joined = " ".join(tokens)
-
         has_num = any(is_numeric_token(t) for t in tokens)
-
-        # --- Doctor Name Merge Fix ---
-        if (not has_num and
-            re.search(r"\bdr\b", joined.lower()) and
-            i + 1 < len(rows)):
-
-            next_tokens = " ".join([c["text"] for c in rows[i + 1]])
-            if not any(is_numeric_token(x) for x in next_tokens.split()):
-                merged_row = row + rows[i + 1]
+        # if row looks pure text and next row contains numbers but short left text tokens, merge
+        if not has_num and i + 1 < len(rows):
+            next_row = rows[i+1]
+            next_tokens = [c["text"] for c in next_row]
+            next_has_num = any(is_numeric_token(t) for t in next_tokens)
+            if next_has_num and len(tokens) >= 2 and len([t for t in next_tokens if not is_numeric_token(t)]) <= 3:
+                merged_row = []
+                min_left = min((c["left"] for c in next_row), default=0)
+                offset = 10
+                for c in row:
+                    newc = c.copy()
+                    newc["left"] = min_left - offset
+                    newc["center_x"] = newc["left"] + newc.get("width", 0) / 2.0
+                    merged_row.append(newc)
+                    offset += 10
+                merged_row.extend(next_row)
+                merged.append(sorted(merged_row, key=lambda cc: cc["left"]))
+                i += 2
+                continue
+        # Additional merge: If a row ends with a trailing token like a doctor's name line with single token and next row also text, merge (helps names split across 2+ lines)
+        if not has_num and i + 1 < len(rows):
+            next_row = rows[i+1]
+            next_tokens = [c["text"] for c in next_row]
+            next_has_num = any(is_numeric_token(t) for t in next_tokens)
+            if not next_has_num and len(tokens) <= 3 and len(next_tokens) <= 3:
+                # merge both textual lines into one (keeps relative left ordering by shifting)
+                merged_row = []
+                min_left = min((c["left"] for c in next_row + row), default=0)
+                offset = 10
+                for c in row + next_row:
+                    newc = c.copy()
+                    if newc["left"] > min_left:
+                        newc["left"] = newc["left"]
+                    else:
+                        newc["left"] = min_left - offset
+                    newc["center_x"] = newc["left"] + newc.get("width", 0) / 2.0
+                    merged_row.append(newc)
+                    offset += 5
                 merged.append(sorted(merged_row, key=lambda cc: cc["left"]))
                 i += 2
                 continue
-
         merged.append(row)
         i += 1
-
     return merged
 
-
-###############################################
-#           DETECT NUMERIC COLUMNS
-###############################################
-
-def detect_numeric_columns(cells, max_cols=4):
+# ---------------- numeric column detection ----------------
+# >>> CHANGE: adaptive clustering (restored to conservative adaptive threshold)
+def detect_numeric_columns(cells: List[Dict[str, Any]], max_columns: int = 4) -> List[float]:
     xs = [c["center_x"] for c in cells if is_numeric_token(c["text"])]
     if not xs:
         return []
-
     xs = sorted(xs)
     if len(xs) == 1:
         return [xs[0]]
-
-    gaps = [xs[i + 1] - xs[i] for i in range(len(xs) - 1)]
+    gaps = [xs[i+1] - xs[i] for i in range(len(xs)-1)]
     mean_gap = float(np.mean(gaps))
     std_gap = float(np.std(gaps)) if len(gaps) > 1 else 0.0
-    thresh = max(30.0, mean_gap + 0.6 * std_gap)
-
+    gap_thresh = max(30.0, mean_gap + 0.6 * std_gap)
     clusters = []
     curr = [xs[0]]
-
     for i, g in enumerate(gaps):
-        if g > thresh and len(clusters) < (max_cols - 1):
+        if g > gap_thresh and len(clusters) < (max_columns - 1):
             clusters.append(curr)
-            curr = [xs[i + 1]]
+            curr = [xs[i+1]]
         else:
-            curr.append(xs[i + 1])
-
+            curr.append(xs[i+1])
     clusters.append(curr)
-
     centers = [float(np.median(c)) for c in clusters]
-    centers = centers[-max_cols:]
+    if len(centers) > max_columns:
+        centers = centers[-max_columns:]
     return sorted(centers)
 
-
-def assign_token_to_column(x, centers):
-    if not centers:
+def assign_token_to_column(token_x: float, column_centers: List[float]) -> Optional[int]:
+    if not column_centers:
         return None
-    dist = [abs(x - c) for c in centers]
-    return int(np.argmin(dist))
-
-
-###############################################
-#      STRONG HEADER DETECTION (PATCHED)
-###############################################
-
-def looks_like_header_text(txt: str, top_of_page=False):
-    if not txt:
-        return False
-
-    t = re.sub(r"\s+", " ", txt.strip().lower())
-
-    patterns = [
-        r"description.*qty",
-        r"qty.*rate",
-        r"rate.*amount",
-        r"net\s*amt",
-        r"discount",
-        r"hrs\s*/\s*qty",
-        r"qty\s*/\s*hrs",
-    ]
-    for p in patterns:
-        if re.search(p, t):
-            return True
-
-    if any(h == t for h in HEADER_PHRASES):
-        return True
+    distances = [abs(token_x - cx) for cx in column_centers]
+    return int(np.argmin(distances))
 
-    hits = sum(1 for k in HEADER_KEYWORDS if k in t)
-    if hits >= 3:
-        return True
-
-    tokens = re.split(r"[ \|,/]+", t)
-    num = sum(1 for tok in tokens if NUM_RE.search(tok))
-    if num >= 3:
-        return True
-
-    if top_of_page and hits >= 2:
-        return True
-
-    return False
-
-
-###############################################
-#             PARSE ROWS INTO ITEMS
-###############################################
-
-def parse_rows_with_columns(rows, cells):
+# ---------------- parsing rows into items ----------------
+def parse_rows_with_columns(rows: List[List[Dict[str, Any]]], page_cells: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    parsed_items = []
     rows = merge_multiline_names(rows)
-    col_centers = detect_numeric_columns(cells)
-
-    parsed = []
+    column_centers = detect_numeric_columns(page_cells, max_columns=4)
 
     for row in rows:
-        texts = [c["text"] for c in row]
-        joined = " ".join(texts).lower()
-
-        if FOOTER_KEYWORDS.search(joined) and not any(is_numeric_token(t) for t in texts):
+        tokens = [c["text"] for c in row]
+        if not tokens:
+            continue
+        joined_lower = " ".join(tokens).lower()
+        if FOOTER_KEYWORDS.search(joined_lower) and not any(is_numeric_token(t) for t in tokens):
             continue
-        if all(not is_numeric_token(t) for t in texts):
+        if all(not is_numeric_token(t) for t in tokens):
             continue
 
+        # gather numeric candidates (unique, filtered)
         numeric_values = []
-        for t in texts:
+        for t in tokens:
             if is_numeric_token(t):
                 v = normalize_num_str(t)
                 if v is not None:
                     numeric_values.append(float(v))
+        # de-duplicate and sort descending (larger candidates first)
+        numeric_values = sorted(list({int(x) if float(x).is_integer() else x for x in numeric_values}), reverse=True)
 
-        # De-duplicate & sort largest first
-        numeric_values = sorted(list({float(v) for v in numeric_values}), reverse=True)
-
-        # Drop tiny noise
-        numeric_values = [v for v in numeric_values if v >= 5 or (v < 5 and len(numeric_values) == 1)]
-
-        if col_centers:
-            left_text = []
-            bucket = {i: [] for i in range(len(col_centers))}
-
+        if column_centers:
+            left_text_parts = []
+            numeric_bucket_map = {i: [] for i in range(len(column_centers))}
             for c in row:
                 t = c["text"]
-                x = c["center_x"]
+                cx = c["center_x"]
                 if is_numeric_token(t):
-                    idx = assign_token_to_column(x, col_centers)
-                    if idx is not None:
-                        bucket[idx].append(t)
+                    col_idx = assign_token_to_column(cx, column_centers)
+                    if col_idx is None:
+                        numeric_bucket_map[len(column_centers) - 1].append(t)
+                    else:
+                        numeric_bucket_map[col_idx].append(t)
                 else:
-                    left_text.append(t)
-
-            name_raw = " ".join(left_text).strip()
-            name = clean_name_text(name_raw)
-
-            N = len(col_centers)
+                    left_text_parts.append(t)
+            raw_name = " ".join(left_text_parts).strip()
+            name = clean_name_text(raw_name) if raw_name else ""
 
-            def pick(k):
-                vals = bucket.get(k, [])
+            num_cols = len(column_centers)
+            def get_bucket(idx):
+                vals = numeric_bucket_map.get(idx, [])
                 return vals[-1] if vals else None
 
-            amount = normalize_num_str(pick(N - 1)) if N >= 1 else None
-            rate   = normalize_num_str(pick(N - 2)) if N >= 2 else None
-            qty    = normalize_num_str(pick(N - 3)) if N >= 3 else None
+            amount = normalize_num_str(get_bucket(num_cols - 1)) if num_cols >= 1 else None
+            rate   = normalize_num_str(get_bucket(num_cols - 2)) if num_cols >= 2 else None
+            qty    = normalize_num_str(get_bucket(num_cols - 3)) if num_cols >= 3 else None
 
-            # fallback amount
             if amount is None:
-                for t in reversed(texts):
+                for t in reversed(tokens):
                     if is_numeric_token(t):
                         amount = normalize_num_str(t)
                         if amount is not None:
                             break
 
-            # strong qty/rate inference
-            if amount is not None and rate is not None:
-                ratio = amount / rate if rate else None
-                if ratio and 1 <= round(ratio) <= 10:
-                    qty = float(round(ratio))
+            # >>> CHANGE: safer inference — skip tiny candidates like 1, enforce qty bounds, require close ratio
+            if amount is not None and numeric_values:
+                # Only accept candidate as rate if candidate >= 2 (or amount is tiny) and candidate < amount
+                for cand in numeric_values:
+                    try:
+                        cand_float = float(cand)
+                    except:
+                        continue
+                    if cand_float <= 1.0:
+                        continue
+                    if amount <= 5 and cand_float < 1.0:
+                        continue
+                    if cand_float >= amount:
+                        continue
+                    ratio = amount / cand_float if cand_float else None
+                    if ratio is None:
+                        continue
+                    r = round(ratio)
+                    if r < 1 or r > 200:
+                        continue
+                    # require relative closeness threshold (adaptive)
+                    if abs(ratio - r) <= max(0.03 * r, 0.15):
+                        # Accept only if qty reasonable (<=100)
+                        if r <= 100:
+                            rate = cand_float
+                            qty = float(r)
+                            break
 
+            # fallback compute rate if qty found but rate missing
+            if (rate is None or rate == 0) and qty and qty != 0 and amount is not None:
+                try:
+                    candidate_rate = amount / qty
+                    # require candidate_rate > 1 (avoid tiny rates) and reasonable
+                    if candidate_rate >= 2:
+                        rate = candidate_rate
+                except Exception:
+                    pass
+
+            # final defaults
             if qty is None:
                 qty = 1.0
 
-            if amount == 0 and rate and qty:
-                amount = rate * qty
-
-            try: amount = float(round(amount, 2))
-            except: continue
-
-            try: rate = float(round(rate or 0.0, 2))
-            except: rate = 0.0
-
-            try: qty = float(qty)
-            except: qty = 1.0
+            # normalize and sanity-check
+            try:
+                amount = float(round(amount, 2))
+            except Exception:
+                continue
+            try:
+                rate = float(round(rate, 2)) if rate is not None else 0.0
+            except Exception:
+                rate = 0.0
+            try:
+                qty = float(qty)
+            except Exception:
+                qty = 1.0
 
-            parsed.append({
+            parsed_items.append({
                 "item_name": name if name else "UNKNOWN",
                 "item_amount": amount,
-                "item_rate": rate,
-                "item_quantity": qty
+                "item_rate": rate if rate is not None else 0.0,
+                "item_quantity": qty if qty is not None else 1.0,
             })
 
         else:
-            idxs = [i for i, t in enumerate(texts) if is_numeric_token(t)]
-            if not idxs:
+            numeric_idxs = [i for i, t in enumerate(tokens) if is_numeric_token(t)]
+            if not numeric_idxs:
                 continue
-
-            amt = normalize_num_str(texts[idxs[-1]])
+            last = numeric_idxs[-1]
+            amt = normalize_num_str(tokens[last])
             if amt is None:
                 continue
-
-            name = " ".join(texts[: idxs[-1]]).strip()
+            name = " ".join(tokens[:last]).strip()
             if not name:
                 continue
+            rate = None; qty = None
 
-            rate = 0.0
-            qty = 1.0
-
-            possible = []
-            for i in idxs:
-                v = normalize_num_str(texts[i])
+            # try to pick rate/qty from previous numeric tokens (right-to-left)
+            # and use the safer inference logic (ignore candidate == 1)
+            right_nums = []
+            for i in numeric_idxs:
+                v = normalize_num_str(tokens[i])
                 if v is not None:
-                    possible.append(float(v))
+                    right_nums.append(float(v))
+            right_nums = sorted(list({int(x) if float(x).is_integer() else x for x in right_nums}), reverse=True)
+
+            # attempt direct mapping: last numeric = amount, previous maybe rate / qty
+            if len(right_nums) >= 2:
+                cand = right_nums[1]
+                if float(cand) > 1 and float(cand) < float(amt):
+                    # check ratio
+                    ratio = float(amt) / float(cand) if cand else None
+                    if ratio:
+                        r = round(ratio)
+                        if 1 <= r <= 200 and abs(ratio - r) <= max(0.03 * r, 0.15) and r <= 100:
+                            rate = float(cand)
+                            qty = float(r)
+            # fallback: conservative search like above
+            if rate is None and right_nums:
+                for cand in right_nums:
+                    if cand <= 1.0 or cand >= float(amt):
+                        continue
+                    ratio = float(amt) / float(cand)
+                    r = round(ratio)
+                    if 1 <= r <= 100 and abs(ratio - r) <= max(0.03 * r, 0.15):
+                        rate = float(cand)
+                        qty = float(r)
+                        break
 
-            possible = sorted(list({v for v in possible}), reverse=True)
+            if qty is None:
+                qty = 1.0
+            if rate is None:
+                rate = 0.0
 
-            for p in possible:
-                if p <= 1 or p >= amt:
-                    continue
-                ratio = amt / p
-                r = round(ratio)
-                if 1 <= r <= 10:
-                    rate = p
-                    qty = r
-                    break
-
-            parsed.append({
+            parsed_items.append({
                 "item_name": clean_name_text(name),
                 "item_amount": float(round(amt, 2)),
                 "item_rate": float(round(rate, 2)),
-                "item_quantity": float(qty)
+                "item_quantity": float(qty),
             })
 
-    return parsed
-
-
-###############################################
-#              FINAL ITEM FILTER
-###############################################
-
-def final_item_filter(item, headers, all_names):
-    name = item["item_name"].strip()
-    ln = name.lower()
-
-    if not name:
-        return False
-
-    for h in headers:
-        if h in ln:
-            return False
-
-    if FOOTER_KEYWORDS.search(ln):
-        return False
-
-    if item["item_amount"] <= 0:
-        return False
-
-    words = ln.split()
-    short = len(words) <= 3
-
-    if any(k in ln for k in ["charges", "services", "room", "radiology", "surgery"]) and short:
-        lower_other = " ".join(all_names).lower()
-        if any(z in lower_other for z in [
-            "rent","ward","nursing","surgeon","anaes","ot","procedure"
-        ]):
-            return False
-
-    rate = item["item_rate"]
-    amt  = item["item_amount"]
-    if rate and rate > amt * 10 and amt < 10000:
-        return False
-
-    return True
-
+    return parsed_items
 
-###############################################
-#           POST VALIDATION (PATCH)
-###############################################
-
-def post_validate_items(items):
+# ---------------- dedupe & totals ----------------
+def dedupe_items(items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    seen = set()
     out = []
     for it in items:
-        amt = it["item_amount"]
-        rate = it["item_rate"]
-        qty = it["item_quantity"]
-
-        if amt == 0 and rate > 0:
-            amt = rate * qty
-
-        if rate > 0:
-            ideal = rate * qty
-            if abs(ideal - amt) > max(2, 0.15 * ideal):
-                q = amt / rate
-                if 1 <= round(q) <= 10:
-                    qty = round(q)
-
-        it["item_amount"] = round(amt, 2)
-        it["item_rate"] = round(rate, 2)
-        it["item_quantity"] = float(qty)
-
+        nm = re.sub(r"\s+", " ", it["item_name"].lower()).strip()
+        key = (nm[:120], round(float(it["item_amount"]), 2))
+        if key in seen:
+            continue
+        seen.add(key)
         out.append(it)
     return out
 
-
-###############################################
-#          SUBTOTAL / FINAL TOTAL DETECTION
-###############################################
-
-def detect_subtotals_and_totals(rows):
-    sub = None
-    final = None
-
-    for rt in rows[::-1]:
-        if not rt.strip():
+def detect_subtotals_and_totals(rows_texts: List[str]) -> Dict[str, Optional[float]]:
+    subtotal = None; final = None
+    for rt in rows_texts[::-1]:
+        if not rt or rt.strip() == "":
             continue
-
         if TOTAL_KEYWORDS.search(rt):
             m = NUM_RE.search(rt)
             if m:
                 v = normalize_num_str(m.group(0))
                 if v is None:
                     continue
-
-                if "sub" in rt.lower():
-                    if sub is None:
-                        sub = round(v, 2)
+                if re.search(r"sub", rt, re.I):
+                    if subtotal is None: subtotal = float(round(v, 2))
                 else:
-                    if final is None:
-                        final = round(v, 2)
-
-    return {"subtotal": sub, "final_total": final}
-
+                    if final is None: final = float(round(v, 2))
+    return {"subtotal": subtotal, "final_total": final}
 
-###############################################
-#          GEMINI REFINER (PATCHED PROMPT)
-###############################################
-
-def refine_with_gemini(items, page_text=""):
+# ---------------- Gemini refinement (deterministic) ----------------
+def refine_with_gemini(page_items: List[Dict[str, Any]], page_text: str = "") -> Tuple[List[Dict[str, Any]], Dict[str, int]]:
+    zero_usage = {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
     if not GEMINI_API_KEY or genai is None:
-        return items, {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
-
-    safe = sanitize_ocr_text(page_text)
-
-    system_prompt = (
-        "You are a strict hospital bill item cleaner.\n"
-        "Return ONLY a JSON array of cleaned line items.\n"
-        "Do NOT include section headers, totals, subtotals, page numbers.\n"
-        "Do NOT invent items.\n"
-    )
-
-    user_prompt = f"""
-Extract ONLY valid line items from the bill.
-
-RULES YOU MUST FOLLOW:
-- Do NOT create new items.
-- Do NOT output section headers (Room Charges, Lab Services, Radiology).
-- Merge broken names (doctor names on multiple lines).
-- Use exact item names from OCR text.
-- Recompute rate/qty if amount = rate×qty is clear.
-- Ignore totals or summary lines.
-- Ignore page numbers.
-- Always output: item_name, item_amount, item_rate, item_quantity.
-
-OCR TEXT:
-{safe}
-
-INITIAL ITEMS:
-{json.dumps(items, ensure_ascii=False)}
-
-Return ONLY a JSON array:
-[
-  {{"item_name":"...","item_amount":float,"item_rate":float,"item_quantity":float}}
-]
-"""
-
+        return page_items, zero_usage
     try:
+        safe_text = sanitize_ocr_text(page_text)
+        system_prompt = (
+            "You are a strict bill-extraction cleaner. Return ONLY a JSON array (no explanation, no backticks). "
+            "Each entry must be an object with keys: item_name (string), item_amount (float), item_rate (float), item_quantity (float). "
+            "Do NOT include subtotal or total lines as items. Do not invent items; only clean/fix/normalize the given items."
+        )
+        user_prompt = (
+            f"page_text='''{safe_text}'''\n"
+            f"items = {json.dumps(page_items, ensure_ascii=False)}\n\n"
+            "Example:\n"
+            "items = [{'item_name':'Consultation Charge | DR PREETHI','item_amount':300.0,'item_rate':0.0,'item_quantity':300.0},\n"
+            "         {'item_name':'Description Qty / Hrs Consultation Rate Discount Net Amt','item_amount':1950.0,'item_rate':1950.0,'item_quantity':1.0}]\n"
+            "=>\n"
+            "[{'item_name':'Consultation Charge | DR PREETHI MARY JOSEPH','item_amount':300.0,'item_rate':300.0,'item_quantity':1.0}]\n\n"
+            "Return only the cleaned JSON array of items."
+        )
         model = genai.GenerativeModel(GEMINI_MODEL_NAME)
-        resp = model.generate_content(
+        response = model.generate_content(
             [
                 {"role": "system", "parts": [system_prompt]},
                 {"role": "user", "parts": [user_prompt]},
             ],
             temperature=0.0,
-            max_output_tokens=1200,
+            max_output_tokens=1000,
         )
-
-        raw = resp.text.strip()
-        raw = raw.replace("```json", "").replace("```", "").strip()
+        raw = response.text.strip()
+        if raw.startswith("```"):
+            raw = re.sub(r"^```[a-zA-Z]*", "", raw)
+            raw = re.sub(r"```$", "", raw).strip()
         parsed = json.loads(raw)
+        if isinstance(parsed, list):
+            cleaned = []
+            for obj in parsed:
+                try:
+                    cleaned.append({
+                        "item_name": str(obj.get("item_name", "")).strip(),
+                        "item_amount": float(obj.get("item_amount", 0.0)),
+                        "item_rate": float(obj.get("item_rate", 0.0) or 0.0),
+                        "item_quantity": float(obj.get("item_quantity", 1.0) or 1.0),
+                    })
+                except Exception:
+                    continue
+            return cleaned, zero_usage
+        return page_items, zero_usage
+    except Exception:
+        return page_items, zero_usage
 
-        cleaned = []
-        for obj in parsed:
-            cleaned.append({
-                "item_name": str(obj.get("item_name", "")).strip(),
-                "item_amount": float(obj.get("item_amount", 0.0)),
-                "item_rate": float(obj.get("item_rate", 0.0)),
-                "item_quantity": float(obj.get("item_quantity", 1.0)),
-            })
-
-        return cleaned, {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
-
-    except:
-        return items, {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
-
+# ---------------- header heuristics & final filter ----------------
+def looks_like_header_text(txt: str, top_of_page: bool = False) -> bool:
+    if not txt:
+        return False
+    t = re.sub(r"\s+", " ", txt.strip().lower())
+    # exact phrase blacklist
+    if any(h == t for h in HEADER_PHRASES):
+        return True
+    hits = sum(1 for k in HEADER_KEYWORDS if k in t)
+    if hits >= 2:
+        return True
+    tokens = re.split(r"[\s\|,/:]+", t)
+    key_hit_count = sum(1 for tok in tokens if tok in HEADER_KEYWORDS)
+    if key_hit_count >= 3:
+        return True
+    if top_of_page and len(tokens) <= 10 and key_hit_count >= 2:
+        return True
+    if ("rate" in t or "net" in t) and "amt" in t and not any(ch.isdigit() for ch in t):
+        return True
+    if t.startswith("description") or t.startswith("qty") or t.startswith("qty /"):
+        return True
+    return False
 
-###############################################
-#              MAIN EXTRACTION ENDPOINT
-###############################################
+def final_item_filter(item: Dict[str, Any], known_page_headers: List[str] = [], other_item_names: List[str] = []) -> bool:
+    name = (item.get("item_name") or "").strip()
+    if not name:
+        return False
+    ln = name.lower()
+    # header exact detection
+    for h in known_page_headers:
+        if h and h.strip() and h.strip().lower() in ln:
+            return False
+    if FOOTER_KEYWORDS.search(ln):
+        return False
+    if item.get("item_amount", 0) > 1_000_000:
+        return False
+    if len(name) <= 2 and not re.search(r"[a-zA-Z]", name):
+        return False
+    # avoid pure section headers (short & header words)
+    words = ln.split()
+    header_word_hits = sum(1 for k in HEADER_KEYWORDS if k in ln)
+    if header_word_hits >= 1 and len(words) <= 3:
+        # if page contains more detailed items with 'room'/'rent'/'nursing' etc, remove this generic header
+        lower_other = " ".join(other_item_names).lower()
+        if any(k in lower_other for k in ["room", "rent", "nursing", "ward", "surgeon", "anaes", "ot", "charges", "procedure", "radiology"]):
+            return False
+        # also if name is exactly one of the short header words, drop
+        if ln in ("charge", "charges", "services", "consultation", "room", "radiology", "surgery"):
+            return False
+    # drop non-informative labels even if they have amount (summary rows)
+    if len(words) <= 4 and re.search(r"\b(charges|services|room|radiolog|laborat|surgery|procedure|rent|nursing)\b", ln):
+        # try to detect if it's a summary (presence of other more specific items)
+        lower_other = " ".join(other_item_names).lower()
+        if any(tok in lower_other for tok in ["rent", "room", "ward", "nursing", "surgeon", "anaes", "ot"]):
+            return False
+    if float(item.get("item_amount", 0)) <= 0.0:
+        return False
+    # sanity check rate vs amount
+    rate = float(item.get("item_rate", 0) or 0)
+    amt = float(item.get("item_amount", 0) or 0)
+    if rate and rate > amt * 10 and amt < 10000:
+        return False
+    return True
 
+# ---------------- main endpoint ----------------
 @app.post("/extract-bill-data")
 async def extract_bill_data(payload: BillRequest):
-
-    url = payload.document
-
-    # download
+    doc_url = payload.document
     try:
-        r = requests.get(url, headers={"User-Agent": "Mozilla"}, timeout=30)
-        if r.status_code != 200:
-            raise RuntimeError("Download failed")
-        data = r.content
-    except:
-        return {
-            "is_success": False,
-            "token_usage": {},
-            "data": {"pagewise_line_items": [], "total_item_count": 0}
-        }
-
-    # load image(s)
+        headers = {"User-Agent": "Mozilla/5.0"}
+        resp = requests.get(doc_url, headers=headers, timeout=30)
+        if resp.status_code != 200:
+            raise RuntimeError(f"download failed status={resp.status_code}")
+        file_bytes = resp.content
+    except Exception:
+        return {"is_success": False, "token_usage": {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}, "data": {"pagewise_line_items": [], "total_item_count": 0}}
+
+    images = []
+    clean_url = doc_url.split("?", 1)[0].lower()
     try:
-        if url.lower().split("?")[0].endswith(".pdf"):
-            imgs = convert_from_bytes(data)
+        if clean_url.endswith(".pdf"):
+            images = convert_from_bytes(file_bytes)
+        elif any(clean_url.endswith(ext) for ext in [".png", ".jpg", ".jpeg", ".tiff", ".bmp"]):
+            images = [Image.open(BytesIO(file_bytes))]
         else:
-            imgs = [Image.open(BytesIO(data))]
-    except:
-        imgs = []
+            try:
+                images = convert_from_bytes(file_bytes)
+            except Exception:
+                images = []
+    except Exception:
+        images = []
 
     pagewise = []
-    total_tokens = {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
-
-    for idx, img in enumerate(imgs, 1):
+    cumulative_token_usage = {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
 
+    for idx, page_img in enumerate(images, start=1):
         try:
-            proc = preprocess_image(img)
+            proc = preprocess_image(page_img)
             cells = image_to_tsv_cells(proc)
-            rows = group_cells_into_rows(cells)
-
-            row_texts = [" ".join([c["text"] for c in r]).strip() for r in rows]
-
-            # remove headers
-            filtered = []
-            for i, (r, t) in enumerate(zip(rows, row_texts)):
-                if looks_like_header_text(t, top_of_page=(i < 5)):
+            rows = group_cells_into_rows(cells, y_tolerance=12)
+            rows_texts = [" ".join([c["text"] for c in r]).strip() for r in rows]
+
+            # === HEADER PREFILTER: remove header-like rows anywhere on page ===
+            rows_filtered = []
+            for i, (r, rt) in enumerate(zip(rows, rows_texts)):
+                top_flag = (i < 6)
+                rt_norm = sanitize_ocr_text(rt).lower()
+                if looks_like_header_text(rt_norm, top_of_page=top_flag):
                     continue
-                filtered.append(r)
-
-            rows = filtered
-            row_texts = [" ".join([c["text"] for c in r]).strip() for r in rows]
-            page_text = " ".join(row_texts)
+                if any(h in rt_norm for h in HEADER_PHRASES):
+                    continue
+                rows_filtered.append(r)
+            # recompute row texts and a simple page_text
+            rows = rows_filtered
+            rows_texts = [" ".join([c["text"] for c in r]).strip() for r in rows]
+            page_text = sanitize_ocr_text(" ".join(rows_texts))
 
+            # detect page-level top headers (for final filtering)
             top_headers = []
-            for t in row_texts[:5]:
-                if looks_like_header_text(t, top_of_page=True):
-                    top_headers.append(t.lower())
+            for i, rt in enumerate(rows_texts[:6]):
+                if looks_like_header_text(rt, top_of_page=(i < 4)):
+                    top_headers.append(rt.strip().lower())
 
             parsed_items = parse_rows_with_columns(rows, cells)
 
-            refined, usage = refine_with_gemini(parsed_items, page_text)
-
-            for k in total_tokens:
-                total_tokens[k] += usage.get(k, 0)
-
-            all_names = [x["item_name"] for x in refined]
+            # ALWAYS attempt Gemini refinement if available (deterministic settings)
+            refined_items, token_u = refine_with_gemini(parsed_items, page_text)
+            for k in cumulative_token_usage:
+                cumulative_token_usage[k] += token_u.get(k, 0)
 
-            cleaned = [
-                x for x in refined
-                if final_item_filter(x, top_headers, all_names)
-            ]
+            # Prepare other_item_names for contextual filtering (helps remove generic section headers)
+            other_item_names = [it.get("item_name","") for it in refined_items]
 
-            cleaned = post_validate_items(cleaned)
-
-            totals = detect_subtotals_and_totals(row_texts)
+            # final cleaning & dedupe
+            cleaned = [p for p in refined_items if final_item_filter(p, known_page_headers=top_headers, other_item_names=other_item_names)]
+            cleaned = dedupe_items(cleaned)
+            cleaned = [p for p in cleaned if not looks_like_header_text(p["item_name"].lower())]
 
             page_type = "Bill Detail"
-            low = page_text.lower()
-            if "pharmacy" in low:
+            page_txt = page_text.lower()
+            if any(x in page_txt for x in ["pharmacy", "medicine", "tablet"]):
                 page_type = "Pharmacy"
-            if "final bill" in low or "grand total" in low:
+            if "final bill" in page_txt or "grand total" in page_txt:
                 page_type = "Final Bill"
 
-            pagewise.append({
-                "page_no": str(idx),
-                "page_type": page_type,
-                "bill_items": cleaned,
-                "subtotal": totals["subtotal"],
-                "final_page_total": totals["final_total"]
-            })
-
-        except:
-            pagewise.append({
-                "page_no": str(idx),
-                "page_type": "Bill Detail",
-                "bill_items": [],
-                "subtotal": None,
-                "final_page_total": None
-            })
-
-    # global final total = sum of all item amounts
-    final_sum = 0.0
-    for p in pagewise:
-        for it in p["bill_items"]:
-            final_sum += it["item_amount"]
-
-    total_item_count = sum(len(p["bill_items"]) for p in pagewise)
-
-    return {
-        "is_success": True,
-        "token_usage": total_tokens,
-        "data": {
-            "pagewise_line_items": pagewise,
-            "total_item_count": total_item_count,
-            "final_total": round(final_sum, 2)
-        }
-    }
+            pagewise.append({"page_no": str(idx), "page_type": page_type, "bill_items": cleaned})
+        except Exception:
+            pagewise.append({"page_no": str(idx), "page_type": "Bill Detail", "bill_items": []})
+            continue
 
+    total_item_count = sum(len(p.get("bill_items", [])) for p in pagewise)
+    if not GEMINI_API_KEY or genai is None:
+        cumulative_token_usage["warning_no_gemini"] = 1
 
-###############################################
-#                  DEBUG ENDPOINT
-###############################################
+    return {"is_success": True, "token_usage": cumulative_token_usage, "data": {"pagewise_line_items": pagewise, "total_item_count": total_item_count}}
 
+# ---------------- debug TSV ----------------
 @app.post("/debug-tsv")
 async def debug_tsv(payload: BillRequest):
+    doc_url = payload.document
     try:
-        r = requests.get(payload.document, timeout=20)
-        img = Image.open(BytesIO(r.content))
-        proc = preprocess_image(img)
-        cells = image_to_tsv_cells(proc)
-        return {"cells": cells}
-    except:
-        return {"error": "debug failed"}
-
-
-###############################################
-#                HEALTH CHECK
-###############################################
+        resp = requests.get(doc_url, timeout=20)
+        if resp.status_code != 200:
+            return {"error": "Download failed"}
+        file_bytes = resp.content
+    except Exception:
+        return {"error": "Download failed"}
+    clean_url = doc_url.split("?", 1)[0].lower()
+    if clean_url.endswith(".pdf"):
+        imgs = convert_from_bytes(file_bytes)
+        img = imgs[0]
+    else:
+        img = Image.open(BytesIO(file_bytes))
+    proc = preprocess_image(img)
+    cells = image_to_tsv_cells(proc)
+    return {"cells": cells}
 
 @app.get("/")
-def ping():
-    msg = "Bill extractor live."
-    if not GEMINI_API_KEY:
-        msg += " (Gemini missing)"
-    return {"status": "ok", "message": msg}
+def health_check():
+    msg = "Bill extraction API (final) live."
+    if not GEMINI_API_KEY or genai is None:
+        msg += " (No GEMINI_API_KEY/configured SDK — LLM refinement skipped.)"
+    return {"status": "ok", "message": msg, "hint": "POST /extract-bill-data with {'document':'<url>'}"}