Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,9 +7,9 @@ from PIL import Image, ImageEnhance
 from streamlit_drawable_canvas import st_canvas
 import pytesseract
 
-# Tesseract is installed via packages.txt on HuggingFace Spaces (Linux)
-# No need to set path - it's in system PATH
-
+# ---------------------------------
+# Page config
+# ---------------------------------
 st.set_page_config(
     page_title="Remittance GT Annotator - Interactive OCR",
     layout="wide"
@@ -17,7 +17,9 @@ st.set_page_config(
 
 st.title("Remittance GT Annotator - Interactive OCR")
 
-# ---- Define fields ----
+# ---------------------------------
+# Field definitions
+# ---------------------------------
 SINGLE_FIELDS = [
     "Remittance Advice Number",
     "Remittance Advice Date",
@@ -63,7 +65,9 @@ COLOR_PALETTE = [
 ALL_BASE_FIELDS = SINGLE_FIELDS + LINE_ITEM_FIELDS
 FIELD_COLORS = {field: COLOR_PALETTE[i % len(COLOR_PALETTE)] for i, field in enumerate(ALL_BASE_FIELDS)}
 
-# ----- JSONL schema helper mappings -----
+# ---------------------------------
+# JSONL schema mappings
+# ---------------------------------
 HEADER_GROUPS = {
     "remittance_advice_details": {
         "Remittance Advice Number": "remittance_advice_number",
@@ -106,70 +110,64 @@ LINE_ITEM_FIELD_KEY_MAP = {
 # Fixed zoom options
 ZOOM_OPTIONS = [25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80, 85, 90, 95, 100, 110, 120, 130, 140, 150]
 
-# ---- Session state ----
+# ---------------------------------
+# Session state init
+# ---------------------------------
 if "field_values" not in st.session_state:
-    st.session_state.field_values = {}
+    st.session_state.field_values = {}  # {image_name: {field_name: value}}
 if "field_rects_orig" not in st.session_state:
-    st.session_state.field_rects_orig = {}
+    st.session_state.field_rects_orig = {}  # {image_name: {field_name: rect_in_orig_coords}}
 if "num_line_items" not in st.session_state:
-    st.session_state.num_line_items = {}
+    st.session_state.num_line_items = {}  # {image_name: int}
 if "selected_image" not in st.session_state:
     st.session_state.selected_image = None
 if "zoom_values" not in st.session_state:
-    st.session_state.zoom_values = {}
+    st.session_state.zoom_values = {}  # {image_name: zoom_int}
 if "rect_version" not in st.session_state:
-    st.session_state.rect_version = {}
+    st.session_state.rect_version = {}  # {image_name: int}
 if "image_data" not in st.session_state:
-    st.session_state.image_data = {}
-
-# Pending delete - process at start before UI
+    st.session_state.image_data = {}  # {image_name: bytes}
 if "pending_delete" not in st.session_state:
     st.session_state.pending_delete = None
 
+# Process pending delete early
 if st.session_state.pending_delete is not None:
     img_name, field_key = st.session_state.pending_delete
     if img_name in st.session_state.field_rects_orig:
         st.session_state.field_rects_orig[img_name].pop(field_key, None)
+    if img_name in st.session_state.field_values:
+        st.session_state.field_values[img_name].pop(field_key, None)
     if img_name in st.session_state.rect_version:
         st.session_state.rect_version[img_name] += 1
     st.session_state.pending_delete = None
+    # Force a quick rerun so canvas reflects deletion
+    st.experimental_rerun()
 
-# --- Helper functions ---
+# ---------------------------------
+# Helper functions
+# ---------------------------------
 @st.cache_data
-def load_image(file_content):
+def load_image(file_content: bytes):
     return Image.open(BytesIO(file_content)).convert("RGB")
 
-def get_display_image_from_bytes(image_bytes, width, height):
-    """Create fresh display image from bytes - no caching to avoid stale PIL references"""
+@st.cache_data
+def get_display_image(image_bytes: bytes, width: int, height: int):
+    """Cached resize + enhancement to minimize flicker on reruns."""
     pil_image = Image.open(BytesIO(image_bytes)).convert("RGB")
     resized = pil_image.resize((width, height), Image.LANCZOS)
     resized = ImageEnhance.Sharpness(resized).enhance(1.2)
     resized = ImageEnhance.Contrast(resized).enhance(1.1)
     return resized
 
-def get_default_zoom(pil_image):
-    """Calculate best fit zoom"""
+def get_default_zoom(pil_image: Image.Image) -> int:
     MAX_WIDTH = 850
     MAX_HEIGHT = 900
     default_scale = min(MAX_WIDTH / pil_image.width, MAX_HEIGHT / pil_image.height, 1.0)
     default_zoom = int(default_scale * 100)
-    # Find closest zoom option
     closest = min(ZOOM_OPTIONS, key=lambda x: abs(x - default_zoom))
     return closest
 
 def build_gt_record_for_file(file_name: str) -> dict:
-    """
-    JSONL record for one remittance image:
-    {
-      "file_name": "<image>",
-      "gt_parse": {
-        "remittance_advice_details": {...},
-        "customer_supplier_details": {...},
-        "bank_details": {...},
-        "line_items": [...]
-      }
-    }
-    """
     values = st.session_state.field_values.get(file_name, {})
     num_items = st.session_state.num_line_items.get(file_name, 1)
 
@@ -207,16 +205,17 @@ def build_gt_record_for_file(file_name: str) -> dict:
     }
 
 def has_any_label(fname: str) -> bool:
-    """Check if file has any labeled values"""
     vals = st.session_state.field_values.get(fname, {})
     return any(str(v).strip() for v in vals.values())
 
-# --- Upload (compact) ---
+# ---------------------------------
+# Upload
+# ---------------------------------
 uploaded_files = st.file_uploader(
     "Upload remittance images",
     type=["png", "jpg", "jpeg"],
     accept_multiple_files=True,
-    label_visibility="collapsed"
+    label_visibility="collapsed",
 )
 
 if not uploaded_files:
@@ -227,7 +226,6 @@ images = []
 for f in uploaded_files:
     f.seek(0)
     content = f.read()
-    # Store image bytes in session state for stability across reruns
     if f.name not in st.session_state.image_data:
         st.session_state.image_data[f.name] = content
     img = load_image(st.session_state.image_data[f.name])
@@ -235,7 +233,6 @@ for f in uploaded_files:
 
 file_names = [img["name"] for img in images]
 
-# Image selector dropdown only (no duplicate list above)
 selected_name = st.selectbox("Select image", file_names, label_visibility="collapsed")
 st.session_state.selected_image = selected_name
 
@@ -243,31 +240,35 @@ selected_img_data = next(img for img in images if img["name"] == selected_name)
 pil_image = selected_img_data["image"]
 image_bytes = selected_img_data["bytes"]
 
-# Init for this image
+# Init per-image state
 if selected_name not in st.session_state.field_values:
     st.session_state.field_values[selected_name] = {}
 if selected_name not in st.session_state.field_rects_orig:
     st.session_state.field_rects_orig[selected_name] = {}
 if selected_name not in st.session_state.num_line_items:
     st.session_state.num_line_items[selected_name] = 1
-if selected_name not in st.session_state.rect_version:
-    st.session_state.rect_version[selected_name] = 0
 if selected_name not in st.session_state.zoom_values:
     st.session_state.zoom_values[selected_name] = get_default_zoom(pil_image)
+if selected_name not in st.session_state.rect_version:
+    st.session_state.rect_version[selected_name] = 0
 
-# ========== MAIN COLUMNS ==========
+# ---------------------------------
+# Layout columns
+# ---------------------------------
 col1, col2 = st.columns([3, 2])
 
-# Initialize field variables with defaults
+# Defaults for current field
 display_field_name = SINGLE_FIELDS[0]
 storage_field_name = SINGLE_FIELDS[0]
 base_field_for_color = SINGLE_FIELDS[0]
 field_color = FIELD_COLORS[base_field_for_color]
 
-# ====== RHS TOP: FIELD SELECTION + ZOOM ======
+# ---------------------------------
+# RHS TOP: Field selection + zoom
+# ---------------------------------
 with col2:
     st.markdown("#### 🎯 Field Selection")
-    
+
     def add_line_item():
         img = st.session_state.selected_image
         if img:
@@ -283,6 +284,7 @@ with col2:
                 st.session_state.field_rects_orig[img].pop(key, None)
             st.session_state.num_line_items[img] -= 1
             st.session_state.rect_version[img] += 1
+            st.experimental_rerun()
 
     field_type = st.radio("Type", ["Single", "Line Item"], horizontal=True, label_visibility="collapsed")
 
@@ -293,19 +295,19 @@ with col2:
         base_field_for_color = field_name
     else:
         num_items = st.session_state.num_line_items[selected_name]
-        
+
         line_col1, add_col, rem_col = st.columns([2, 1, 1])
         with line_col1:
             line_item_options = [f"Line {i+1}" for i in range(num_items)]
             selected_line_item = st.selectbox("Line", line_item_options, label_visibility="collapsed")
             line_item_num = int(selected_line_item.split()[1])
-        
+
         with add_col:
             st.button("➕", key=f"addli_{selected_name}", on_click=add_line_item, help="Add line item")
         with rem_col:
             if st.session_state.num_line_items[selected_name] > 1:
                 st.button("➖", key=f"remli_{selected_name}", on_click=remove_line_item, help="Remove line item")
-        
+
         base_field = st.selectbox("Field", LINE_ITEM_FIELDS, label_visibility="collapsed")
         display_field_name = f"{selected_line_item}: {base_field}"
         storage_field_name = f"Line {line_item_num}: {base_field}"
@@ -321,33 +323,33 @@ with col2:
     )
 
     st.markdown("#### 🔍 Zoom")
-    
+
     current_zoom = st.session_state.zoom_values[selected_name]
     zoom_index = ZOOM_OPTIONS.index(current_zoom) if current_zoom in ZOOM_OPTIONS else 0
-    
+
     def do_zoom_out():
         img = st.session_state.selected_image
         curr = st.session_state.zoom_values[img]
         idx = ZOOM_OPTIONS.index(curr) if curr in ZOOM_OPTIONS else 0
         if idx > 0:
             st.session_state.zoom_values[img] = ZOOM_OPTIONS[idx - 1]
-    
+
     def do_zoom_in():
         img = st.session_state.selected_image
         curr = st.session_state.zoom_values[img]
         idx = ZOOM_OPTIONS.index(curr) if curr in ZOOM_OPTIONS else 0
         if idx < len(ZOOM_OPTIONS) - 1:
             st.session_state.zoom_values[img] = ZOOM_OPTIONS[idx + 1]
-    
+
     def do_zoom_fit():
         img = st.session_state.selected_image
         img_bytes = st.session_state.image_data.get(img)
         if img_bytes:
             pil_img = load_image(img_bytes)
             st.session_state.zoom_values[img] = get_default_zoom(pil_img)
-    
+
     zoom_col1, zoom_col2, zoom_col3, zoom_col4 = st.columns([2, 1, 1, 1])
-    
+
     with zoom_col1:
         zoom = st.selectbox(
             "Zoom",
@@ -355,29 +357,31 @@ with col2:
             index=zoom_index,
             format_func=lambda x: f"{x}%",
             key=f"zoom_select_{selected_name}",
-            label_visibility="collapsed"
+            label_visibility="collapsed",
         )
         st.session_state.zoom_values[selected_name] = zoom
-    
+
     with zoom_col2:
-        st.button("➖", key="zoom_out", help="Zoom out", on_click=do_zoom_out)
-    
+        st.button("➖", key=f"zoom_out_{selected_name}", help="Zoom out", on_click=do_zoom_out)
+
     with zoom_col3:
-        st.button("➕", key="zoom_in", help="Zoom in", on_click=do_zoom_in)
-    
+        st.button("➕", key=f"zoom_in_{selected_name}", help="Zoom in", on_click=do_zoom_in)
+
     with zoom_col4:
-        st.button("Fit", key="zoom_fit", help="Fit to screen", on_click=do_zoom_fit)
+        st.button("Fit", key=f"zoom_fit_{selected_name}", help="Fit to screen", on_click=do_zoom_fit)
 
     st.caption(f"Original: {pil_image.width}×{pil_image.height}")
 
-# ====== LHS: CANVAS / IMAGE ======
+# ---------------------------------
+# LHS: Canvas / Image
+# ---------------------------------
 with col1:
     zoom = st.session_state.zoom_values[selected_name]
     scale = zoom / 100.0
     disp_w = int(pil_image.width * scale)
     disp_h = int(pil_image.height * scale)
-    
-    display_image = get_display_image_from_bytes(image_bytes, disp_w, disp_h)
+
+    display_image = get_display_image(image_bytes, disp_w, disp_h)
 
     def orig_to_display(rect_orig, s):
         return {
@@ -405,9 +409,10 @@ with col1:
             "strokeWidth": rect_display.get("strokeWidth", 2),
         }
 
-    # Build display objects from stored rectangles
+    # Build rectangles from state (one per field)
     all_display_objects = []
-    for fld, rect_orig in st.session_state.field_rects_orig[selected_name].items():
+    rects_for_image = st.session_state.field_rects_orig[selected_name]
+    for fld, rect_orig in rects_for_image.items():
         disp_rect = orig_to_display(rect_orig, scale)
         base = fld.split(": ", 1)[1] if ": " in fld else fld
         disp_rect["stroke"] = FIELD_COLORS.get(base, "#FF0000")
@@ -418,7 +423,7 @@ with col1:
     expected_count = len(all_display_objects)
 
     rect_ver = st.session_state.rect_version[selected_name]
-    num_rects = len(st.session_state.field_rects_orig[selected_name])
+    num_rects = len(rects_for_image)
     canvas_key = f"canvas_{selected_name}_z{zoom}_rv{rect_ver}_n{num_rects}"
 
     canvas_result = st_canvas(
@@ -434,18 +439,19 @@ with col1:
         key=canvas_key,
     )
 
-    # Detect new rectangle and auto-OCR
+    # Detect new rectangle
     if canvas_result.json_data is not None:
-        objs = canvas_result.json_data.get("objects", [])
+        objs = canvas_result.json_data.get("objects", []) or []
         if len(objs) > expected_count:
             new_rect_display = objs[-1]
             new_rect_orig = display_to_orig(new_rect_display, scale)
             new_rect_orig["stroke"] = field_color
 
-            # overwrite previous rect for this field
+            # Overwrite previous rect for this field (so old one disappears)
             st.session_state.field_rects_orig[selected_name][storage_field_name] = new_rect_orig
             st.session_state.rect_version[selected_name] += 1
 
+            # Auto OCR
             x1 = max(0, int(new_rect_orig["left"]))
             y1 = max(0, int(new_rect_orig["top"]))
             x2 = min(pil_image.width, int(new_rect_orig["left"] + new_rect_orig["width"]))
@@ -456,7 +462,6 @@ with col1:
                 try:
                     text = pytesseract.image_to_string(crop, config="--psm 6").strip()
                     if text:
-                        # update field values and the text-area state key
                         st.session_state.field_values[selected_name][storage_field_name] = text
                         value_state_key = f"value_{selected_name}_{storage_field_name}"
                         st.session_state[value_state_key] = text
@@ -468,14 +473,17 @@ with col1:
             else:
                 st.toast("✅ Rectangle saved")
 
-# ====== RHS BOTTOM: OCR VALUE + EXPORT ======
+            # Rerun once so canvas remounts with cleaned rectangles (no old ones)
+            st.experimental_rerun()
+
+# ---------------------------------
+# RHS BOTTOM: OCR value, all values, export
+# ---------------------------------
 with col2:
     st.markdown("#### ✏️ OCR & Value")
 
     current_rect_orig = st.session_state.field_rects_orig[selected_name].get(storage_field_name)
     value_state_key = f"value_{selected_name}_{storage_field_name}"
-
-    # initialise text state from saved field value on first use
     if value_state_key not in st.session_state:
         st.session_state[value_state_key] = st.session_state.field_values[selected_name].get(
             storage_field_name, ""
@@ -513,7 +521,6 @@ with col2:
         if current_rect_orig:
             st.button("🗑️ Delete", on_click=delete_rect)
 
-    # Text area bound to state key (so it reflects auto-OCR & Re-OCR without reruns)
     st.text_area(
         "Value (auto-filled by OCR)",
         key=value_state_key,
@@ -522,7 +529,7 @@ with col2:
         placeholder="Value (auto-filled by OCR)",
     )
 
-    # ========== ALL VALUES SECTION ==========
+    # All values
     with st.expander("📋 All Values"):
         for f in SINGLE_FIELDS:
             v = st.session_state.field_values[selected_name].get(f, "")
@@ -540,10 +547,9 @@ with col2:
                 for lif, v in vals:
                     st.write(f"  {lif}: {v}")
 
-    # ========== EXPORT SECTION ==========
+    # Export
     st.markdown("#### 📤 JSONL Export")
 
-    # Export ALL labeled remittances
     records_all = [
         build_gt_record_for_file(img["name"])
         for img in images
@@ -551,9 +557,7 @@ with col2:
     ]
 
     if records_all:
-        all_jsonl_str = "\n".join(
-            json.dumps(rec, ensure_ascii=False) for rec in records_all
-        )
+        all_jsonl_str = "\n".join(json.dumps(rec, ensure_ascii=False) for rec in records_all)
         st.download_button(
             "⬇️ Export ALL labeled (JSONL)",
             data=all_jsonl_str.encode("utf-8"),
@@ -563,7 +567,6 @@ with col2:
     else:
         st.caption("No labeled remittances yet.")
 
-    # Export CURRENT remittance
     current_record = build_gt_record_for_file(selected_name)
     with st.expander("Preview CURRENT JSON"):
         st.json(current_record)