Update src/streamlit_app.py

src/streamlit_app.py

@@ -33,14 +33,12 @@ from datetime import datetime
 if os.name == 'nt':  # Windows
     pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
 else:  # Linux/Mac (HF Spaces uses Linux)
-    # On HF Spaces with packages.txt, tesseract is in system PATH
-    # No need to set path explicitly
     pass
 
 # Page configuration
 st.set_page_config(page_title="Remittance Data Viewer", layout="wide")
 
-# Custom CSS to reduce gaps between form fields and style buttons
+# Custom CSS
 st.markdown("""
 <style>
     /* Reduce spacing between form fields */
@@ -56,21 +54,31 @@ st.markdown("""
         margin-bottom: 4px !important;
     }
     
-    /* Reduce gap between selectbox and following elements */
     .stSelectbox {
         margin-bottom: 4px !important;
     }
     
-    /* Style for small buttons */
+    /* Style for OCR buttons - small and compact */
     .stButton > button {
         padding: 0.25rem 0.5rem !important;
-        font-size: 1.2rem !important;
+        font-size: 0.85rem !important;
         line-height: 1 !important;
-        min-height: 2rem !important;
-        height: 2rem !important;
+        min-height: 1.8rem !important;
+        height: 1.8rem !important;
+    }
+    
+    /* Active OCR buttons - RED highlight */
+    .stButton > button[kind="primary"] {
+        background-color: #FF0000 !important;
+        border-color: #FF0000 !important;
+        color: white !important;
+    }
+    
+    .stButton > button[kind="primary"]:hover {
+        background-color: #CC0000 !important;
+        border-color: #CC0000 !important;
     }
     
-    /* Reduce padding in form containers */
     [data-testid="stVerticalBlock"] > [data-testid="stVerticalBlock"] {
         gap: 0.25rem !important;
     }
@@ -84,6 +92,12 @@ st.markdown("""
     [data-testid="stHorizontalBlock"] {
         gap: 0.5rem !important;
     }
+    
+    /* Active OCR field highlighting */
+    .ocr-active {
+        border: 2px solid #ff4b4b !important;
+        box-shadow: 0 0 5px rgba(255, 75, 75, 0.5) !important;
+    }
 </style>
 """, unsafe_allow_html=True)
 
@@ -104,17 +118,10 @@ def save_to_jsonl(data):
 def perform_ocr(image, bbox):
     """Perform OCR on the selected region of the image"""
     try:
-        # bbox is [x1, y1, x2, y2]
         x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
-        
-        # Ensure coordinates are within image bounds
         x1, y1 = max(0, x1), max(0, y1)
         x2, y2 = min(image.width, x2), min(image.height, y2)
-        
-        # Crop the image
         cropped = image.crop((x1, y1, x2, y2))
-        
-        # Perform OCR
         text = pytesseract.image_to_string(cropped, config='--psm 6').strip()
         return text
     except Exception as e:
@@ -122,42 +129,28 @@ def perform_ocr(image, bbox):
 
 def scale_image_to_fixed_size(image, target_width=700, target_height=900):
     """Scale and pad image to exact fixed size while maintaining aspect ratio and quality"""
-    # Convert image to RGB if it's not already (handles RGBA, L, etc.)
     if image.mode not in ('RGB', 'RGBA'):
         image = image.convert('RGB')
     elif image.mode == 'RGBA':
-        # Create white background for transparent images
         background = Image.new('RGB', image.size, (255, 255, 255))
-        background.paste(image, mask=image.split()[3])  # Use alpha channel as mask
+        background.paste(image, mask=image.split()[3])
         image = background
     
-    # Calculate scaling ratio to fit within target dimensions
     width_ratio = target_width / image.width
     height_ratio = target_height / image.height
-    
-    # Use the smaller ratio to ensure image fits within both constraints
     ratio = min(width_ratio, height_ratio)
     
-    # Calculate new dimensions
     new_width = int(image.width * ratio)
     new_height = int(image.height * ratio)
     
-    # Resize image with high-quality LANCZOS resampling
-    # Only resize if needed (don't upscale small images too much)
     if ratio < 1.0 or (ratio > 1.0 and ratio < 1.5):
         resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
     else:
-        # For significant upscaling, use BICUBIC which can be sharper
         resized_image = image.resize((new_width, new_height), Image.Resampling.BICUBIC)
     
-    # Create a new image with target size and white background
     final_image = Image.new('RGB', (target_width, target_height), (255, 255, 255))
-    
-    # Calculate position to paste resized image (center it)
     paste_x = (target_width - new_width) // 2
     paste_y = (target_height - new_height) // 2
-    
-    # Paste resized image onto white background
     final_image.paste(resized_image, (paste_x, paste_y))
     
     return final_image, ratio, paste_x, paste_y
@@ -194,27 +187,30 @@ if 'save_message_time' not in st.session_state:
 if 'just_saved' not in st.session_state:
     st.session_state.just_saved = False
 
-# Auto-save function
 def auto_save(index):
     """Automatically save changes to session state and mark as modified"""
     if st.session_state.edited_data:
         st.session_state.data = st.session_state.edited_data.copy()
         st.session_state.modified_indices.add(index)
 
-# Save button callback
-def save_changes_callback():
-    """Callback function for save button"""
-    auto_save(st.session_state.current_index)
-    st.session_state.save_message = "✅ Changes saved successfully!"
-    st.session_state.save_message_time = time.time()
+def activate_ocr_field(section, field, row_idx=None):
+    """Activate OCR for a specific field"""
+    st.session_state.ocr_active_section = section
+    st.session_state.ocr_active_field = field
+    st.session_state.ocr_line_item_row = row_idx
+    st.rerun()
+
+def is_ocr_active(section, field, row_idx=None):
+    """Check if this OCR button is currently active"""
+    return (st.session_state.ocr_active_section == section and 
+            st.session_state.ocr_active_field == field and 
+            st.session_state.ocr_line_item_row == row_idx)
 
 # PAGE 1: Upload Page
 if st.session_state.page == 'upload':
     st.title("📤 Remittance Data Viewer with OCR")
     st.markdown("### Upload your files to begin")
     
-    # Step 1: Upload JSONL
-    
     st.markdown("**Step 1: Upload JSONL File**")
     uploaded_file = st.file_uploader("Choose a JSONL file", type=['jsonl', 'json'])
     
@@ -227,11 +223,8 @@ if st.session_state.page == 'upload':
         except Exception as e:
             st.error(f"Error loading file: {str(e)}")
     
-    # Step 2: Upload Images
-    
     st.markdown("**Step 2: Upload Images Folder**")
     
-    
     uploaded_images = st.file_uploader(
         "Choose image files", 
         type=['png', 'jpg', 'jpeg', 'tiff', 'tif', 'bmp'],
@@ -240,7 +233,6 @@ if st.session_state.page == 'upload':
     )
     
     if uploaded_images:
-        # Load images into session state
         images_dict = {}
         for img_file in uploaded_images:
             try:
@@ -250,22 +242,18 @@ if st.session_state.page == 'upload':
                 st.warning(f"Could not load image {img_file.name}: {str(e)}")
         
         st.session_state.images = images_dict
-        # Show summary of loaded images and matches with ground truth
+        
         if st.session_state.data is not None:
-            # gather ground truth file names
             gt_file_names = [rec.get('file_name', '') for rec in st.session_state.data]
             matched_images = set()
             unmatched_gt_files = []
             
-            # Find matched images - CASE SENSITIVE EXACT MATCH ONLY
             for fname in gt_file_names:
                 if not fname:
                     continue
-                # Check for exact match in uploaded images
                 if fname in images_dict:
                     matched_images.add(fname)
             
-            # Find unmatched ground truth file names
             for fname in gt_file_names:
                 if fname and fname not in matched_images:
                     unmatched_gt_files.append(fname)
@@ -273,7 +261,6 @@ if st.session_state.page == 'upload':
             st.success(f"✅ Successfully loaded {len(images_dict)} images!")
             st.info(f"🔎 Exact matches: {len(matched_images)}/{len([f for f in gt_file_names if f])}")
             
-            # Show unmatched files
             if unmatched_gt_files:
                 st.warning(f"⚠️ {len(unmatched_gt_files)} file(s) from JSONL not matched to images:")
                 with st.expander(f"Show {len(unmatched_gt_files)} unmatched file names"):
@@ -285,8 +272,6 @@ if st.session_state.page == 'upload':
             st.success(f"✅ Successfully loaded {len(images_dict)} images!")
             st.info("ℹ️ Upload a JSONL file to see how many images match the ground truth 'file_name' field.")
     
-    # Continue Button
-    
     if st.session_state.data is not None:
         col1, col2, col3 = st.columns([1, 1, 1])
         with col2:
@@ -297,16 +282,13 @@ if st.session_state.page == 'upload':
 
 # PAGE 2: Viewer Page
 elif st.session_state.page == 'viewer':
-    # Clear old save messages (after 3 seconds)
     if st.session_state.save_message_time is not None:
         if time.time() - st.session_state.save_message_time > 3:
             st.session_state.save_message = None
             st.session_state.save_message_time = None
     
-    # Get today's date for filename
     today_date = datetime.now().strftime("%Y-%m-%d")
     
-    # Header with back button and download options
     col1, col2, col3, col4 = st.columns([1, 2, 2, 2])
     
     with col1:
@@ -318,7 +300,6 @@ elif st.session_state.page == 'viewer':
             st.session_state.save_message_time = None
             st.rerun()
     
-    # Download modified records and unmodified records separately
     with col2:
         if st.session_state.modified_indices:
             modified_data = [st.session_state.edited_data[i] for i in sorted(st.session_state.modified_indices)]
@@ -332,16 +313,10 @@ elif st.session_state.page == 'viewer':
                 use_container_width=True
             )
         else:
-            st.button(
-                "⬇️ No Modified Records",
-                disabled=True,
-                use_container_width=True
-            )
+            st.button("⬇️ No Modified Records", disabled=True, use_container_width=True)
     
-    # Download unmodified records (original data excluding modified)
     with col3:
         if st.session_state.modified_indices:
-            # Get original unmodified data
             unmodified_data = [st.session_state.data[i] for i in range(len(st.session_state.data)) 
                               if i not in st.session_state.modified_indices]
             jsonl_unmodified = save_to_jsonl(unmodified_data)
@@ -353,13 +328,8 @@ elif st.session_state.page == 'viewer':
                 use_container_width=True
             )
         else:
-            st.button(
-                "⬇️ No Unmodified Records",
-                disabled=True,
-                use_container_width=True
-            )
+            st.button("⬇️ No Unmodified Records", disabled=True, use_container_width=True)
     
-    # Download all edited data
     with col4:
         jsonl_all = save_to_jsonl(st.session_state.edited_data)
         st.download_button(
@@ -369,11 +339,7 @@ elif st.session_state.page == 'viewer':
             mime="application/jsonl",
             use_container_width=True
         )
-
     
-    
-    
-    # File selector dropdown
     file_names = [record.get('file_name', f'Record {i}') for i, record in enumerate(st.session_state.data)]
     
     selected_file = st.selectbox(
@@ -386,9 +352,6 @@ elif st.session_state.page == 'viewer':
     st.session_state.current_index = selected_file
     current_record = st.session_state.edited_data[selected_file]
     
-    
-    
-    # Main layout: LHS (Image) and RHS (Details) - REDUCED GAP
     left_col, right_col = st.columns([1.3, 1], gap="small")
     
     # LEFT SIDE: Image Display with OCR Canvas
@@ -400,7 +363,6 @@ elif st.session_state.page == 'viewer':
         if file_name:
             st.caption(f"**File:** {file_name}")
             
-            # Try to find matching image - CASE SENSITIVE EXACT MATCH ONLY
             current_image = None
             if file_name in st.session_state.images:
                 current_image = st.session_state.images[file_name]
@@ -414,10 +376,8 @@ elif st.session_state.page == 'viewer':
                         st.text(f"  ... and {len(st.session_state.images) - 20} more")
             
             if current_image:
-                # Scale image to fixed size
                 scaled_image, scale_ratio, paste_x, paste_y = scale_image_to_fixed_size(current_image)
                 
-                # Always show canvas for drawing rectangles
                 canvas_result = st_canvas(
                     fill_color="rgba(255, 165, 0, 0.3)",
                     stroke_width=2,
@@ -430,14 +390,11 @@ elif st.session_state.page == 'viewer':
                     key=f"canvas_{selected_file}_{st.session_state.canvas_key}",
                 )
                 
-                # Process OCR when rectangle is drawn and field is selected
                 if canvas_result.json_data is not None and st.session_state.ocr_active_field:
                     objects = canvas_result.json_data["objects"]
                     if len(objects) > 0:
-                        # Get the last drawn rectangle
                         rect = objects[-1]
                         
-                        # Adjust coordinates for padding and scale back to original image coordinates
                         bbox = [
                             (rect["left"] - paste_x) / scale_ratio,
                             (rect["top"] - paste_y) / scale_ratio,
@@ -445,25 +402,21 @@ elif st.session_state.page == 'viewer':
                             (rect["top"] + rect["height"] - paste_y) / scale_ratio
                         ]
                         
-                        # Perform OCR on original image
                         with st.spinner("Performing OCR..."):
                             ocr_text = perform_ocr(current_image, bbox)
                         
                         if ocr_text and not ocr_text.startswith("OCR Error"):
                             st.success(f"✅ OCR Result: {ocr_text}")
                             
-                            # Update the field value
                             gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
                             
                             if st.session_state.ocr_active_section == 'Line_items':
-                                # Handle line items
                                 line_items = gt_parse.get('Line_items', [])
                                 row_idx = st.session_state.ocr_line_item_row
                                 if row_idx is not None and row_idx < len(line_items):
                                     line_items[row_idx][st.session_state.ocr_active_field] = ocr_text
                                     gt_parse['Line_items'] = line_items
                             else:
-                                # Handle other sections
                                 section = st.session_state.ocr_active_section
                                 field = st.session_state.ocr_active_field
                                 if section not in gt_parse:
@@ -471,10 +424,10 @@ elif st.session_state.page == 'viewer':
                                 gt_parse[section][field] = ocr_text
                             
                             st.session_state.edited_data[selected_file]['gt_parse'] = gt_parse
+                            st.session_state.modified_indices.add(selected_file)
                             
-                            # Clear canvas and reset
+                            # Clear canvas for next OCR
                             st.session_state.canvas_key += 1
-                            time.sleep(0.3)
                             st.rerun()
                         else:
                             st.error(ocr_text)
@@ -485,9 +438,8 @@ elif st.session_state.page == 'viewer':
     with right_col:
         st.markdown("### 📝 Document Details")
         
-        gt_parse = current_record.get('gt_parse', {})
+        gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
         
-        # Create tabs for each section
         tab1, tab2, tab3, tab4 = st.tabs([
             "📄 Remittance Details", 
             "👥 Party Details", 
@@ -497,340 +449,487 @@ elif st.session_state.page == 'viewer':
         
         # TAB 1: Remittance Details
         with tab1:
-            
-            
-            # OCR Field Selector
-            remittance_fields = [
-                'Select fields',
-                'Remittance_adv_no',
-                'Remittance_adv_date',
-                'Payment_method',
-                'FCY',
-                'Total_payment_amt_FCY',
-                'Payment_date',
-                'Payment_ref_no'
-            ]
-            
-            selected_rem_field = st.selectbox(
-                "🔍 Select field to populate via OCR:",
-                options=remittance_fields,
-                key=f"rem_ocr_select_{selected_file}"
-            )
-            
-            if selected_rem_field != 'Select fields':
-                st.session_state.ocr_active_section = 'Remittance_details'
-                st.session_state.ocr_active_field = selected_rem_field
-                st.session_state.ocr_line_item_row = None
-            else:
-                if st.session_state.ocr_active_section == 'Remittance_details':
-                    st.session_state.ocr_active_section = None
-                    st.session_state.ocr_active_field = None
-            
             remittance = gt_parse.get('Remittance_details', {})
             
-            remittance['Remittance_adv_no'] = st.text_input(
-                "Remittance Advice No", 
-                value=remittance.get('Remittance_adv_no', ''),
-                key=f"rem_adv_no_{selected_file}"
-            )
-            remittance['Remittance_adv_date'] = st.text_input(
-                "Remittance Advice Date", 
-                value=remittance.get('Remittance_adv_date', ''),
-                key=f"rem_adv_date_{selected_file}"
-            )
-            remittance['Payment_method'] = st.text_input(
-                "Payment Method", 
-                value=remittance.get('Payment_method', ''),
-                key=f"payment_method_{selected_file}"
-            )
-            remittance['FCY'] = st.text_input(
-                "FCY (Foreign Currency)", 
-                value=remittance.get('FCY', ''),
-                key=f"fcy_{selected_file}"
-            )
-            remittance['Total_payment_amt_FCY'] = st.text_input(
-                "Total Payment Amount (FCY)", 
-                value=remittance.get('Total_payment_amt_FCY', ''),
-                key=f"total_payment_{selected_file}"
-            )
-            remittance['Payment_date'] = st.text_input(
-                "Payment Date", 
-                value=remittance.get('Payment_date', ''),
-                key=f"payment_date_{selected_file}"
-            )
-            remittance['Payment_ref_no'] = st.text_input(
-                "Payment Reference No", 
-                value=remittance.get('Payment_ref_no', ''),
-                key=f"payment_ref_{selected_file}"
-            )
+            # Each field with OCR button
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Remittance_adv_no'] = st.text_input(
+                    "Remittance Advice No", 
+                    value=remittance.get('Remittance_adv_no', ''),
+                    key=f"rem_adv_no_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_rem_adv_no_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Remittance_adv_no') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Remittance_details', 'Remittance_adv_no')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Remittance_adv_date'] = st.text_input(
+                    "Remittance Advice Date", 
+                    value=remittance.get('Remittance_adv_date', ''),
+                    key=f"rem_adv_date_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_rem_adv_date_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Remittance_adv_date') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Remittance_details', 'Remittance_adv_date')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Payment_method'] = st.text_input(
+                    "Payment Method", 
+                    value=remittance.get('Payment_method', ''),
+                    key=f"payment_method_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_payment_method_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Payment_method') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Remittance_details', 'Payment_method')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['FCY'] = st.text_input(
+                    "FCY (Foreign Currency)", 
+                    value=remittance.get('FCY', ''),
+                    key=f"fcy_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_fcy_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'FCY') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Remittance_details', 'FCY')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Total_payment_amt_FCY'] = st.text_input(
+                    "Total Payment Amount (FCY)", 
+                    value=remittance.get('Total_payment_amt_FCY', ''),
+                    key=f"total_payment_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_total_payment_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Total_payment_amt_FCY') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Remittance_details', 'Total_payment_amt_FCY')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Payment_date'] = st.text_input(
+                    "Payment Date", 
+                    value=remittance.get('Payment_date', ''),
+                    key=f"payment_date_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_payment_date_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Payment_date') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Remittance_details', 'Payment_date')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                remittance['Payment_ref_no'] = st.text_input(
+                    "Payment Reference No", 
+                    value=remittance.get('Payment_ref_no', ''),
+                    key=f"payment_ref_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_payment_ref_{selected_file}", 
+                           type="primary" if is_ocr_active('Remittance_details', 'Payment_ref_no') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Remittance_details', 'Payment_ref_no')
             
             gt_parse['Remittance_details'] = remittance
         
         # TAB 2: Customer/Supplier Details
         with tab2:
-            
-            
-            # OCR Field Selector
-            customer_fields = [
-                'Select fields',
-                'Customer_name',
-                'Customer_address',
-                'Customer_contact_info',
-                'Supplier_name',
-                'Supplier_address',
-                'Supplier_contact_info'
-            ]
-            
-            selected_cust_field = st.selectbox(
-                "🔍 Select field to populate via OCR:",
-                options=customer_fields,
-                key=f"cust_ocr_select_{selected_file}"
-            )
-            
-            if selected_cust_field != 'Select fields':
-                st.session_state.ocr_active_section = 'Customer_supplier_details'
-                st.session_state.ocr_active_field = selected_cust_field
-                st.session_state.ocr_line_item_row = None
-            else:
-                if st.session_state.ocr_active_section == 'Customer_supplier_details':
-                    st.session_state.ocr_active_section = None
-                    st.session_state.ocr_active_field = None
-            
             st.markdown("**Customer Details**")
             customer_supplier = gt_parse.get('Customer_supplier_details', {})
             
-            customer_supplier['Customer_name'] = st.text_input(
-                "Customer Name", 
-                value=customer_supplier.get('Customer_name', ''),
-                key=f"cust_name_{selected_file}"
-            )
-            customer_supplier['Customer_address'] = st.text_area(
-                "Customer Address", 
-                value=customer_supplier.get('Customer_address', ''),
-                key=f"cust_addr_{selected_file}",
-                height=60
-            )
-            customer_supplier['Customer_contact_info'] = st.text_input(
-                "Customer Contact Info", 
-                value=customer_supplier.get('Customer_contact_info', ''),
-                key=f"cust_contact_{selected_file}"
-            )
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Customer_name'] = st.text_input(
+                    "Customer Name", 
+                    value=customer_supplier.get('Customer_name', ''),
+                    key=f"cust_name_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_cust_name_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Customer_name') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Customer_supplier_details', 'Customer_name')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Customer_address'] = st.text_area(
+                    "Customer Address", 
+                    value=customer_supplier.get('Customer_address', ''),
+                    key=f"cust_addr_{selected_file}",
+                    height=60
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_cust_addr_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Customer_address') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Customer_supplier_details', 'Customer_address')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Customer_contact_info'] = st.text_input(
+                    "Customer Contact Info", 
+                    value=customer_supplier.get('Customer_contact_info', ''),
+                    key=f"cust_contact_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_cust_contact_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Customer_contact_info') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Customer_supplier_details', 'Customer_contact_info')
             
             st.markdown("**Supplier Details**")
-            customer_supplier['Supplier_name'] = st.text_input(
-                "Supplier Name", 
-                value=customer_supplier.get('Supplier_name', ''),
-                key=f"supp_name_{selected_file}"
-            )
-            customer_supplier['Supplier_address'] = st.text_area(
-                "Supplier Address", 
-                value=customer_supplier.get('Supplier_address', ''),
-                key=f"supp_addr_{selected_file}",
-                height=60
-            )
-            customer_supplier['Supplier_contact_info'] = st.text_input(
-                "Supplier Contact Info", 
-                value=customer_supplier.get('Supplier_contact_info', ''),
-                key=f"supp_contact_{selected_file}"
-            )
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Supplier_name'] = st.text_input(
+                    "Supplier Name", 
+                    value=customer_supplier.get('Supplier_name', ''),
+                    key=f"supp_name_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_supp_name_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Supplier_name') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Customer_supplier_details', 'Supplier_name')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Supplier_address'] = st.text_area(
+                    "Supplier Address", 
+                    value=customer_supplier.get('Supplier_address', ''),
+                    key=f"supp_addr_{selected_file}",
+                    height=60
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_supp_addr_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Supplier_address') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Customer_supplier_details', 'Supplier_address')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                customer_supplier['Supplier_contact_info'] = st.text_input(
+                    "Supplier Contact Info", 
+                    value=customer_supplier.get('Supplier_contact_info', ''),
+                    key=f"supp_contact_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_supp_contact_{selected_file}", 
+                           type="primary" if is_ocr_active('Customer_supplier_details', 'Supplier_contact_info') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Customer_supplier_details', 'Supplier_contact_info')
             
             gt_parse['Customer_supplier_details'] = customer_supplier
         
         # TAB 3: Bank Details
         with tab3:
-            
-            
-            # OCR Field Selector
-            bank_fields = [
-                'Select fields',
-                'Bank_name',
-                'Bank_acc_no',
-                'Bank_routing_no',
-                'Swift_code'
-            ]
-            
-            selected_bank_field = st.selectbox(
-                "🔍 Select field to populate via OCR:",
-                options=bank_fields,
-                key=f"bank_ocr_select_{selected_file}"
-            )
-            
-            if selected_bank_field != 'Select fields':
-                st.session_state.ocr_active_section = 'Bank_details'
-                st.session_state.ocr_active_field = selected_bank_field
-                st.session_state.ocr_line_item_row = None
-            else:
-                if st.session_state.ocr_active_section == 'Bank_details':
-                    st.session_state.ocr_active_section = None
-                    st.session_state.ocr_active_field = None
-            
             bank = gt_parse.get('Bank_details', {})
             
-            bank['Bank_name'] = st.text_input(
-                "Bank Name", 
-                value=bank.get('Bank_name', ''),
-                key=f"bank_name_{selected_file}"
-            )
-            bank['Bank_acc_no'] = st.text_input(
-                "Bank Account No", 
-                value=bank.get('Bank_acc_no', ''),
-                key=f"bank_acc_{selected_file}"
-            )
-            bank['Bank_routing_no'] = st.text_input(
-                "Bank Routing No", 
-                value=bank.get('Bank_routing_no', ''),
-                key=f"bank_routing_{selected_file}"
-            )
-            bank['Swift_code'] = st.text_input(
-                "SWIFT Code", 
-                value=bank.get('Swift_code', ''),
-                key=f"swift_{selected_file}"
-            )
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                bank['Bank_name'] = st.text_input(
+                    "Bank Name", 
+                    value=bank.get('Bank_name', ''),
+                    key=f"bank_name_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_bank_name_{selected_file}", 
+                           type="primary" if is_ocr_active('Bank_details', 'Bank_name') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Bank_details', 'Bank_name')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                bank['Bank_acc_no'] = st.text_input(
+                    "Bank Account No", 
+                    value=bank.get('Bank_acc_no', ''),
+                    key=f"bank_acc_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_bank_acc_{selected_file}", 
+                           type="primary" if is_ocr_active('Bank_details', 'Bank_acc_no') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Bank_details', 'Bank_acc_no')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                bank['Bank_routing_no'] = st.text_input(
+                    "Bank Routing No", 
+                    value=bank.get('Bank_routing_no', ''),
+                    key=f"bank_routing_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_bank_routing_{selected_file}", 
+                           type="primary" if is_ocr_active('Bank_details', 'Bank_routing_no') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Bank_details', 'Bank_routing_no')
+            
+            col_input, col_btn = st.columns([5, 1])
+            with col_input:
+                bank['Swift_code'] = st.text_input(
+                    "SWIFT Code", 
+                    value=bank.get('Swift_code', ''),
+                    key=f"swift_{selected_file}"
+                )
+            with col_btn:
+                st.markdown("<br>", unsafe_allow_html=True)
+                if st.button("🔍", key=f"ocr_swift_{selected_file}", 
+                           type="primary" if is_ocr_active('Bank_details', 'Swift_code') else "secondary",
+                           help="OCR this field"):
+                    activate_ocr_field('Bank_details', 'Swift_code')
             
             gt_parse['Bank_details'] = bank
         
         # TAB 4: Line Items
         with tab4:
+            current_gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
+            line_items = current_gt_parse.get('Line_items', [])
             
-            
-            # OCR Controls for Line Items - Fixed layout
-            line_items = gt_parse.get('Line_items', [])
-            
-            # Adjusted column widths - all controls in single compact line
-            col_field, col_row, col_add, col_remove = st.columns([1.5, 0.7, 0.30, 0.30])
-            
-            line_item_fields = [
-                'Select fields',
-                'Po_number',
-                'Invoice_no',
-                'Other_doc_ref_no',
-                'Invoice_date',
-                'Invoice_amount_FCY',
-                'Amount_paid_for_each_invoice',
-                'Outstanding_balance_FCY',
-                'Discounts_taken_FCY',
-                'Adjustments(without_holding_tax)_FCY',
-                'Descriptions'
-            ]
-            
-            with col_field:
-                selected_line_field = st.selectbox(
-                    "🔍 Field:",
-                    options=line_item_fields,
-                    key=f"line_ocr_field_{selected_file}"
-                )
-            
-            with col_row:
-                if len(line_items) > 0:
-                    selected_row = st.selectbox(
-                        "Row:",
-                        options=list(range(len(line_items))),
-                        format_func=lambda x: f"Row {x + 1}",
-                        key=f"line_ocr_row_{selected_file}"
-                    )
-                else:
-                    st.selectbox("Row:", options=[], disabled=True, key=f"line_ocr_row_empty_{selected_file}")
-                    selected_row = None
-            
+            # Add/Remove row buttons
+            col_add, col_remove = st.columns([1, 1])
             with col_add:
-                # Use button with on_click callback to prevent loop
-                if st.button("➕", key=f"add_row_{selected_file}", help="Add new row"):
+                if st.button("➕ Add New Row", key=f"add_row_{selected_file}", use_container_width=True):
                     if not st.session_state.button_clicked:
                         st.session_state.button_clicked = True
                         new_row = {
-                            "Po_number": "",
-                            "Invoice_no": "",
-                            "Other_doc_ref_no": "",
-                            "Invoice_date": "",
-                            "Invoice_amount_FCY": "",
-                            "Amount_paid_for_each_invoice": "",
-                            "Outstanding_balance_FCY": "",
-                            "Discounts_taken_FCY": "",
-                            "Adjustments(without_holding_tax)_FCY": "",
+                            "Po_number": "", "Invoice_no": "", "Other_doc_ref_no": "",
+                            "Invoice_date": "", "Invoice_amount_FCY": "",
+                            "Amount_paid_for_each_invoice": "", "Outstanding_balance_FCY": "",
+                            "Discounts_taken_FCY": "", "Adjustments(without_holding_tax)_FCY": "",
                             "Descriptions": ""
                         }
-                        line_items.append(new_row)
-                        gt_parse['Line_items'] = line_items
-                        st.session_state.edited_data[selected_file]['gt_parse'] = gt_parse
+                        current_gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
+                        current_line_items = current_gt_parse.get('Line_items', [])
+                        current_line_items.append(new_row)
+                        current_gt_parse['Line_items'] = current_line_items
+                        st.session_state.edited_data[selected_file]['gt_parse'] = current_gt_parse
                         st.session_state.modified_indices.add(selected_file)
                         st.rerun()
             
             with col_remove:
-                if st.button("➖", key=f"remove_row_{selected_file}", help="Remove selected row", disabled=(len(line_items) == 0)):
-                    if not st.session_state.button_clicked and len(line_items) > 0 and selected_row is not None:
+                if st.button("➖ Remove Last Row", key=f"remove_row_{selected_file}", 
+                           disabled=(len(line_items) == 0), use_container_width=True):
+                    if not st.session_state.button_clicked and len(line_items) > 0:
                         st.session_state.button_clicked = True
-                        line_items.pop(selected_row)
-                        gt_parse['Line_items'] = line_items
-                        st.session_state.edited_data[selected_file]['gt_parse'] = gt_parse
+                        current_gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
+                        current_line_items = current_gt_parse.get('Line_items', [])
+                        current_line_items.pop()
+                        current_gt_parse['Line_items'] = current_line_items
+                        st.session_state.edited_data[selected_file]['gt_parse'] = current_gt_parse
                         st.session_state.modified_indices.add(selected_file)
                         st.rerun()
             
-            # Reset button clicked flag after processing
             if st.session_state.button_clicked:
                 st.session_state.button_clicked = False
             
-            # Set OCR state for line items
-            if selected_line_field != 'Select fields' and selected_row is not None:
-                st.session_state.ocr_active_section = 'Line_items'
-                st.session_state.ocr_active_field = selected_line_field
-                st.session_state.ocr_line_item_row = selected_row
-            else:
-                if st.session_state.ocr_active_section == 'Line_items':
-                    st.session_state.ocr_active_section = None
-                    st.session_state.ocr_active_field = None
-                    st.session_state.ocr_line_item_row = None
-            
+            st.markdown("---")
             
+            # Display each row as an expander with OCR buttons
+            current_gt_parse = st.session_state.edited_data[selected_file].get('gt_parse', {})
+            line_items = current_gt_parse.get('Line_items', [])
             
-            # Display line items table
             if line_items:
-                df = pd.DataFrame(line_items)
-                
-                # Set index to start from 1 instead of 0
-                df.index = range(1, len(df) + 1)
+                for idx, item in enumerate(line_items):
+                    with st.expander(f"**Row {idx + 1}** - Invoice: {item.get('Invoice_no', 'N/A')}", expanded=(idx == 0)):
+                        # PO Number
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Po_number'] = st.text_input(
+                                "PO Number", 
+                                value=item.get('Po_number', ''),
+                                key=f"po_num_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_po_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Po_number', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Po_number', idx)
+                        
+                        # Invoice No
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Invoice_no'] = st.text_input(
+                                "Invoice No", 
+                                value=item.get('Invoice_no', ''),
+                                key=f"inv_no_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_inv_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Invoice_no', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Invoice_no', idx)
+                        
+                        # Other Doc Ref No
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Other_doc_ref_no'] = st.text_input(
+                                "Other Doc Ref No", 
+                                value=item.get('Other_doc_ref_no', ''),
+                                key=f"other_doc_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_other_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Other_doc_ref_no', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Other_doc_ref_no', idx)
+                        
+                        # Invoice Date
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Invoice_date'] = st.text_input(
+                                "Invoice Date", 
+                                value=item.get('Invoice_date', ''),
+                                key=f"inv_date_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_inv_date_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Invoice_date', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Invoice_date', idx)
+                        
+                        # Invoice Amount FCY
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Invoice_amount_FCY'] = st.text_input(
+                                "Invoice Amount FCY", 
+                                value=item.get('Invoice_amount_FCY', ''),
+                                key=f"inv_amt_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_inv_amt_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Invoice_amount_FCY', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Invoice_amount_FCY', idx)
+                        
+                        # Amount Paid
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Amount_paid_for_each_invoice'] = st.text_input(
+                                "Amount Paid", 
+                                value=item.get('Amount_paid_for_each_invoice', ''),
+                                key=f"amt_paid_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_amt_paid_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Amount_paid_for_each_invoice', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Amount_paid_for_each_invoice', idx)
+                        
+                        # Outstanding Balance
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Outstanding_balance_FCY'] = st.text_input(
+                                "Outstanding Balance FCY", 
+                                value=item.get('Outstanding_balance_FCY', ''),
+                                key=f"out_bal_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_out_bal_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Outstanding_balance_FCY', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Outstanding_balance_FCY', idx)
+                        
+                        # Discounts
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Discounts_taken_FCY'] = st.text_input(
+                                "Discounts Taken FCY", 
+                                value=item.get('Discounts_taken_FCY', ''),
+                                key=f"disc_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_disc_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Discounts_taken_FCY', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Discounts_taken_FCY', idx)
+                        
+                        # Adjustments
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Adjustments(without_holding_tax)_FCY'] = st.text_input(
+                                "Adjustments FCY", 
+                                value=item.get('Adjustments(without_holding_tax)_FCY', ''),
+                                key=f"adj_{selected_file}_{idx}"
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_adj_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Adjustments(without_holding_tax)_FCY', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Adjustments(without_holding_tax)_FCY', idx)
+                        
+                        # Descriptions
+                        col_input, col_btn = st.columns([5, 1])
+                        with col_input:
+                            item['Descriptions'] = st.text_area(
+                                "Descriptions", 
+                                value=item.get('Descriptions', ''),
+                                key=f"desc_{selected_file}_{idx}",
+                                height=60
+                            )
+                        with col_btn:
+                            st.markdown("<br>", unsafe_allow_html=True)
+                            if st.button("🔍", key=f"ocr_desc_{selected_file}_{idx}", 
+                                       type="primary" if is_ocr_active('Line_items', 'Descriptions', idx) else "secondary",
+                                       help="OCR this field"):
+                                activate_ocr_field('Line_items', 'Descriptions', idx)
                 
-                # Convert amount fields to numeric
-                amount_fields = ['Invoice_amount_FCY', 'Amount_paid_for_each_invoice', 
-                               'Outstanding_balance_FCY', 'Discounts_taken_FCY', 
-                               'Adjustments(without_holding_tax)_FCY']
+                # Update line items back to gt_parse
+                current_gt_parse['Line_items'] = line_items
                 
-                for field in amount_fields:
-                    if field in df.columns:
-                        df[field] = pd.to_numeric(df[field].replace('', None), errors='coerce')
+                st.markdown("---")
+                st.markdown("**📊 Line Items Summary Table**")
                 
-                column_config = {
-                    "Po_number": st.column_config.TextColumn("PO Number", width="medium"),
-                    "Invoice_no": st.column_config.TextColumn("Invoice No", width="medium"),
-                    "Other_doc_ref_no": st.column_config.TextColumn("Other Doc Ref No", width="medium"),
-                    "Invoice_date": st.column_config.TextColumn("Invoice Date", width="medium"),
-                    "Invoice_amount_FCY": st.column_config.NumberColumn("Invoice Amt FCY", width="medium", format="%.2f"),
-                    "Amount_paid_for_each_invoice": st.column_config.NumberColumn("Amount Paid", width="medium", format="%.2f"),
-                    "Outstanding_balance_FCY": st.column_config.NumberColumn("Outstanding FCY", width="medium", format="%.2f"),
-                    "Discounts_taken_FCY": st.column_config.NumberColumn("Discounts FCY", width="medium", format="%.2f"),
-                    "Adjustments(without_holding_tax)_FCY": st.column_config.NumberColumn("Adjustments FCY", width="medium", format="%.2f"),
-                    "Descriptions": st.column_config.TextColumn("Descriptions", width="medium"),
-                }
+                # Display summary table with index starting from 1
+                df = pd.DataFrame(line_items)
+                df.index = df.index + 1  # Start index from 1
+                df.index.name = 'SL No'
                 
-                edited_df = st.data_editor(
+                st.dataframe(
                     df,
-                    column_config=column_config,
-                    num_rows="fixed",
                     use_container_width=True,
-                    key=f"line_items_table_{selected_file}",
-                    hide_index=False
+                    height=300
                 )
-                
-                # Convert back to string and reset index to 0-based for storage
-                edited_df.index = range(len(edited_df))
-                for field in amount_fields:
-                    if field in edited_df.columns:
-                        edited_df[field] = edited_df[field].apply(lambda x: str(x) if pd.notna(x) else '')
-                
-                gt_parse['Line_items'] = edited_df.to_dict('records')
             else:
-                st.info("No line items. Click ➕ to add a new row.")
+                st.info("No line items. Click '➕ Add New Row' to add a new row.")
         
-        # Update the edited data
         st.session_state.edited_data[selected_file]['gt_parse'] = gt_parse
         
         # Save button
@@ -845,10 +944,8 @@ elif st.session_state.page == 'viewer':
                     st.session_state.save_message_time = time.time()
                     st.rerun()
         
-        # Reset the just_saved flag after rerun
         if st.session_state.just_saved:
             st.session_state.just_saved = False
         
-        # Display save message under the button (appears after rerun)
         if st.session_state.save_message:
-            st.success(st.session_state.save_message)
+            st.success(st.session_state.save_message)