Update app.py

app.py

@@ -2,10 +2,20 @@
 ValuationAI — Nairobi Valuation Sheet OCR
 Model: rasmodev/Handwriting_trocr_model
 
+PDF processing matches notebook exactly:
+- fitz opened via temp file (not stream) matching how training data was built
+- Matrix(200/72, 200/72) — same DPI as training
+- get_pixmap(matrix=mat, alpha=False) — same as training
+- Image.open(...).convert('RGB') — same as training
+
+Inference matches notebook exactly:
+- processor(images=img.convert('RGB'), return_tensors='pt').pixel_values
+- model.generate(pixel_values=pv, max_new_tokens=64, num_beams=1)
+
 Label format from training:
-PLOT: LR 209/617 | LOC: STATE HOUSE AVENUE | AREA: 0.06 | AMT: 52000000 | DATE: 2008-06-17 | VOS: 3872
+- PLOT: LR 209/617 | LOC: STATE HOUSE AVENUE | AREA: 0.06 | AMT: 52000000 | DATE: 2008-06-17 | VOS: 3872
 """
-import io, time, logging
+import io, time, logging, tempfile, os
 import streamlit as st
 import pandas as pd
 from PIL import Image
@@ -33,7 +43,6 @@ html, body, [class*="css"], .stApp {
 }
 #MainMenu, footer, header { visibility: hidden; }
 
-/* Top bar */
 .topbar {
     display: flex; align-items: flex-end;
     justify-content: space-between;
@@ -46,15 +55,12 @@ html, body, [class*="css"], .stApp {
 .model-ref { font-size: 0.7rem; color: #9CA3AF; font-weight: 400; letter-spacing: 0.04em; text-align: right; }
 .model-ref strong { color: #2563EB; font-weight: 600; }
 
-/* Headline */
 .headline { font-family: 'Cormorant Garamond', serif; font-size: 3.4rem; font-weight: 700; line-height: 1.08; letter-spacing: -0.03em; color: #1A1A2E; margin-bottom: 1rem; max-width: 700px; }
 .headline em { font-style: italic; color: #2563EB; }
 .subline { font-size: 0.95rem; font-weight: 300; color: #6B7280; line-height: 1.7; max-width: 500px; margin-bottom: 3rem; }
 
-/* Step label */
 .step { font-size: 0.65rem; font-weight: 700; letter-spacing: 0.18em; text-transform: uppercase; color: #2563EB; margin-bottom: 0.5rem; }
 
-/* File uploader */
 [data-testid="stFileUploader"] section {
     background: #fff !important;
     border: 2px dashed #D1D5DB !important;
@@ -69,10 +75,8 @@ html, body, [class*="css"], .stApp {
 }
 [data-testid="stFileUploader"] label { color: #6B7280 !important; font-size: 0.9rem !important; }
 
-/* File chip */
 .fchip { display: inline-flex; align-items: center; gap: 5px; background: #EFF6FF; border: 1px solid #BFDBFE; color: #1D4ED8; padding: 0.25rem 0.7rem; border-radius: 6px; font-size: 0.73rem; font-weight: 500; margin: 2px; }
 
-/* Button */
 .stButton > button {
     background: #1A1A2E !important; color: #fff !important; border: none !important;
     border-radius: 8px !important; padding: 0.85rem 2.5rem !important;
@@ -84,22 +88,18 @@ html, body, [class*="css"], .stApp {
 .stButton > button:hover { background: #2563EB !important; box-shadow: 0 4px 16px rgba(37,99,235,0.3) !important; transform: translateY(-1px) !important; }
 .stButton > button:disabled { background: #E5E7EB !important; color: #9CA3AF !important; box-shadow: none !important; transform: none !important; }
 
-/* Progress */
 .stProgress > div > div > div { background: #2563EB !important; border-radius: 4px !important; }
 .stProgress > div > div { background: #E5E7EB !important; border-radius: 4px !important; height: 4px !important; }
 
-/* Stats */
 .stats-strip { display: flex; background: #1A1A2E; border-radius: 12px; overflow: hidden; margin: 2.5rem 0 2rem; }
 .stat-item { flex: 1; padding: 1.6rem 2rem; border-right: 1px solid rgba(255,255,255,0.08); }
 .stat-item:last-child { border-right: none; }
 .stat-n { font-family: 'Cormorant Garamond', serif; font-size: 2.6rem; font-weight: 700; color: #fff; line-height: 1; margin-bottom: 0.3rem; }
 .stat-l { font-size: 0.68rem; font-weight: 500; letter-spacing: 0.12em; text-transform: uppercase; color: #6B7280; }
 
-/* Section head */
 .section-head { display: flex; align-items: center; justify-content: space-between; margin-bottom: 1rem; padding-bottom: 0.75rem; border-bottom: 1px solid #E5E7EB; }
 .section-title { font-family: 'Cormorant Garamond', serif; font-size: 1.5rem; font-weight: 600; color: #1A1A2E; letter-spacing: -0.01em; }
 
-/* Download button */
 div[data-testid="stDownloadButton"] > button {
     background: #fff !important; border: 1.5px solid #1A1A2E !important; color: #1A1A2E !important;
     border-radius: 8px !important; padding: 0.6rem 1.4rem !important;
@@ -109,14 +109,13 @@ div[data-testid="stDownloadButton"] > button {
 }
 div[data-testid="stDownloadButton"] > button:hover { background: #1A1A2E !important; color: #fff !important; }
 
-/* Dataframe */
 [data-testid="stDataFrame"] { border-radius: 10px !important; border: 1px solid #E5E7EB !important; overflow: hidden !important; box-shadow: 0 1px 4px rgba(0,0,0,0.05) !important; }
 </style>
 """, unsafe_allow_html=True)
 
 
 # ═══════════════════════════════════════════════════════════
-# MODEL LOADING
+# MODEL — matches notebook Cell 13 + Cell 28
 # ═══════════════════════════════════════════════════════════
 @st.cache_resource(show_spinner="Loading recognition model…")
 def load_model():
@@ -131,28 +130,59 @@ def load_model():
 
 
 # ═══════════════════════════════════════════════════════════
-# OCR — matches training output format exactly
-# Training label format:
-# PLOT: LR 209/617 | LOC: STATE HOUSE AVENUE | AREA: 0.06 | AMT: 52000000 | DATE: 2008-06-17 | VOS: 3872
+# PDF → IMAGES — matches notebook Cell 10 exactly
+# Uses temp file not stream — same as training
+# Matrix(200/72, 200/72), get_pixmap(alpha=False), convert('RGB')
+# ═══════════════════════════════════════════════════════════
+def pdf_to_images(file_bytes: bytes) -> list:
+    import fitz
+    images = []
+    # Write to temp file — same as training which used file paths
+    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
+        tmp.write(file_bytes)
+        tmp_path = tmp.name
+    try:
+        doc = fitz.open(tmp_path)                      # open from path like training
+        mat = fitz.Matrix(200/72, 200/72)              # same DPI as training
+        for page in doc:
+            pix = page.get_pixmap(matrix=mat, alpha=False)  # same as training
+            img = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB")  # same as training
+            images.append(img)
+            pix = None                                 # free memory immediately like training
+        doc.close()
+    finally:
+        os.unlink(tmp_path)
+    return images
+
+
+# ═══════════════════════════════════════════════════════════
+# OCR — matches notebook Cell 18 + Cell 20 inference
+# processor(images=img.convert('RGB')) then model.generate
+# max_new_tokens=64, num_beams=1 (greedy — fast)
 # ═══════════════════════════════════════════════════════════
 def ocr_page(img: Image.Image) -> str:
-    """Run the fine-tuned model on one page image."""
     import torch
     processor, model, device = load_model()
-    pv = processor(
+    # Exactly as in ValuationDataset.__getitem__
+    pixel_values = processor(
         images=img.convert("RGB"),
         return_tensors="pt"
     ).pixel_values.to(device)
+
     with torch.no_grad():
-        gen = model.generate(pv, max_new_tokens=128)
-    return processor.batch_decode(gen, skip_special_tokens=True)[0].strip()
+        generated = model.generate(
+            pixel_values=pixel_values,
+            max_new_tokens=64,
+            num_beams=1,            # greedy — fast, matches validation in notebook
+        )
+    return processor.batch_decode(generated, skip_special_tokens=True)[0].strip()
 
 
+# ═══════════════════════════════════════════════════════════
+# PARSE LABEL — matches row_to_label() from notebook Cell 10
+# Format: PLOT: ... | LOC: ... | AREA: ... | AMT: ... | DATE: ... | VOS: ...
+# ═══════════════════════════════════════════════════════════
 def parse_label(raw_text: str, filename: str) -> dict:
-    """
-    Parse the pipe-delimited label that the model was trained to output.
-    Format: PLOT: ... | LOC: ... | AREA: ... | AMT: ... | DATE: ... | VOS: ...
-    """
     record = {
         "File":         filename,
         "Plot Number":  "",
@@ -163,17 +193,13 @@ def parse_label(raw_text: str, filename: str) -> dict:
         "VOS":          "",
         "Raw Output":   raw_text,
     }
-
-    # Split on pipe delimiter
-    parts = raw_text.split("|")
-    for part in parts:
+    for part in raw_text.split("|"):
         part = part.strip()
         if ":" not in part:
             continue
         key, _, val = part.partition(":")
         key = key.strip().upper()
         val = val.strip()
-
         if key == "PLOT":
             record["Plot Number"] = val
         elif key == "LOC":
@@ -181,7 +207,6 @@ def parse_label(raw_text: str, filename: str) -> dict:
         elif key == "AREA":
             record["Area"] = val
         elif key == "AMT":
-            # Remove commas and convert to int
             try:
                 record["Amount (KES)"] = int(val.replace(",", "").replace(" ", ""))
             except ValueError:
@@ -190,54 +215,34 @@ def parse_label(raw_text: str, filename: str) -> dict:
             record["Date"] = val
         elif key == "VOS":
             record["VOS"] = val
-
     return record
 
 
-def pdf_to_images(file_bytes: bytes) -> list:
-    """Convert all pages of a PDF to PIL images at 200 DPI."""
-    import fitz
-    doc  = fitz.open(stream=file_bytes, filetype="pdf")
-    mat  = fitz.Matrix(200/72, 200/72)
-    imgs = []
-    for i in range(len(doc)):
-        pix = doc[i].get_pixmap(matrix=mat, alpha=False)
-        imgs.append(Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB"))
-    doc.close()
-    return imgs
-
-
+# ═══════════════════════════════════════════════════════════
+# EXCEL EXPORT
+# ═══════════════════════════════════════════════════════════
 def make_excel(records: list) -> bytes:
-    """Export records to a formatted Excel workbook."""
     from openpyxl import load_workbook
     from openpyxl.styles import Font, PatternFill, Alignment
     from openpyxl.utils import get_column_letter
-
-    # Drop raw output from Excel — it's only for debugging
     clean = [{k: v for k, v in r.items() if k != "Raw Output"} for r in records]
-
     buf = io.BytesIO()
     pd.DataFrame(clean).to_excel(buf, index=False, sheet_name="Valuation Data")
     buf.seek(0)
-
     wb  = load_workbook(buf)
     ws  = wb.active
     hdr = PatternFill("solid", start_color="1A1A2E")
-
     for ci, cell in enumerate(ws[1], 1):
         cell.font      = Font(name="Calibri", bold=True, color="FFFFFF", size=11)
         cell.fill      = hdr
         cell.alignment = Alignment(horizontal="center", vertical="center")
         ws.column_dimensions[get_column_letter(ci)].width = 26
-
     ws.row_dimensions[1].height = 30
-
     for row in ws.iter_rows(min_row=2):
         for cell in row:
             cell.alignment = Alignment(vertical="center", wrap_text=True)
             if cell.row % 2 == 0:
                 cell.fill = PatternFill("solid", start_color="F0F4FF")
-
     ws.freeze_panes = "A2"
     out = io.BytesIO()
     wb.save(out)
@@ -253,7 +258,7 @@ for k, v in [("records",[]),("excel",None),("done",False),("errors",[])]:
 
 
 # ═══════════════════════════════════════════════════════════
-# UI — TOP BAR
+# UI
 # ═══════════════════════════════════════════════════════════
 st.markdown("""
 <div class="topbar">
@@ -268,10 +273,6 @@ st.markdown("""
 </div>
 """, unsafe_allow_html=True)
 
-
-# ═══════════════════════════════════════════════════════════
-# UI — HEADLINE
-# ═══════════════════════════════════════════════════════════
 st.markdown("""
 <div class="headline">
   Digitise handwritten<br>valuation sheets <em>instantly.</em>
@@ -283,10 +284,6 @@ st.markdown("""
 </div>
 """, unsafe_allow_html=True)
 
-
-# ═══════════════════════════════════════════════════════════
-# UI — UPLOAD
-# ═══════════════════════════════════════════════════════════
 st.markdown('<div class="step">Step 1 — Upload Documents</div>', unsafe_allow_html=True)
 
 uploaded = st.file_uploader(
@@ -333,17 +330,18 @@ if run and uploaded:
         try:
             ext = fname.lower().rsplit(".", 1)[-1]
 
-            # Get page images
             if ext == "pdf":
                 imgs = pdf_to_images(raw)
             else:
                 imgs = [Image.open(io.BytesIO(raw)).convert("RGB")]
 
-            # Run OCR on each page
+            if not imgs:
+                st.session_state.errors.append(f"{fname}: no pages could be extracted")
+                continue
+
             for pi, img in enumerate(imgs, 1):
                 status.caption(
-                    f"Running recognition on **{fname}** — "
-                    f"page {pi} of {len(imgs)}"
+                    f"Processing **{fname}** — page {pi} of {len(imgs)}"
                 )
                 raw_text = ocr_page(img)
                 record   = parse_label(raw_text, fname)
@@ -374,45 +372,25 @@ if run and uploaded:
 if st.session_state.done and st.session_state.records:
     records = st.session_state.records
     df      = pd.DataFrame(records)
-
-    # Display columns — exclude raw output from table
     display_cols = [c for c in df.columns if c != "Raw Output"]
     df_display   = df[display_cols]
 
-    # Stats
     n_plots   = df["Plot Number"].astype(bool).sum()
     n_amounts = pd.to_numeric(df["Amount (KES)"], errors="coerce").notna().sum()
     n_dates   = df["Date"].astype(bool).sum()
 
     st.markdown(f"""
     <div class="stats-strip">
-      <div class="stat-item">
-        <div class="stat-n">{len(records)}</div>
-        <div class="stat-l">Pages processed</div>
-      </div>
-      <div class="stat-item">
-        <div class="stat-n">{n_plots}</div>
-        <div class="stat-l">Plot numbers</div>
-      </div>
-      <div class="stat-item">
-        <div class="stat-n">{n_amounts}</div>
-        <div class="stat-l">Amounts extracted</div>
-      </div>
-      <div class="stat-item">
-        <div class="stat-n">{n_dates}</div>
-        <div class="stat-l">Dates captured</div>
-      </div>
+      <div class="stat-item"><div class="stat-n">{len(records)}</div><div class="stat-l">Pages processed</div></div>
+      <div class="stat-item"><div class="stat-n">{n_plots}</div><div class="stat-l">Plot numbers</div></div>
+      <div class="stat-item"><div class="stat-n">{n_amounts}</div><div class="stat-l">Amounts extracted</div></div>
+      <div class="stat-item"><div class="stat-n">{n_dates}</div><div class="stat-l">Dates captured</div></div>
     </div>
     """, unsafe_allow_html=True)
 
-    # Table header + download side by side
     col_t, col_d = st.columns([5, 1])
     with col_t:
-        st.markdown("""
-        <div class="section-head">
-          <div class="section-title">Extracted Records</div>
-        </div>
-        """, unsafe_allow_html=True)
+        st.markdown('<div class="section-head"><div class="section-title">Extracted Records</div></div>', unsafe_allow_html=True)
     with col_d:
         st.markdown('<div style="padding-top:0.3rem"></div>', unsafe_allow_html=True)
         if st.session_state.excel:
@@ -423,16 +401,10 @@ if st.session_state.done and st.session_state.records:
                 mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
             )
 
-    # Dataframe
-    st.dataframe(
-        df_display,
-        use_container_width=True,
-        height=min(80 + len(df) * 38, 560),
-        hide_index=True,
-    )
+    st.dataframe(df_display, use_container_width=True,
+                 height=min(80 + len(df)*38, 560), hide_index=True)
 
-    # Raw OCR output — for verification/debugging
-    with st.expander("🔍 View raw model output (for verification)"):
+    with st.expander("🔍 Raw model output (for verification)"):
         for r in records:
             st.markdown(
                 f'<div style="font-family:monospace;font-size:0.78rem;'
@@ -441,8 +413,7 @@ if st.session_state.done and st.session_state.records:
                 unsafe_allow_html=True,
             )
 
-    # Errors
     if st.session_state.errors:
-        with st.expander(f"⚠  {len(st.session_state.errors)} file(s) could not be processed"):
+        with st.expander(f"⚠ {len(st.session_state.errors)} file(s) could not be processed"):
             for e in st.session_state.errors:
-                st.caption(e)
+                st.caption(e)