Improved efficiency of review page navigation, especially for large documents. Updated user guide

.dockerignore

@@ -17,4 +17,6 @@ dist/*
 build_deps/*
 logs/*
 config/*
-user_guide/*
+user_guide/*
+cdk/*
+web/*

.gitignore

@@ -18,4 +18,6 @@ build_deps/*
 logs/*
 config/*
 doc_redaction_amplify_app/*
-user_guide/*
+user_guide/*
+cdk/*
+web/*

README.md

@@ -20,6 +20,12 @@ NOTE: The app is not 100% accurate, and it will miss some personal information.
 
 # USER GUIDE
 
+## Experiment with the test (public) version of the app
+You can test out many of the features described in this user guide at the [public test version of the app](https://huggingface.co/spaces/seanpedrickcase/document_redaction), which is free. AWS functions (e.g. Textract, Comprehend) are not enabled (unless you have valid API keys).
+
+## Chat over this user guide
+You can now [speak with a chat bot about this user guide](https://huggingface.co/spaces/seanpedrickcase/Light-PDF-Web-QA-Chatbot) (beta!)
+
 ## Table of contents
 
 - [Example data files](#example-data-files)
@@ -35,57 +41,97 @@ NOTE: The app is not 100% accurate, and it will miss some personal information.
 - [Reviewing and modifying suggested redactions](#reviewing-and-modifying-suggested-redactions)
 
 See the [advanced user guide here](#advanced-user-guide):
-- [Modifying and merging redaction review files](#modifying-and-merging-redaction-review-files)
-    - [Modifying existing redaction review files](#modifying-existing-redaction-review-files)
-    - [Merging existing redaction review files](#merging-existing-redaction-review-files)
+- [Merging redaction review files](#merging-redaction-review-files)
 - [Identifying and redacting duplicate pages](#identifying-and-redacting-duplicate-pages)
 - [Fuzzy search and redaction](#fuzzy-search-and-redaction)
 - [Export redactions to and import from Adobe Acrobat](#export-to-and-import-from-adobe)
     - [Exporting to Adobe Acrobat](#exporting-to-adobe-acrobat)
     - [Importing from Adobe Acrobat](#importing-from-adobe-acrobat)
+- [Using the AWS Textract document API](#using-the-aws-textract-document-api)
 - [Using AWS Textract and Comprehend when not running in an AWS environment](#using-aws-textract-and-comprehend-when-not-running-in-an-aws-environment)
+- [Modifying existing redaction review files](#modifying-existing-redaction-review-files)
 
 ## Example data files
 
-Please refer to these example files to follow this guide: 
+Please try these example files to follow along with this guide: 
 - [Example of files sent to a professor before applying](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/example_of_emails_sent_to_a_professor_before_applying.pdf)
 - [Example complaint letter (jpg)](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/example_complaint_letter.jpg)
 - [Partnership Agreement Toolkit (for signatures and more advanced usage)](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/Partnership-Agreement-Toolkit_0_0.pdf)
+- [Dummy case note data](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/combined_case_notes.csv)
 
 ## Basic redaction
 
-The document redaction app can detect personally-identifiable information (PII) in documents. Documents can be redacted directly, or suggested redactions can be reviewed and modified using a grapical user interface.
+The document redaction app can detect personally-identifiable information (PII) in documents. Documents can be redacted directly, or suggested redactions can be reviewed and modified using a grapical user interface. Basic document redaction can be performed quickly using the default options.
 
 Download the example PDFs above to your computer. Open up the redaction app with the link provided by email.
 
 ![Upload files](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/file_upload_highlight.PNG)
 
-Click on the upload files area, and select the three different files (they should all be stored in the same folder if you want them to be redacted at the same time).
+### Upload files to the app
+
+The 'Redact PDFs/images tab' currently accepts PDFs and image files (JPG, PNG) for redaction. Click on the 'Drop files here or Click to Upload' area of the screen, and select one of the three different [example files](#example-data-files) (they should all be stored in the same folder if you want them to be redacted at the same time).
+
+### Text extraction
 
-First, select one of the three text extraction options below:
-- 'Local model - selectable text' - This will read text directly from PDFs that have selectable text to redact (using PikePDF). This is fine for most PDFs, but will find nothing if the PDF does not have selectable text, and it is not good for handwriting or signatures. If it encounters an image file, it will send it onto the second option below.
-- 'Local OCR model - PDFs without selectable text' - This option will use a simple Optical Character Recognition (OCR) model (Tesseract) to pull out text from a PDF/image that it 'sees'. This can handle most typed text in PDFs/images without selectable text, but struggles with handwriting/signatures. If you are interested in the latter, then you should use the third option if available.
-- 'AWS Textract service - all PDF types' - Only available for instances of the app running on AWS. AWS Textract is a service that performs OCR on documents within their secure service. This is a more advanced version of OCR compared to the local option, and carries a (relatively small) cost. Textract excels in complex documents based on images, or documents that contain a lot of handwriting and signatures.
+First, select one of the three text extraction options:
+- **'Local model - selectable text'** - This will read text directly from PDFs that have selectable text to redact (using PikePDF). This is fine for most PDFs, but will find nothing if the PDF does not have selectable text, and it is not good for handwriting or signatures. If it encounters an image file, it will send it onto the second option below.
+- **'Local OCR model - PDFs without selectable text'** - This option will use a simple Optical Character Recognition (OCR) model (Tesseract) to pull out text from a PDF/image that it 'sees'. This can handle most typed text in PDFs/images without selectable text, but struggles with handwriting/signatures. If you are interested in the latter, then you should use the third option if available.
+- **'AWS Textract service - all PDF types'** - Only available for instances of the app running on AWS. AWS Textract is a service that performs OCR on documents within their secure service. This is a more advanced version of OCR compared to the local option, and carries a (relatively small) cost. Textract excels in complex documents based on images, or documents that contain a lot of handwriting and signatures.
+
+### Optional - select signature extraction
+If you chose the AWS Textract service above, you can choose if you want handwriting and/or signatures redacted by default. Choosing signatures here will have a cost implication, as identifying signatures will cost ~£2.66 ($3.50) per 1,000 pages vs ~£1.14 ($1.50) per 1,000 pages without signature detection. 
+
+![AWS Textract handwriting and signature options](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/textract_handwriting_signatures.PNG)
+
+### PII redaction method
 
 If you are running with the AWS service enabled, here you will also have a choice for PII redaction method:
-- 'Local' - This uses the spacy package to rapidly detect PII in extracted text. This method is often sufficient if you are just interested in redacting specific terms defined in a custom list. 
-- 'AWS Comprehend' - This method calls an AWS service to provide more accurate identification of PII in extracted text.
+- **'Only extract text - (no redaction)'** - If you are only interested in getting the text out of the document for further processing (e.g. to find duplicate pages, or to review text on the Review redactions page)
+- **'Local'** - This uses the spacy package to rapidly detect PII in extracted text. This method is often sufficient if you are just interested in redacting specific terms defined in a custom list. 
+- **'AWS Comprehend'** - This method calls an AWS service to provide more accurate identification of PII in extracted text.
+
+### Optional - costs and time estimation
+If the option is enabled (by your system admin, in the config file), you will see a cost and time estimate for the redaction process. 'Existing Textract output file found' will be checked automatically if previous Textract text extraction files exist in the output folder, or have been [previously uploaded by the user](#aws-textract-outputs) (saving time and money for redaction). 
+
+![Cost and time estimation](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/costs_and_time.PNG)
+
+### Optional - cost code selection
+If the option is enabled (by your system admin, in the config file), you may be prompted to select a cost code before continuing with the redaction task.
+
+![Cost code selection](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/cost_code_selection.PNG)
+
+The relevant cost code can be found either by: 1. Using the search bar above the data table to find relevant cost codes, then clicking on the relevant row, or 2. typing it directly into the dropdown to the right, where it should filter as you type.
+
+### Optional - Submit whole documents to Textract API
+If this option is enabled (by your system admin, in the config file), you will have the option to submit whole documents in quick succession to the AWS Textract service to get extracted text outputs quickly (faster than using the 'Redact document' process described here). This feature is described in more detail in the [advanced user guide](#using-the-aws-textract-document-api).
+
+![Textract document API](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/textract_document_api.PNG)
+
+### Redact the document
 
-Hit 'Redact document'. After loading in the document, the app should be able to process about 30 pages per minute (depending on redaction methods chose above). When ready, you should see a message saying that processing is complete, with output files appearing in the bottom right.
+Click 'Redact document'. After loading in the document, the app should be able to process about 30 pages per minute (depending on redaction methods chose above). When ready, you should see a message saying that processing is complete, with output files appearing in the bottom right.
+
+### Redaction outputs
 
 ![Redaction outputs](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/redaction_outputs.PNG)
 
-- '...redacted.pdf' files contain the original pdf with suggested redacted text deleted and replaced by a black box on top of the document.
-- '...ocr_results.csv' files contain the line-by-line text outputs from the entire document. This file can be useful for later searching through for any terms of interest in the document (e.g. using Excel or a similar program).
-- '...review_file.csv' files are the review files that contain details and locations of all of the suggested redactions in the document. This file is key to the [review process](#reviewing-and-modifying-suggested-redactions), and should be downloaded to use later for this.
+- **'...redacted.pdf'** files contain the original pdf with suggested redacted text deleted and replaced by a black box on top of the document.
+- **'...ocr_results.csv'** files contain the line-by-line text outputs from the entire document. This file can be useful for later searching through for any terms of interest in the document (e.g. using Excel or a similar program).
+- **'...review_file.csv'** files are the review files that contain details and locations of all of the suggested redactions in the document. This file is key to the [review process](#reviewing-and-modifying-suggested-redactions), and should be downloaded to use later for this.
+
+### Additional AWS Textract outputs
+
+If you have used the AWS Textract option for extracting text, you may also see a '..._textract.json' file. This file contains all the relevant extracted text information that comes from the AWS Textract service. You can keep this file and upload it at a later date alongside your input document, which will enable you to skip calling AWS Textract every single time you want to do a redaction task, as follows:
+
+![Document upload alongside Textract](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/document_upload_with_textract.PNG)
 
-Additional outputs are available under the 'Redaction settings' tab. Scroll to the bottom and you should see more files: 
+### Downloading output files from previous redaction tasks
 
-![Additional processing outputs](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/redaction_additional_outputs.PNG)
+If you are logged in via AWS Cognito and you lose your app page for some reason (e.g. from a crash, reloading), it is possible recover your previous output files, provided the server has not been shut down since you redacted the document. Go to 'Redaction settings', then scroll to the bottom to see 'View all output files from this session'.
 
-- '...review_file.json' is the same file as the review file above, but in .json format.
-- '...decision_process_output.csv' is also similar to the review file above, with a few more details on the location and scores of identified PII in the document.
-- If you are using AWS Textract, you should also get a .json file with the Textract outputs. It could be useful to retain this document to avoid having to repeatedly analyse the same document in future (this .json file can be uploaded into the app on the first redaction tab to load into local memory before redaction).
+![View all output files](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/view_all_output_files.PNG)
+
+### Basic redaction summary
 
 We have covered redacting documents with the default redaction options. The '...redacted.pdf' file output may be enough for your purposes. But it is very likely that you will need to customise your redaction options, which we will cover below. 
 
@@ -126,6 +172,16 @@ There may be full pages in a document that you want to redact. The app also prov
 
 Using the above approaches to allow, deny, and full page redaction lists will give you an output [like this](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/allow_list/Partnership-Agreement-Toolkit_0_0_redacted.pdf).
 
+#### Adding to the loaded allow, deny, and whole page lists in-app
+
+If you open the accordion below the allow list options called 'Manually modify custom allow...', you should be able to see a few tables with options to add new rows:
+
+![Manually modify allow or deny list](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/allow_list/manually_modify.PNG)
+
+If the table is empty, you can add a new entry, you can add a new row by clicking on the '+' item below each table header. If there is existing data, you may need to click on the three dots to the right and select 'Add row below'. Type the item you wish to keep/remove in the cell, and then (important) press enter to add this new item to the allow/deny/whole page list. Your output tables should look something like below.
+
+![Manually modify allow or deny list filled](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/allow_list/manually_modify_filled.PNG)
+
 ### Redacting additional types of personal information
 
 You may want to redact additional types of information beyond the defaults, or you may not be interested in default suggested entity types. There are dates in the example complaint letter. Say we wanted to redact those dates also?
@@ -146,7 +202,9 @@ Say also we are only interested in redacting page 1 of the loaded documents. On
 
 ## Handwriting and signature redaction
 
-The file [Partnership Agreement Toolkit (for signatures and more advanced usage)](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/Partnership-Agreement-Toolkit_0_0.pdf) is provided as an example document to test AWS Textract + redaction with a document that has signatures in. If you have access to AWS Textract in the app, try removing all entity types from redaction on the Redaction settings and clicking the big X to the right of 'Entities to redact'. Ensure that handwriting and signatures are enabled for redaction on the Redaction Settings tab(enabled by default):
+The file [Partnership Agreement Toolkit (for signatures and more advanced usage)](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/Partnership-Agreement-Toolkit_0_0.pdf) is provided as an example document to test AWS Textract + redaction with a document that has signatures in. If you have access to AWS Textract in the app, try removing all entity types from redaction on the Redaction settings and clicking the big X to the right of 'Entities to redact'. 
+
+To ensure that handwriting and signatures are enabled (enabled by default), on the front screen go the 'AWS Textract signature detection' to enable/disable the following options :
 
 ![Handwriting and signatures](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/textract_handwriting_signatures.PNG)
 
@@ -156,72 +214,135 @@ The outputs should show handwriting/signatures redacted (see pages 5 - 7), which
 
 ## Reviewing and modifying suggested redactions
 
-Quite often there are certain terms suggested for redaction by the model that don't match quite what you intended. The app allows you to review and modify suggested redactions for the last file redacted. Refresh your browser tab. On the first tab 'PDFs/images' upload the ['Example of files sent to a professor before applying.pdf'](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/example_of_emails_sent_to_a_professor_before_applying.pdf) file. Let's stick with the 'Local model - selectable text' option, and click 'Redact document'. Once the outputs are created, go to the 'Review redactions' tab.
+Sometimes the app will suggest redactions that are incorrect, or will miss personal information entirely. The app allows you to review and modify suggested redactions to compensate for this. You can do this on the 'Review redactions' tab.
+
+We will go through ways to review suggested redactions with an example.On the first tab 'PDFs/images' upload the ['Example of files sent to a professor before applying.pdf'](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/example_of_emails_sent_to_a_professor_before_applying.pdf) file. Let's stick with the 'Local model - selectable text' option, and click 'Redact document'. Once the outputs are created, go to the 'Review redactions' tab.
 
-On this tab you have a visual interface that allows you to inspect and modify redactions suggested by the app. 
+On the 'Review redactions' tab you have a visual interface that allows you to inspect and modify redactions suggested by the app. There are quite a few options to look at, so we'll go from top to bottom.
 
 ![Review redactions](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/review_redactions.PNG)
 
+### Uploading documents for review
+
+The top area has a file upload area where you can upload original, unredacted PDFs, alongside the '..._review_file.csv' that is produced by the redaction process. Once you have uploaded these two files, click the 'Review PDF...' button to load in the files for review. This will allow you to visualise and modify the suggested redactions using the interface below.
+
+Optionally, you can also upload one of the '..._ocr_output.csv' files here that comes out of a redaction task, so that you can navigate the extracted text from the document.
+
+![Search extracted text](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/search_extracted_text.PNG)
+
+You can upload the three review files in the box (unredacted document, '..._review_file.csv' and '..._ocr_output.csv' file) before clicking 'Review PDF...', as in the image below:
+
+![Upload three files for review](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/upload_three_files.PNG)
+
+**NOTE:** ensure you upload the ***unredacted*** document here and not the redacted version, otherwise you will be checking over a document that already has redaction boxes applied!
+
+### Page navigation
+
 You can change the page viewed either by clicking 'Previous page' or 'Next page', or by typing a specific page number in the 'Current page' box and pressing Enter on your keyboard. Each time you switch page, it will save redactions you have made on the page you are moving from, so you will not lose changes you have made.
 
-On your selected page, each redaction is highlighted with a box next to its suggested entity type. By default the interface allows you to modify existing redaction boxes. Click and hold on an existing box to move it. Click on one of the small boxes at the edges to change the size of the box. To delete a box, click on it to highlight it, then press delete on your keyboard. Alternatively, double click on a box and click 'Remove' on the box that appears.
+You can also navigate to different pages by clicking on rows in the tables under 'Search suggested redactions' to the right, or 'search all extracted text' (if enabled) beneath that.
+
+### The document viewer pane
+
+On the selected page, each redaction is highlighted with a box next to its suggested redaction label (e.g. person, email).
+
+![Document view pane](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/document_viewer_pane.PNG)
 
-To change to 'add new redactions' mode, scroll to the bottom of the page. Click on the box icon, and your cursor will change into a crosshair. Now you can add new redaction boxes where you wish.
+There are a number of different options to add and modify redaction boxes and page on the document viewer pane. To zoom in and out of the page, use your mouse wheel. To move around the page while zoomed, you need to be in modify mode. Scroll to the bottom of the document viewer to see the relevant controls. You should see a box icon, a hand icon, and two arrows pointing counter-clockwise and clockwise.
 
 ![Change redaction mode](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/change_review_mode.PNG)
 
-On the right of the screen there is a dropdown and table where you can filter to entity types that have been found throughout the document. You can choose a specific entity type to see which pages the entity is present on. If you want to go to the page specified in the table, you can click on a cell in the table and the review page will be changed to that page.
+Click on the hand icon to go into modify mode. When you click and hold on the document viewer, This will allow you to move around the page when zoomed in. To rotate the page, you can click on either of the round arrow buttons to turn in that direction. 
+
+**NOTE:** When you switch page, the viewer will stay in your selected orientation, so if it looks strange, just rotate the page again and hopefully it will look correct!
+
+#### Modify existing redactions (hand icon)
+
+After clicking on the hand icon, the interface allows you to modify existing redaction boxes. When in this mode, you can click and hold on an existing box to move it. 
+
+![Modify existing redaction box](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/modify_existing_redaction_box.PNG)
+
+Click on one of the small boxes at the edges to change the size of the box. To delete a box, click on it to highlight it, then press delete on your keyboard. Alternatively, double click on a box and click 'Remove' on the box that appears.
+
+![Remove existing redaction box](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/existing_redaction_box_remove.PNG)
+
+#### Add new redaction boxes (box icon)
 
-![Change redaction mode](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/list_find_labels.PNG)
+To change to 'add redaction boxes' mode, scroll to the bottom of the page. Click on the box icon, and your cursor will change into a crosshair. Now you can add new redaction boxes where you wish. A popup will appear when you create a new box so you can select a label and colour for the new box.
 
-Note that the table currently only shows entity types, and not specific found text. So for instance if you provide a list of specific terms to redact in the [deny list](#deny-list-example), they will all be labelled just as 'CUSTOM'. A feature to include in the near term will include being able to view specific redacted text in this table to get a better sense of the PII entities found.
+#### 'Locking in' new redaction box format
 
-Once you happy with your modified changes throughout the document, click 'Apply revised redactions' at the top of the page. The app will then run through all the pages in the document to update the redactions, and will output a modified PDF file. The modified PDF will appear at the top of the page in the file area. It will also output a revised '...review_file.csv' that you can then use for future review tasks.
+It is possible to lock in a chosen format for new redaction boxes so that you don't have the popup appearing each time. When you make a new box, select the options for your 'locked' format, and then click on the lock icon on the left side of the popup, which should turn blue.
+
+![Lock redaction box format](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/new_redaction_box_lock_mode.PNG)
+
+You can now add new redaction boxes without a popup appearing. If you want to change or 'unlock' the your chosen box format, you can click on the new icon that has appeared at the bottom of the document viewer pane that looks a little like a gift tag. You can then change the defaults, or click on the lock icon again to 'unlock' the new box format - then popups will appear again each time you create a new box.
+
+![Change or unlock redaction box format](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/change_review_mode_with_lock.PNG)
+
+### Apply redactions to PDF and Save changes on current page
+
+Once you have reviewed all the redactions in your document and you are happy with the outputs, you can click 'Apply revised redactions to PDF' to create a new '_redacted.pdf' output alongside a new '_review_file.csv' output.
+
+If you are working on a page and haven't saved for a while, you can click 'Save changes on current page to file' to ensure that they are saved to an updated 'review_file.csv' output.
 
 ![Review modified outputs](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/review_mod_outputs.PNG)
 
-Any feedback or comments on the app, please get in touch!
+### Selecting and removing redaction boxes using the 'Search suggested redactions' table
 
-# ADVANCED USER GUIDE
+The table shows a list of all the suggested redactions in the document alongside the page, label, and text (if available).
 
-This advanced user guide will go over some of the features recently added to the app, including: modifying and merging redaction review files, identifying and redacting duplicate pages across multiple PDFs, 'fuzzy' search and redact, and exporting redactions to Adobe Acrobat.
+![Search suggested redaction area](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/list_find_labels.PNG)
 
-## Table of contents
+If you click on one of the rows in this table, you will be taken to the page of the redaction. Clicking on a redaction row on the same page *should* change the colour of redaction box to blue to help you locate it in the document viewer (just in app, not in redaction output PDFs).
 
-- [Modifying and merging redaction review files](#modifying-and-merging-redaction-review-files)
-    - [Modifying existing redaction review files](#modifying-existing-redaction-review-files)
-    - [Merging existing redaction review files](#merging-existing-redaction-review-files)
-- [Identifying and redacting duplicate pages](#identifying-and-redacting-duplicate-pages)
-- [Fuzzy search and redaction](#fuzzy-search-and-redaction)
-- [Export redactions to and import from Adobe Acrobat](#export-to-and-import-from-adobe)
-    - [Exporting to Adobe Acrobat](#exporting-to-adobe-acrobat)
-    - [Importing from Adobe Acrobat](#importing-from-adobe-acrobat)
+![Search suggested redaction area](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/review_row_highlight.PNG)
 
+You can choose a specific entity type to see which pages the entity is present on. If you want to go to the page specified in the table, you can click on a cell in the table and the review page will be changed to that page.
 
-## Modifying and merging redaction review files
+To filter the 'Search suggested redactions' table you can:
+1. Click on one of the dropdowns (Redaction category, Page, Text), and select an option, or
+2. Write text in the 'Filter' box just above the table. Click the blue box to apply the filter to the table.
 
-You can find the folder containing the files discussed in this section [here](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/merge_review_files/).
+Once you have filtered the table, you have a few options underneath on what you can do with the filtered rows:
 
-As well as serving as inputs to the document redaction app's review function, the 'review_file.csv' output can be modified outside of the app, and also merged with others from multiple redaction attempts on the same file. This gives you the flexibility to change redaction details outside of the app.
+- Click the 'Exclude specific row from redactions' button to remove only the redaction from the last row you clicked on from the document. 
+- Click the 'Exclude all items in table from redactions' button to remove all redactions visible in the table from the document. 
 
-### Modifying existing redaction review files
-If you open up a 'review_file' csv output using a spreadsheet software program such as Microsoft Excel you can easily modify redaction properties. Open the file '[Partnership-Agreement-Toolkit_0_0_redacted.pdf_review_file_local.csv](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/merge_review_files/Partnership-Agreement-Toolkit_0_0.pdf_review_file_local.csv)', and you should see a spreadshet with just four suggested redactions (see below). The following instructions are for using Excel.
+**NOTE**: After excluding redactions using either of the above options, click the 'Reset filters' button below to ensure that the dropdowns and table return to seeing all remaining redactions in the document.
 
-![Review file before](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/merge_review_files/img/review_file_before.PNG)
+If you made a mistake, click the 'Undo last element removal' button to restore the Search suggested redactions table to its previous state (can only undo the last action).
 
-The first thing we can do is remove the first row - 'et' is suggested as a person, but is obviously not a genuine instance of personal information. Right click on the row number and select delete on this menu. Next, let's imagine that what the app identified as a 'phone number' was in fact another type of number and so we wanted to change the label. Simply click on the relevant label cells, let's change it to 'SECURITY_NUMBER'. You could also use 'Finad & Select' -> 'Replace' from the top ribbon menu if you wanted to change a number of labels simultaneously.
+### Navigating through the document using the 'Search all extracted text' 
 
-How about we wanted to change the colour of the 'email address' entry on the redaction review tab of the redaction app? The colours in a review file are based on an RGB scale with three numbers ranging from 0-255. [You can find suitable colours here](https://rgbcolorpicker.com). Using this scale, if I wanted my review box to be pure blue, I can change the cell value to (0,0,255).
+The 'search all extracted text' table will contain text if you have just redacted a document, or if you have uploaded a '..._ocr_output.csv' file alongside a document file and review file on the Review redactions tab as [described above](#uploading-documents-for-review). 
 
-Imagine that a redaction box was slightly too small, and I didn't want to use the in-app options to change the size. In the review file csv, we can modify e.g. the ymin and ymax values for any box to increase the extent of the redaction box. For the 'email address' entry, let's decrease ymin by 5, and increase ymax by 5.
+You can navigate through the document using this table. When you click on a row, the Document viewer pane to the left will change to the selected page. 
 
-I have saved an output file following the above steps as '[Partnership-Agreement-Toolkit_0_0_redacted.pdf_review_file_local_mod.csv](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/merge_review_files/outputs/Partnership-Agreement-Toolkit_0_0.pdf_review_file_local_mod.csv)' in the same folder that the original was found. Let's upload this file to the app along with the original pdf to see how the redactions look now.
+![Search suggested redaction area](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/select_extracted_text.PNG)
 
-![Review file after modification](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/merge_review_files/img/partnership_redactions_after.PNG)
+You can search through the extracted text by using the search bar just above the table, which should filter as you type. To apply the filter and 'cut' the table, click on the blue tick inside the box next to your search term. To return the table to its original content, click the button below the table 'Reset OCR output table filter'.
+
+![Search suggested redaction area](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/review_redactions/search_extracted_text.PNG)
+
+# ADVANCED USER GUIDE
+
+This advanced user guide will go over some of the features recently added to the app, including: modifying and merging redaction review files, identifying and redacting duplicate pages across multiple PDFs, 'fuzzy' search and redact, and exporting redactions to Adobe Acrobat.
 
-We can see from the above that we have successfully removed a redaction box, changed labels, colours, and redaction box sizes.
+## Table of contents
 
-### Merging existing redaction review files
+- [Merging redaction review files](#merging-redaction-review-files)
+- [Identifying and redacting duplicate pages](#identifying-and-redacting-duplicate-pages)
+- [Fuzzy search and redaction](#fuzzy-search-and-redaction)
+- [Export redactions to and import from Adobe Acrobat](#export-to-and-import-from-adobe)
+    - [Exporting to Adobe Acrobat](#exporting-to-adobe-acrobat)
+    - [Importing from Adobe Acrobat](#importing-from-adobe-acrobat)
+- [Using the AWS Textract document API](#using-the-aws-textract-document-api)
+- [Using AWS Textract and Comprehend when not running in an AWS environment](#using-aws-textract-and-comprehend-when-not-running-in-an-aws-environment)
+- [Modifying existing redaction review files](#modifying-existing-redaction-review-files)
+
+
+## Merging redaction review files
 
 Say you have run multiple redaction tasks on the same document, and you want to merge all of these redactions together. You could do this in your spreadsheet editor, but this could be fiddly especially if dealing with multiple review files or large numbers of redactions. The app has a feature to combine multiple review files together to create a 'merged' review file.
 
@@ -303,6 +424,30 @@ When you click the 'convert .xfdf comment file to review_file.csv' button, the a
 
 ![Outputs from Adobe import](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/export_to_adobe/img/import_from_adobe_interface_outputs.PNG)
 
+## Using the AWS Textract document API
+
+This option can be enabled by your system admin, in the config file ('SHOW_BULK_TEXTRACT_CALL_OPTIONS' environment variable, and subsequent variables). Using this, you will have the option to submit whole documents in quick succession to the AWS Textract service to get extracted text outputs quickly (faster than using the 'Redact document' process described here).
+
+### Starting a new Textract API job
+
+To use this feature, first upload a document file in the file input box [in the usual way](#upload-files-to-the-app) on the first tab of the app. Under AWS Textract signature detection you can select whether or not you would like to analyse signatures or not (with a [cost implication](#optional---select-signature-extraction)).
+
+Then, open the section under the heading 'Submit whole document to AWS Textract API...'.
+
+![Textract document API menu](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/quick_start/textract_document_api.PNG)
+
+Click 'Analyse document with AWS Textract API call'. After a few seconds, the job should be submitted to the AWS Textract service. The box 'Job ID to check status' should now have an ID filled in. If it is not already filled with previous jobs (up to seven days old), the table should have a row added with details of the new API job.
+
+Click the button underneath, 'Check status of Textract job and download', to see progress on the job. Processing will continue in the background until the job is ready, so it is worth periodically clicking this button to see if the outputs are ready. In testing, and as a rough estimate, it seems like this process takes about five seconds per page. However, this has not been tested with very large documents. Once ready, the '_textract.json' output should appear below.
+
+### Textract API job outputs
+
+The '_textract.json' output can be used to speed up further redaction tasks as [described previously](#optional---costs-and-time-estimation), the 'Existing Textract output file found' flag should now be ticked.
+
+![Textract document API initial ouputs](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/textract_api/textract_api_initial_outputs.PNG)
+
+You can now easily get the '..._ocr_output.csv' redaction output based on this '_textract.json' (described in [Redaction outputs](#redaction-outputs)) by clicking on the button 'Convert Textract job outputs to OCR results'. You can now use this file e.g. for [identifying duplicate pages](#identifying-and-redacting-duplicate-pages), or for redaction review.
+
 ## Using AWS Textract and Comprehend when not running in an AWS environment
 
 AWS Textract and Comprehend give much better results for text extraction and document redaction than the local model options in the app. The most secure way to access them in the Redaction app is to run the app in a secure AWS environment with relevant permissions. Alternatively, you could run the app on your own system while logged in to AWS SSO with relevant permissions.
@@ -322,4 +467,27 @@ AWS_SECRET_KEY= your-secret-key
 
 The app should then pick up these keys when trying to access the AWS Textract and Comprehend services during redaction.
 
-Again, a lot can potentially go wrong with AWS solutions that are insecure, so before trying the above please consult with your AWS and data security teams.
+Again, a lot can potentially go wrong with AWS solutions that are insecure, so before trying the above please consult with your AWS and data security teams.
+
+## Modifying and merging redaction review files
+
+You can find the folder containing the files discussed in this section [here](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/merge_review_files/).
+
+As well as serving as inputs to the document redaction app's review function, the 'review_file.csv' output can be modified outside of the app, and also merged with others from multiple redaction attempts on the same file. This gives you the flexibility to change redaction details outside of the app.
+
+### Modifying existing redaction review files
+If you open up a 'review_file' csv output using a spreadsheet software program such as Microsoft Excel you can easily modify redaction properties. Open the file '[Partnership-Agreement-Toolkit_0_0_redacted.pdf_review_file_local.csv](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/merge_review_files/Partnership-Agreement-Toolkit_0_0.pdf_review_file_local.csv)', and you should see a spreadshet with just four suggested redactions (see below). The following instructions are for using Excel.
+
+![Review file before](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/merge_review_files/img/review_file_before.PNG)
+
+The first thing we can do is remove the first row - 'et' is suggested as a person, but is obviously not a genuine instance of personal information. Right click on the row number and select delete on this menu. Next, let's imagine that what the app identified as a 'phone number' was in fact another type of number and so we wanted to change the label. Simply click on the relevant label cells, let's change it to 'SECURITY_NUMBER'. You could also use 'Finad & Select' -> 'Replace' from the top ribbon menu if you wanted to change a number of labels simultaneously.
+
+How about we wanted to change the colour of the 'email address' entry on the redaction review tab of the redaction app? The colours in a review file are based on an RGB scale with three numbers ranging from 0-255. [You can find suitable colours here](https://rgbcolorpicker.com). Using this scale, if I wanted my review box to be pure blue, I can change the cell value to (0,0,255).
+
+Imagine that a redaction box was slightly too small, and I didn't want to use the in-app options to change the size. In the review file csv, we can modify e.g. the ymin and ymax values for any box to increase the extent of the redaction box. For the 'email address' entry, let's decrease ymin by 5, and increase ymax by 5.
+
+I have saved an output file following the above steps as '[Partnership-Agreement-Toolkit_0_0_redacted.pdf_review_file_local_mod.csv](https://github.com/seanpedrick-case/document_redaction_examples/blob/main/merge_review_files/outputs/Partnership-Agreement-Toolkit_0_0.pdf_review_file_local_mod.csv)' in the same folder that the original was found. Let's upload this file to the app along with the original pdf to see how the redactions look now.
+
+![Review file after modification](https://raw.githubusercontent.com/seanpedrick-case/document_redaction_examples/main/merge_review_files/img/partnership_redactions_after.PNG)
+
+We can see from the above that we have successfully removed a redaction box, changed labels, colours, and redaction box sizes.

app.py

@@ -4,11 +4,11 @@ import pandas as pd
 import gradio as gr
 from gradio_image_annotation import image_annotator
 
-from tools.config import OUTPUT_FOLDER, INPUT_FOLDER, RUN_DIRECT_MODE, MAX_QUEUE_SIZE, DEFAULT_CONCURRENCY_LIMIT, MAX_FILE_SIZE, GRADIO_SERVER_PORT, ROOT_PATH, GET_DEFAULT_ALLOW_LIST, ALLOW_LIST_PATH, S3_ALLOW_LIST_PATH, FEEDBACK_LOGS_FOLDER, ACCESS_LOGS_FOLDER, USAGE_LOGS_FOLDER, TESSERACT_FOLDER, POPPLER_FOLDER, REDACTION_LANGUAGE, GET_COST_CODES, COST_CODES_PATH, S3_COST_CODES_PATH, ENFORCE_COST_CODES, DISPLAY_FILE_NAMES_IN_LOGS, SHOW_COSTS, RUN_AWS_FUNCTIONS, DOCUMENT_REDACTION_BUCKET, SHOW_BULK_TEXTRACT_CALL_OPTIONS, TEXTRACT_BULK_ANALYSIS_BUCKET, TEXTRACT_BULK_ANALYSIS_INPUT_SUBFOLDER, TEXTRACT_BULK_ANALYSIS_OUTPUT_SUBFOLDER, SESSION_OUTPUT_FOLDER, LOAD_PREVIOUS_TEXTRACT_JOBS_S3, TEXTRACT_JOBS_S3_LOC, TEXTRACT_JOBS_LOCAL_LOC, HOST_NAME, DEFAULT_COST_CODE, OUTPUT_COST_CODES_PATH, OUTPUT_ALLOW_LIST_PATH
+from tools.config import OUTPUT_FOLDER, INPUT_FOLDER, RUN_DIRECT_MODE, MAX_QUEUE_SIZE, DEFAULT_CONCURRENCY_LIMIT, MAX_FILE_SIZE, GRADIO_SERVER_PORT, ROOT_PATH, GET_DEFAULT_ALLOW_LIST, ALLOW_LIST_PATH, S3_ALLOW_LIST_PATH, FEEDBACK_LOGS_FOLDER, ACCESS_LOGS_FOLDER, USAGE_LOGS_FOLDER, TESSERACT_FOLDER, POPPLER_FOLDER, REDACTION_LANGUAGE, GET_COST_CODES, COST_CODES_PATH, S3_COST_CODES_PATH, ENFORCE_COST_CODES, DISPLAY_FILE_NAMES_IN_LOGS, SHOW_COSTS, RUN_AWS_FUNCTIONS, DOCUMENT_REDACTION_BUCKET, SHOW_BULK_TEXTRACT_CALL_OPTIONS, TEXTRACT_BULK_ANALYSIS_BUCKET, TEXTRACT_BULK_ANALYSIS_INPUT_SUBFOLDER, TEXTRACT_BULK_ANALYSIS_OUTPUT_SUBFOLDER, SESSION_OUTPUT_FOLDER, LOAD_PREVIOUS_TEXTRACT_JOBS_S3, TEXTRACT_JOBS_S3_LOC, TEXTRACT_JOBS_LOCAL_LOC, HOST_NAME, DEFAULT_COST_CODE, OUTPUT_COST_CODES_PATH, OUTPUT_ALLOW_LIST_PATH, COGNITO_AUTH
 from tools.helper_functions import put_columns_in_df, get_connection_params, reveal_feedback_buttons, custom_regex_load, reset_state_vars, load_in_default_allow_list, tesseract_ocr_option, text_ocr_option, textract_option, local_pii_detector, aws_pii_detector, no_redaction_option, reset_review_vars, merge_csv_files, load_all_output_files, update_dataframe, check_for_existing_textract_file, load_in_default_cost_codes, enforce_cost_codes, calculate_aws_costs, calculate_time_taken, reset_base_dataframe, reset_ocr_base_dataframe, update_cost_code_dataframe_from_dropdown_select
 from tools.aws_functions import upload_file_to_s3, download_file_from_s3
 from tools.file_redaction import choose_and_run_redactor
-from tools.file_conversion import prepare_image_or_pdf, get_input_file_names, convert_review_df_to_annotation_json
+from tools.file_conversion import prepare_image_or_pdf, get_input_file_names
 from tools.redaction_review import apply_redactions_to_review_df_and_files, update_all_page_annotation_object_based_on_previous_page, decrease_page, increase_page, update_annotator_object_and_filter_df, update_entities_df_recogniser_entities, update_entities_df_page, update_entities_df_text, df_select_callback, convert_df_to_xfdf, convert_xfdf_to_dataframe, reset_dropdowns, exclude_selected_items_from_redaction, undo_last_removal, update_selected_review_df_row_colour, update_all_entity_df_dropdowns, df_select_callback_cost, update_other_annotator_number_from_current, update_annotator_page_from_review_df, df_select_callback_ocr, df_select_callback_textract_api
 from tools.data_anonymise import anonymise_data_files
 from tools.auth import authenticate_user
@@ -572,9 +572,9 @@ with app:
     text_entity_dropdown.select(update_entities_df_text, inputs=[text_entity_dropdown, recogniser_entity_dataframe_base, recogniser_entity_dropdown, page_entity_dropdown], outputs=[recogniser_entity_dataframe, recogniser_entity_dropdown, page_entity_dropdown])
 
     # Clicking on a cell in the recogniser entity dataframe will take you to that page, and also highlight the target redaction box in blue
-    recogniser_entity_dataframe.select(df_select_callback, inputs=[recogniser_entity_dataframe], outputs=[annotate_current_page, selected_entity_dataframe_row]).\
-        success(update_selected_review_df_row_colour, inputs=[selected_entity_dataframe_row, review_file_state, selected_entity_id, selected_entity_colour, page_sizes], outputs=[review_file_state, selected_entity_id, selected_entity_colour]).\
-        success(update_annotator_page_from_review_df, inputs=[review_file_state, images_pdf_state, page_sizes, annotate_current_page, annotate_previous_page, all_image_annotations_state, annotator], outputs=[annotator, all_image_annotations_state])
+    recogniser_entity_dataframe.select(df_select_callback, inputs=[recogniser_entity_dataframe], outputs=[selected_entity_dataframe_row]).\
+        success(update_selected_review_df_row_colour, inputs=[selected_entity_dataframe_row, review_file_state, selected_entity_id, selected_entity_colour], outputs=[review_file_state, selected_entity_id, selected_entity_colour]).\
+        success(update_annotator_page_from_review_df, inputs=[review_file_state, images_pdf_state, page_sizes, all_image_annotations_state, annotator, selected_entity_dataframe_row, input_folder_textbox, doc_full_file_name_textbox], outputs=[annotator, all_image_annotations_state, annotate_current_page, page_sizes, review_file_state, annotate_previous_page])
    
     reset_dropdowns_btn.click(reset_dropdowns, inputs=[recogniser_entity_dataframe_base], outputs=[recogniser_entity_dropdown, text_entity_dropdown, page_entity_dropdown]).\
         success(update_annotator_object_and_filter_df, inputs=[all_image_annotations_state, annotate_current_page, recogniser_entity_dropdown, page_entity_dropdown, text_entity_dropdown, recogniser_entity_dataframe_base, annotator_zoom_number, review_file_state, page_sizes, doc_full_file_name_textbox, input_folder_textbox], outputs = [annotator, annotate_current_page, annotate_current_page_bottom, annotate_previous_page, recogniser_entity_dropdown, recogniser_entity_dataframe, recogniser_entity_dataframe_base, text_entity_dropdown, page_entity_dropdown, page_sizes, all_image_annotations_state])
@@ -733,7 +733,7 @@ with app:
 if __name__ == "__main__":
     if RUN_DIRECT_MODE == "0":
         
-        if os.environ['COGNITO_AUTH'] == "1":
+        if COGNITO_AUTH == "1":
             app.queue(max_size=int(MAX_QUEUE_SIZE), default_concurrency_limit=int(DEFAULT_CONCURRENCY_LIMIT)).launch(show_error=True, inbrowser=True, auth=authenticate_user, max_file_size=MAX_FILE_SIZE, server_port=GRADIO_SERVER_PORT, root_path=ROOT_PATH)
         else:
             app.queue(max_size=int(MAX_QUEUE_SIZE), default_concurrency_limit=int(DEFAULT_CONCURRENCY_LIMIT)).launch(show_error=True, inbrowser=True, max_file_size=MAX_FILE_SIZE, server_port=GRADIO_SERVER_PORT, root_path=ROOT_PATH)

tools/config.py

@@ -237,7 +237,7 @@ else: OUTPUT_ALLOW_LIST_PATH = 'config/default_allow_list.csv'
 
 SHOW_COSTS = get_or_create_env_var('SHOW_COSTS', 'False')
 
-GET_COST_CODES = get_or_create_env_var('GET_COST_CODES', 'False')
+GET_COST_CODES = get_or_create_env_var('GET_COST_CODES', 'True')
 
 DEFAULT_COST_CODE = get_or_create_env_var('DEFAULT_COST_CODE', '')
 
@@ -246,7 +246,7 @@ COST_CODES_PATH = get_or_create_env_var('COST_CODES_PATH', '') # 'config/COST_CE
 S3_COST_CODES_PATH = get_or_create_env_var('S3_COST_CODES_PATH', '') # COST_CENTRES.csv # This is a path within the DOCUMENT_REDACTION_BUCKET
 
 if COST_CODES_PATH: OUTPUT_COST_CODES_PATH = COST_CODES_PATH
-else: OUTPUT_COST_CODES_PATH = 'config/COST_CENTRES.csv'
+else: OUTPUT_COST_CODES_PATH = ''
 
 ENFORCE_COST_CODES = get_or_create_env_var('ENFORCE_COST_CODES', 'False') # If you have cost codes listed, is it compulsory to choose one before redacting?
 

tools/file_conversion.py

@@ -21,6 +21,7 @@ from PIL import Image
 from scipy.spatial import cKDTree
 import random
 import string
+import warnings # To warn about potential type changes
 
 IMAGE_NUM_REGEX = re.compile(r'_(\d+)\.png$')
 
@@ -617,11 +618,10 @@ def prepare_image_or_pdf(
 
         elif file_extension in ['.csv']:
             if '_review_file' in file_path_without_ext:
-                #print("file_path:", file_path)
                 review_file_csv = read_file(file_path)
                 all_annotations_object = convert_review_df_to_annotation_json(review_file_csv, image_file_paths, page_sizes)
                 json_from_csv = True
-                print("Converted CSV review file to image annotation object")
+                #print("Converted CSV review file to image annotation object")
             elif '_ocr_output' in file_path_without_ext:
                 all_line_level_ocr_results_df = read_file(file_path)
                 json_from_csv = False
@@ -850,121 +850,246 @@ def remove_duplicate_images_with_blank_boxes(data: List[dict]) -> List[dict]:
 
     return result
 
-def divide_coordinates_by_page_sizes(review_file_df:pd.DataFrame, page_sizes_df:pd.DataFrame, xmin="xmin", xmax="xmax", ymin="ymin", ymax="ymax"):
-
-    '''Convert data to same coordinate system. If all coordinates all greater than one, this is a absolute image coordinates - change back to relative coordinates.'''
-
-    review_file_df_out = review_file_df
-
-    if xmin in review_file_df.columns and not review_file_df.empty:
-        coord_cols = [xmin, xmax, ymin, ymax]
-        for col in coord_cols:
-            review_file_df.loc[:, col] = pd.to_numeric(review_file_df[col], errors="coerce")
-
-        review_file_df_orig = review_file_df.copy().loc[(review_file_df[xmin] <= 1) & (review_file_df[xmax] <= 1) & (review_file_df[ymin] <= 1) & (review_file_df[ymax] <= 1),:]
-
-        #print("review_file_df_orig:", review_file_df_orig)
-        
-        review_file_df_div = review_file_df.loc[(review_file_df[xmin] > 1) & (review_file_df[xmax] > 1) & (review_file_df[ymin] > 1) & (review_file_df[ymax] > 1),:]
-
-        #print("review_file_df_div:", review_file_df_div)
-
-        review_file_df_div.loc[:, "page"] = pd.to_numeric(review_file_df_div["page"], errors="coerce")
+def divide_coordinates_by_page_sizes(
+    review_file_df: pd.DataFrame,
+    page_sizes_df: pd.DataFrame,
+    xmin="xmin", xmax="xmax", ymin="ymin", ymax="ymax"
+) -> pd.DataFrame:
+    """
+    Optimized function to convert absolute image coordinates (>1) to relative coordinates (<=1).
 
-        if "image_width" not in review_file_df_div.columns and not page_sizes_df.empty:  
+    Identifies rows with absolute coordinates, merges page size information,
+    divides coordinates by dimensions, and combines with already-relative rows.
 
-            page_sizes_df["image_width"] = page_sizes_df["image_width"].replace("<NA>", pd.NA)
-            page_sizes_df["image_height"] = page_sizes_df["image_height"].replace("<NA>", pd.NA)
-            review_file_df_div = review_file_df_div.merge(page_sizes_df[["page", "image_width", "image_height", "mediabox_width", "mediabox_height"]], on="page", how="left")
+    Args:
+        review_file_df: Input DataFrame with potentially mixed coordinate systems.
+        page_sizes_df: DataFrame with page dimensions ('page', 'image_width',
+                       'image_height', 'mediabox_width', 'mediabox_height').
+        xmin, xmax, ymin, ymax: Names of the coordinate columns.
 
-        if "image_width" in review_file_df_div.columns:
-            if review_file_df_div["image_width"].isna().all():  # Check if all are NaN values. If so, assume we only have mediabox coordinates available
-                review_file_df_div["image_width"] = review_file_df_div["image_width"].fillna(review_file_df_div["mediabox_width"]).infer_objects()
-                review_file_df_div["image_height"] = review_file_df_div["image_height"].fillna(review_file_df_div["mediabox_height"]).infer_objects()
+    Returns:
+        DataFrame with coordinates converted to relative system, sorted.
+    """
+    if review_file_df.empty or xmin not in review_file_df.columns:
+        return review_file_df # Return early if empty or key column missing
 
-            convert_type_cols = ["image_width", "image_height", xmin, xmax, ymin, ymax]
-            review_file_df_div[convert_type_cols] = review_file_df_div[convert_type_cols].apply(pd.to_numeric, errors="coerce")  
+    # --- Initial Type Conversion ---
+    coord_cols = [xmin, xmax, ymin, ymax]
+    cols_to_convert = coord_cols + ["page"]
+    temp_df = review_file_df.copy() # Work on a copy initially
 
-            review_file_df_div[xmin] = review_file_df_div[xmin] / review_file_df_div["image_width"]
-            review_file_df_div[xmax] = review_file_df_div[xmax] / review_file_df_div["image_width"]
-            review_file_df_div[ymin] = review_file_df_div[ymin] / review_file_df_div["image_height"]
-            review_file_df_div[ymax] = review_file_df_div[ymax] / review_file_df_div["image_height"]
+    for col in cols_to_convert:
+        if col in temp_df.columns:
+            temp_df[col] = pd.to_numeric(temp_df[col], errors="coerce")
+        else:
+            # If essential 'page' or coord column missing, cannot proceed meaningfully
+            if col == 'page' or col in coord_cols:
+                 print(f"Warning: Required column '{col}' not found in review_file_df. Returning original DataFrame.")
+                 return review_file_df
+
+    # --- Identify Absolute Coordinates ---
+    # Create mask for rows where *all* coordinates are potentially absolute (> 1)
+    # Handle potential NaNs introduced by to_numeric - treat NaN as not absolute.
+    is_absolute_mask = (
+        (temp_df[xmin] > 1) & (temp_df[xmin].notna()) &
+        (temp_df[xmax] > 1) & (temp_df[xmax].notna()) &
+        (temp_df[ymin] > 1) & (temp_df[ymin].notna()) &
+        (temp_df[ymax] > 1) & (temp_df[ymax].notna())
+    )
 
-        # Concatenate the original and modified DataFrames
-        dfs_to_concat = [df for df in [review_file_df_orig, review_file_df_div] if not df.empty]
-        if dfs_to_concat:  # Ensure there's at least one non-empty DataFrame
-            review_file_df_out = pd.concat(dfs_to_concat)
+    # --- Separate DataFrames ---
+    df_rel = temp_df[~is_absolute_mask] # Rows already relative or with NaN/mixed coords
+    df_abs = temp_df[is_absolute_mask].copy() # Absolute rows - COPY here to allow modifications
+
+    # --- Process Absolute Coordinates ---
+    if not df_abs.empty:
+        # Merge page sizes if necessary
+        if "image_width" not in df_abs.columns and not page_sizes_df.empty:
+            ps_df_copy = page_sizes_df.copy() # Work on a copy of page sizes
+
+            # Ensure page is numeric for merge key matching
+            ps_df_copy['page'] = pd.to_numeric(ps_df_copy['page'], errors='coerce')
+
+            # Columns to merge from page_sizes
+            merge_cols = ['page', 'image_width', 'image_height', 'mediabox_width', 'mediabox_height']
+            available_merge_cols = [col for col in merge_cols if col in ps_df_copy.columns]
+
+            # Prepare dimension columns in the copy
+            for col in ['image_width', 'image_height', 'mediabox_width', 'mediabox_height']:
+                 if col in ps_df_copy.columns:
+                     # Replace "<NA>" string if present
+                     if ps_df_copy[col].dtype == 'object':
+                          ps_df_copy[col] = ps_df_copy[col].replace("<NA>", pd.NA)
+                     # Convert to numeric
+                     ps_df_copy[col] = pd.to_numeric(ps_df_copy[col], errors='coerce')
+
+            # Perform the merge
+            if 'page' in available_merge_cols: # Check if page exists for merging
+                df_abs = df_abs.merge(
+                    ps_df_copy[available_merge_cols],
+                    on="page",
+                    how="left"
+                )
+            else:
+                 print("Warning: 'page' column not found in page_sizes_df. Cannot merge dimensions.")
+
+
+        # Fallback to mediabox dimensions if image dimensions are missing
+        if "image_width" in df_abs.columns and "mediabox_width" in df_abs.columns:
+             # Check if image_width mostly missing - use .isna().all() or check percentage
+             if df_abs["image_width"].isna().all():
+                 print("Falling back to mediabox dimensions as image_width is entirely missing.")
+                 df_abs["image_width"] = df_abs["image_width"].fillna(df_abs["mediabox_width"])
+                 df_abs["image_height"] = df_abs["image_height"].fillna(df_abs["mediabox_height"])
+             else:
+                  # Optional: Fill only missing image dims if some exist?
+                  # df_abs["image_width"].fillna(df_abs["mediabox_width"], inplace=True)
+                  # df_abs["image_height"].fillna(df_abs["mediabox_height"], inplace=True)
+                  pass # Current logic only falls back if ALL image_width are NaN
+
+        # Ensure divisor columns are numeric before division
+        divisors_numeric = True
+        for col in ["image_width", "image_height"]:
+            if col in df_abs.columns:
+                 df_abs[col] = pd.to_numeric(df_abs[col], errors='coerce')
+            else:
+                 print(f"Warning: Dimension column '{col}' missing. Cannot perform division.")
+                 divisors_numeric = False
+
+
+        # Perform division if dimensions are available and numeric
+        if divisors_numeric and "image_width" in df_abs.columns and "image_height" in df_abs.columns:
+             # Use np.errstate to suppress warnings about division by zero or NaN if desired
+             with np.errstate(divide='ignore', invalid='ignore'):
+                df_abs[xmin] = df_abs[xmin] / df_abs["image_width"]
+                df_abs[xmax] = df_abs[xmax] / df_abs["image_width"]
+                df_abs[ymin] = df_abs[ymin] / df_abs["image_height"]
+                df_abs[ymax] = df_abs[ymax] / df_abs["image_height"]
+                # Replace potential infinities with NaN (optional, depending on desired outcome)
+                df_abs.replace([np.inf, -np.inf], np.nan, inplace=True)
         else:
-            review_file_df_out = review_file_df  # Return an original DataFrame instead of raising an error
+             print("Skipping coordinate division due to missing or non-numeric dimension columns.")
 
-        # Only sort if the DataFrame is not empty and contains the required columns
-        required_sort_columns = {"page", xmin, ymin}
-        if not review_file_df_out.empty and required_sort_columns.issubset(review_file_df_out.columns):
-            review_file_df_out.sort_values(["page", ymin, xmin], inplace=True)
 
-    review_file_df_out.drop(["image_width", "image_height", "mediabox_width", "mediabox_height"], axis=1, errors="ignore")
+    # --- Combine Relative and Processed Absolute DataFrames ---
+    dfs_to_concat = [df for df in [df_rel, df_abs] if not df.empty]
 
-    return review_file_df_out
+    if dfs_to_concat:
+        final_df = pd.concat(dfs_to_concat, ignore_index=True)
+    else:
+        # If both splits were empty, return an empty DF with original columns
+        print("Warning: Both relative and absolute splits resulted in empty DataFrames.")
+        final_df = pd.DataFrame(columns=review_file_df.columns)
 
-def multiply_coordinates_by_page_sizes(review_file_df: pd.DataFrame, page_sizes_df: pd.DataFrame, xmin="xmin", xmax="xmax", ymin="ymin", ymax="ymax"):
 
+    # --- Final Sort ---
+    required_sort_columns = {"page", xmin, ymin}
+    if not final_df.empty and required_sort_columns.issubset(final_df.columns):
+        # Ensure sort columns are numeric before sorting
+        final_df['page'] = pd.to_numeric(final_df['page'], errors='coerce')
+        final_df[ymin] = pd.to_numeric(final_df[ymin], errors='coerce')
+        final_df[xmin] = pd.to_numeric(final_df[xmin], errors='coerce')
+        # Sort by page, ymin, xmin (note order compared to multiply function)
+        final_df.sort_values(["page", ymin, xmin], inplace=True, na_position='last')
 
-    if xmin in review_file_df.columns and not review_file_df.empty:
 
-        coord_cols = [xmin, xmax, ymin, ymax]
-        for col in coord_cols:
-            review_file_df.loc[:, col] = pd.to_numeric(review_file_df[col], errors="coerce")
+    # --- Clean Up Columns ---
+    # Correctly drop columns and reassign the result
+    cols_to_drop = ["image_width", "image_height", "mediabox_width", "mediabox_height"]
+    final_df = final_df.drop(columns=cols_to_drop, errors="ignore")
 
-        # Separate absolute vs relative coordinates
-        review_file_df_orig = review_file_df.loc[
-            (review_file_df[xmin] > 1) & (review_file_df[xmax] > 1) & 
-            (review_file_df[ymin] > 1) & (review_file_df[ymax] > 1), :].copy()
+    return final_df
 
-        review_file_df = review_file_df.loc[
-            (review_file_df[xmin] <= 1) & (review_file_df[xmax] <= 1) & 
-            (review_file_df[ymin] <= 1) & (review_file_df[ymax] <= 1), :].copy()
+def multiply_coordinates_by_page_sizes(
+    review_file_df: pd.DataFrame,
+    page_sizes_df: pd.DataFrame,
+    xmin="xmin", xmax="xmax", ymin="ymin", ymax="ymax"
+):
+    """
+    Optimized function to convert relative coordinates to absolute based on page sizes.
 
-        if review_file_df.empty:
-            return review_file_df_orig  # If nothing is left, return the original absolute-coordinates DataFrame
+    Separates relative (<=1) and absolute (>1) coordinates, merges page sizes
+    for relative coordinates, calculates absolute pixel values, and recombines.
+    """
+    if review_file_df.empty or xmin not in review_file_df.columns:
+        return review_file_df # Return early if empty or key column missing
+
+    coord_cols = [xmin, xmax, ymin, ymax]
+    # Initial type conversion for coordinates and page
+    for col in coord_cols + ["page"]:
+        if col in review_file_df.columns:
+             # Use astype for potentially faster conversion if confident,
+             # but to_numeric is safer for mixed types/errors
+            review_file_df[col] = pd.to_numeric(review_file_df[col], errors="coerce")
+
+    # --- Identify relative coordinates ---
+    # Create mask for rows where *all* coordinates are potentially relative (<= 1)
+    # Handle potential NaNs introduced by to_numeric - treat NaN as not relative here.
+    is_relative_mask = (
+        (review_file_df[xmin].le(1) & review_file_df[xmin].notna()) &
+        (review_file_df[xmax].le(1) & review_file_df[xmax].notna()) &
+        (review_file_df[ymin].le(1) & review_file_df[ymin].notna()) &
+        (review_file_df[ymax].le(1) & review_file_df[ymax].notna())
+    )
 
-        review_file_df.loc[:, "page"] = pd.to_numeric(review_file_df["page"], errors="coerce")
+    # Separate DataFrames (minimal copies)
+    df_abs = review_file_df[~is_relative_mask].copy() # Keep absolute rows separately
+    df_rel = review_file_df[is_relative_mask].copy()  # Work only with relative rows
+
+    if df_rel.empty:
+        # If no relative coordinates, just sort and return absolute ones (if any)
+        if not df_abs.empty and {"page", xmin, ymin}.issubset(df_abs.columns):
+             df_abs.sort_values(["page", xmin, ymin], inplace=True, na_position='last')
+        return df_abs
+
+    # --- Process relative coordinates ---
+    if "image_width" not in df_rel.columns and not page_sizes_df.empty:
+        # Prepare page_sizes_df for merge
+        page_sizes_df = page_sizes_df.copy() # Avoid modifying original page_sizes_df
+        page_sizes_df['page'] = pd.to_numeric(page_sizes_df['page'], errors='coerce')
+        # Ensure proper NA handling for image dimensions
+        page_sizes_df[['image_width', 'image_height']] = page_sizes_df[['image_width','image_height']].replace("<NA>", pd.NA)
+        page_sizes_df['image_width'] = pd.to_numeric(page_sizes_df['image_width'], errors='coerce')
+        page_sizes_df['image_height'] = pd.to_numeric(page_sizes_df['image_height'], errors='coerce')
+
+        # Merge page sizes
+        df_rel = df_rel.merge(
+            page_sizes_df[['page', 'image_width', 'image_height']],
+            on="page",
+            how="left"
+        )
 
-        if "image_width" not in review_file_df.columns and not page_sizes_df.empty:
-            page_sizes_df[['image_width', 'image_height']] = page_sizes_df[['image_width','image_height']].replace("<NA>", pd.NA)  # Ensure proper NA handling
-            review_file_df = review_file_df.merge(page_sizes_df, on="page", how="left")
+    # Multiply coordinates where image dimensions are available
+    if "image_width" in df_rel.columns:
+        # Create mask for rows in df_rel that have valid image dimensions
+        has_size_mask = df_rel["image_width"].notna() & df_rel["image_height"].notna()
 
-        if "image_width" in review_file_df.columns:
-            # Split into rows with/without image size info
-            review_file_df_not_na = review_file_df.loc[review_file_df["image_width"].notna()].copy()
-            review_file_df_na = review_file_df.loc[review_file_df["image_width"].isna()].copy()
+        # Apply multiplication using .loc and the mask (vectorized and efficient)
+        # Ensure columns are numeric before multiplication (might be redundant if types are good)
+        # df_rel.loc[has_size_mask, coord_cols + ['image_width', 'image_height']] = df_rel.loc[has_size_mask, coord_cols + ['image_width', 'image_height']].apply(pd.to_numeric, errors='coerce')
 
-            if not review_file_df_not_na.empty:
-                convert_type_cols = ["image_width", "image_height", xmin, xmax, ymin, ymax]
-                review_file_df_not_na[convert_type_cols] = review_file_df_not_na[convert_type_cols].apply(pd.to_numeric, errors="coerce")
+        df_rel.loc[has_size_mask, xmin] *= df_rel.loc[has_size_mask, "image_width"]
+        df_rel.loc[has_size_mask, xmax] *= df_rel.loc[has_size_mask, "image_width"]
+        df_rel.loc[has_size_mask, ymin] *= df_rel.loc[has_size_mask, "image_height"]
+        df_rel.loc[has_size_mask, ymax] *= df_rel.loc[has_size_mask, "image_height"]
 
-                # Multiply coordinates by image sizes
-                review_file_df_not_na[xmin] *= review_file_df_not_na["image_width"]
-                review_file_df_not_na[xmax] *= review_file_df_not_na["image_width"]
-                review_file_df_not_na[ymin] *= review_file_df_not_na["image_height"]
-                review_file_df_not_na[ymax] *= review_file_df_not_na["image_height"]
 
-            # Concatenate the modified and unmodified data
-            review_file_df = pd.concat([df for df in [review_file_df_not_na, review_file_df_na] if not df.empty])
+    # --- Combine absolute and processed relative DataFrames ---
+    # Use list comprehension to handle potentially empty DataFrames
+    dfs_to_concat = [df for df in [df_abs, df_rel] if not df.empty]
 
-        # Merge with the original absolute-coordinates DataFrame
-        dfs_to_concat = [df for df in [review_file_df_orig, review_file_df] if not df.empty]
-        if dfs_to_concat:  # Ensure there's at least one non-empty DataFrame
-            review_file_df = pd.concat(dfs_to_concat)
-        else:
-            review_file_df = pd.DataFrame()  # Return an empty DataFrame instead of raising an error
+    if not dfs_to_concat:
+        return pd.DataFrame() # Return empty if both are empty
 
-        # Only sort if the DataFrame is not empty and contains the required columns
-        required_sort_columns = {"page", "xmin", "ymin"}
-        if not review_file_df.empty and required_sort_columns.issubset(review_file_df.columns):
-            review_file_df.sort_values(["page", "xmin", "ymin"], inplace=True)
+    final_df = pd.concat(dfs_to_concat, ignore_index=True) # ignore_index is good practice after filtering/concat
 
-    return review_file_df
+    # --- Final Sort ---
+    required_sort_columns = {"page", xmin, ymin}
+    if not final_df.empty and required_sort_columns.issubset(final_df.columns):
+        # Handle potential NaNs in sort columns gracefully
+        final_df.sort_values(["page", xmin, ymin], inplace=True, na_position='last')
 
+    return final_df
 
 def do_proximity_match_by_page_for_text(df1:pd.DataFrame, df2:pd.DataFrame):
     '''
@@ -1018,7 +1143,6 @@ def do_proximity_match_by_page_for_text(df1:pd.DataFrame, df2:pd.DataFrame):
 
     return merged_df
 
-
 def do_proximity_match_all_pages_for_text(df1:pd.DataFrame, df2:pd.DataFrame, threshold:float=0.03):
     '''
     Match text from one dataframe to another based on proximity matching of coordinates across all pages.
@@ -1142,12 +1266,12 @@ def convert_annotation_data_to_dataframe(all_annotations: List[Dict[str, Any]]):
     # prevents this from being necessary.
 
     # 7. Ensure essential columns exist and set column order
-    essential_box_cols = ["xmin", "xmax", "ymin", "ymax", "text", "id"]
+    essential_box_cols = ["xmin", "xmax", "ymin", "ymax", "text", "id", "label"]
     for col in essential_box_cols:
         if col not in final_df.columns:
             final_df[col] = pd.NA # Add column with NA if it wasn't present in any box
 
-    base_cols = ["image", "page"]
+    base_cols = ["image"]
     extra_box_cols = [col for col in final_df.columns if col not in base_cols and col not in essential_box_cols]
     final_col_order = base_cols + essential_box_cols + sorted(extra_box_cols)
 
@@ -1185,7 +1309,8 @@ def create_annotation_dicts_from_annotation_df(
     available_cols = [col for col in box_cols if col in all_image_annotations_df.columns]
 
     if 'text' in all_image_annotations_df.columns:
-        all_image_annotations_df.loc[all_image_annotations_df['text'].isnull(), 'text'] = ''
+        all_image_annotations_df['text'] = all_image_annotations_df['text'].fillna('')
+        #all_image_annotations_df.loc[all_image_annotations_df['text'].isnull(), 'text'] = ''
 
     if not available_cols:
         print(f"Warning: None of the expected box columns ({box_cols}) found in DataFrame.")
@@ -1226,85 +1351,84 @@ def create_annotation_dicts_from_annotation_df(
 
     return result
 
-def convert_annotation_json_to_review_df(all_annotations: List[dict],
-                                         redaction_decision_output: pd.DataFrame = pd.DataFrame(),
-                                         page_sizes: List[dict] = [],
-                                         do_proximity_match: bool = True) -> pd.DataFrame:
+def convert_annotation_json_to_review_df(
+    all_annotations: List[dict],
+    redaction_decision_output: pd.DataFrame = pd.DataFrame(),
+    page_sizes: List[dict] = [],
+    do_proximity_match: bool = True
+) -> pd.DataFrame:
     '''
     Convert the annotation json data to a dataframe format.
     Add on any text from the initial review_file dataframe by joining based on 'id' if available
     in both sources, otherwise falling back to joining on pages/co-ordinates (if option selected).
+
+    Refactored for improved efficiency, prioritizing ID-based join and conditionally applying
+    coordinate division and proximity matching.
     '''
 
     # 1. Convert annotations to DataFrame
-    # Ensure convert_annotation_data_to_dataframe populates the 'id' column
-    # if 'id' exists in the dictionaries within all_annotations.
-
     review_file_df = convert_annotation_data_to_dataframe(all_annotations)
 
-    # Only keep rows in review_df where there are coordinates
-    review_file_df.dropna(subset='xmin', axis=0, inplace=True)
+    # Only keep rows in review_df where there are coordinates (assuming xmin is representative)
+    # Use .notna() for robustness with potential None or NaN values
+    review_file_df.dropna(subset=['xmin', 'ymin', 'xmax', 'ymax'], how='any', inplace=True)
 
     # Exit early if the initial conversion results in an empty DataFrame
     if review_file_df.empty:
         # Define standard columns for an empty return DataFrame
-        check_columns = ["image", "page", "label", "color", "xmin", "ymin", "xmax", "ymax", "text", "id"]
-        # Ensure 'id' is included if it might have been expected
-        return pd.DataFrame(columns=[col for col in check_columns if col != 'id' or 'id' in review_file_df.columns])
-
-    # 2. Handle page sizes if provided
-    if not page_sizes:
-        page_sizes_df = pd.DataFrame(page_sizes) # Ensure it's a DataFrame
-        # Safely convert page column to numeric
-        page_sizes_df["page"] = pd.to_numeric(page_sizes_df["page"], errors="coerce")
-        page_sizes_df.dropna(subset=["page"], inplace=True) # Drop rows where conversion failed
-        page_sizes_df["page"] = page_sizes_df["page"].astype(int) # Convert to int after handling errors/NaNs
+        # Ensure 'id' is included if it was potentially expected based on input structure
+        # We don't know the columns from convert_annotation_data_to_dataframe without seeing it,
+        # but let's assume a standard set and add 'id' if it appeared.
+        standard_cols = ["image", "page", "label", "color", "xmin", "ymin", "xmax", "ymax", "text"]
+        if 'id' in review_file_df.columns:
+             standard_cols.append('id')
+        return pd.DataFrame(columns=standard_cols)
+
+    # Ensure 'id' column exists for logic flow, even if empty
+    if 'id' not in review_file_df.columns:
+        review_file_df['id'] = ''
+    # Do the same for redaction_decision_output if it's not empty
+    if not redaction_decision_output.empty and 'id' not in redaction_decision_output.columns:
+         redaction_decision_output['id'] = ''
 
 
-        # Apply coordinate division if page_sizes_df is not empty after processing
+    # 2. Process page sizes if provided - needed potentially for coordinate division later
+    # Process this once upfront if the data is available
+    page_sizes_df = pd.DataFrame() # Initialize as empty
+    if page_sizes:
+        page_sizes_df = pd.DataFrame(page_sizes)
         if not page_sizes_df.empty:
-            # Ensure 'page' column in review_file_df is numeric for merging
-            if 'page' in review_file_df.columns:
-                 review_file_df['page'] = pd.to_numeric(review_file_df['page'], errors='coerce')
-                 # Drop rows with invalid pages before division
-                 review_file_df.dropna(subset=['page'], inplace=True)
-                 review_file_df['page'] = review_file_df['page'].astype(int)
-                 review_file_df = divide_coordinates_by_page_sizes(review_file_df, page_sizes_df)
-
-                 print("review_file_df after coord divide:", review_file_df)
-
-            # Also apply to redaction_decision_output if it's not empty and has page numbers
-            if not redaction_decision_output.empty and 'page' in redaction_decision_output.columns:
-                redaction_decision_output['page'] = pd.to_numeric(redaction_decision_output['page'], errors='coerce')
-                # Drop rows with invalid pages before division
-                redaction_decision_output.dropna(subset=['page'], inplace=True)
-                redaction_decision_output['page'] = redaction_decision_output['page'].astype(int)
-                redaction_decision_output = divide_coordinates_by_page_sizes(redaction_decision_output, page_sizes_df)
-
-                print("redaction_decision_output after coord divide:", redaction_decision_output)
-        else:
-             print("Warning: Page sizes DataFrame became empty after processing, skipping coordinate division.")
+            # Safely convert page column to numeric and then int
+            page_sizes_df["page"] = pd.to_numeric(page_sizes_df["page"], errors="coerce")
+            page_sizes_df.dropna(subset=["page"], inplace=True)
+            if not page_sizes_df.empty: # Check again after dropping NaNs
+                page_sizes_df["page"] = page_sizes_df["page"].astype(int)
+            else:
+                 print("Warning: Page sizes DataFrame became empty after processing, coordinate division will be skipped.")
 
 
     # 3. Join additional data from redaction_decision_output if provided
+    text_added_successfully = False # Flag to track if text was added by any method
+
     if not redaction_decision_output.empty:
-        # --- NEW LOGIC: Prioritize joining by 'id' ---
-        id_col_exists_in_review = 'id' in review_file_df.columns
-        id_col_exists_in_redaction = 'id' in redaction_decision_output.columns
-        joined_by_id = False # Flag to track if ID join was successful
+        # --- Attempt to join data based on 'id' column first ---
+
+        # Check if 'id' columns are present and have non-null values in *both* dataframes
+        id_col_exists_in_review = 'id' in review_file_df.columns and not review_file_df['id'].isnull().all() and not (review_file_df['id'] == '').all()
+        id_col_exists_in_redaction = 'id' in redaction_decision_output.columns and not redaction_decision_output['id'].isnull().all() and not (redaction_decision_output['id'] == '').all()
+
 
         if id_col_exists_in_review and id_col_exists_in_redaction:
             #print("Attempting to join data based on 'id' column.")
             try:
-                # Ensure 'id' columns are of compatible types (e.g., string) to avoid merge errors
+                # Ensure 'id' columns are of string type for robust merging
                 review_file_df['id'] = review_file_df['id'].astype(str)
-                # Make a copy to avoid SettingWithCopyWarning if redaction_decision_output is used elsewhere
+                # Make a copy if needed, but try to avoid if redaction_decision_output isn't modified later
+                # Let's use a copy for safety as in the original code
                 redaction_copy = redaction_decision_output.copy()
                 redaction_copy['id'] = redaction_copy['id'].astype(str)
 
-                # Select columns to merge from redaction output.
-                # Primarily interested in 'text', but keep 'id' for the merge key.
-                # Add other columns from redaction_copy if needed.
+                # Select columns to merge from redaction output. Prioritize 'text'.
                 cols_to_merge = ['id']
                 if 'text' in redaction_copy.columns:
                     cols_to_merge.append('text')
@@ -1312,83 +1436,128 @@ def convert_annotation_json_to_review_df(all_annotations: List[dict],
                     print("Warning: 'text' column not found in redaction_decision_output. Cannot merge text using 'id'.")
 
                 # Perform a left merge to keep all annotations and add matching text
-                # Suffixes prevent collision if 'text' already exists and we want to compare/choose
-                original_cols = review_file_df.columns.tolist()
+                # Use a suffix for the text column from the right DataFrame
+                original_text_col_exists = 'text' in review_file_df.columns
+                merge_suffix = '_redaction' if original_text_col_exists else ''
+
                 merged_df = pd.merge(
                     review_file_df,
                     redaction_copy[cols_to_merge],
                     on='id',
                     how='left',
-                    suffixes=('', '_redaction') # Suffix applied to columns from right df if names clash
+                    suffixes=('', merge_suffix)
                 )
 
-                # Update the original 'text' column. Prioritize text from redaction output.
-                # If redaction output had 'text', a 'text_redaction' column now exists.
-                if 'text_redaction' in merged_df.columns:
-                     if 'text' not in merged_df.columns: # If review_file_df didn't have text initially
-                         merged_df['text'] = merged_df['text_redaction']
-                     else:
-                         # Use text from redaction where available, otherwise keep original text
-                         merged_df['text'] = merged_df['text_redaction'].combine_first(merged_df['text'])
-
-                     # Remove the temporary column
-                     merged_df = merged_df.drop(columns=['text_redaction'])
+                # Update the 'text' column if a new one was brought in
+                if 'text' + merge_suffix in merged_df.columns:
+                    redaction_text_col = 'text' + merge_suffix
+                    if original_text_col_exists:
+                         # Combine: Use text from redaction where available, otherwise keep original
+                         merged_df['text'] = merged_df[redaction_text_col].combine_first(merged_df['text'])
+                         # Drop the temporary column
+                         merged_df = merged_df.drop(columns=[redaction_text_col])
+                    else:
+                         # Redaction output had text, but review_file_df didn't. Rename the new column.
+                         merged_df = merged_df.rename(columns={redaction_text_col: 'text'})
 
-                # Ensure final columns match original expectation + potentially new 'text'
-                final_cols = original_cols
-                if 'text' not in final_cols and 'text' in merged_df.columns:
-                    final_cols.append('text') # Make sure text column is kept if newly added
-                 # Reorder/select columns if necessary, ensuring 'id' is kept
-                review_file_df = merged_df[[col for col in final_cols if col in merged_df.columns] + (['id'] if 'id' not in final_cols else [])]
+                    text_added_successfully = True # Indicate text was potentially added
 
+                review_file_df = merged_df # Update the main DataFrame
 
-                #print("Successfully joined data using 'id'.")
-                joined_by_id = True
+                #print("Successfully attempted to join data using 'id'.") # Note: Text might not have been in redaction data
 
             except Exception as e:
-                print(f"Error during 'id'-based merge: {e}. Falling back to proximity match if enabled.")
-                # Fall through to proximity match below if an error occurred
-
-        # --- Fallback to proximity match ---
-        if not joined_by_id and do_proximity_match:
-            if not id_col_exists_in_review or not id_col_exists_in_redaction:
-                 print("Could not join by 'id' (column missing in one or both sources).")
-            print("Performing proximity match to add text data.")
-            # Match text to review file using proximity
-
-            review_file_df = do_proximity_match_all_pages_for_text(df1=review_file_df.copy(), df2=redaction_decision_output.copy())
-        elif not joined_by_id and not do_proximity_match:
-             print("Skipping joining text data (ID join not possible, proximity match disabled).")
-        # --- End of join logic ---
-
-    # 4. Ensure required columns exist, filling with blank if they don't
-    # Define base required columns, 'id' might or might not be present initially
-    required_columns = ["image", "page", "label", "color", "xmin", "ymin", "xmax", "ymax", "text"]
-    # Add 'id' to required list if it exists in the dataframe at this point
+                print(f"Error during 'id'-based merge: {e}. Checking for proximity match fallback.")
+                # Fall through to proximity match logic below
+
+        # --- Fallback to proximity match if ID join wasn't possible/successful and enabled ---
+        # Note: If id_col_exists_in_review or id_col_exists_in_redaction was False,
+        # the block above was skipped, and we naturally fall here.
+        # If an error occurred in the try block, joined_by_id would implicitly be False
+        # because text_added_successfully wasn't set to True.
+
+        # Only attempt proximity match if text wasn't added by ID join and proximity is requested
+        if not text_added_successfully and do_proximity_match:
+             print("Attempting proximity match to add text data.")
+
+             # Ensure 'page' columns are numeric before coordinate division and proximity match
+             # (Assuming divide_coordinates_by_page_sizes and do_proximity_match_all_pages_for_text need this)
+             if 'page' in review_file_df.columns:
+                 review_file_df['page'] = pd.to_numeric(review_file_df['page'], errors='coerce').fillna(-1).astype(int) # Use -1 for NaN pages
+                 review_file_df = review_file_df[review_file_df['page'] != -1] # Drop rows where page conversion failed
+             if not redaction_decision_output.empty and 'page' in redaction_decision_output.columns:
+                  redaction_decision_output['page'] = pd.to_numeric(redaction_decision_output['page'], errors='coerce').fillna(-1).astype(int)
+                  redaction_decision_output = redaction_decision_output[redaction_decision_output['page'] != -1]
+
+             # Perform coordinate division IF page_sizes were processed and DataFrame is not empty
+             if not page_sizes_df.empty:
+                  # Apply coordinate division *before* proximity match
+                  review_file_df = divide_coordinates_by_page_sizes(review_file_df, page_sizes_df)
+                  if not redaction_decision_output.empty:
+                       redaction_decision_output = divide_coordinates_by_page_sizes(redaction_decision_output, page_sizes_df)
+
+             # Now perform the proximity match
+             # Note: Potential DataFrame copies happen inside do_proximity_match based on its implementation
+             if not redaction_decision_output.empty:
+                try:
+                    review_file_df = do_proximity_match_all_pages_for_text(
+                        df1=review_file_df, # Pass directly, avoid caller copy if possible by modifying function signature
+                        df2=redaction_decision_output # Pass directly
+                    )
+                    # Assuming do_proximity_match_all_pages_for_text adds/updates the 'text' column
+                    if 'text' in review_file_df.columns:
+                         text_added_successfully = True
+                    print("Proximity match completed.")
+                except Exception as e:
+                    print(f"Error during proximity match: {e}. Text data may not be added.")
+
+        elif not text_added_successfully and not do_proximity_match:
+             print("Skipping joining text data (ID join not possible/failed, proximity match disabled).")
+
+    # 4. Ensure required columns exist and are ordered
+    # Define base required columns. 'id' and 'text' are conditionally added.
+    required_columns_base = ["image", "page", "label", "color", "xmin", "ymin", "xmax", "ymax"]
+    final_columns = required_columns_base[:] # Start with base columns
+
+    # Add 'id' and 'text' if they exist in the DataFrame at this point
     if 'id' in review_file_df.columns:
-        required_columns.append('id')
+        final_columns.append('id')
+    if 'text' in review_file_df.columns:
+        final_columns.append('text') # Add text column if it was created/merged
 
-    for col in required_columns:
+    # Add any missing required columns with a default value (e.g., blank string)
+    for col in final_columns:
         if col not in review_file_df.columns:
-            # Decide default value based on column type (e.g., '' for text, np.nan for numeric?)
-            # Using '' for simplicity here.
-            review_file_df[col] = ''
+            # Use appropriate default based on expected type, '' for text/id, np.nan for coords?
+            # Sticking to '' as in original for simplicity, but consider data types.
+            review_file_df[col] = '' # Or np.nan for numerical, but coords already checked by dropna
 
     # Select and order the final set of columns
-    review_file_df = review_file_df[required_columns]
+    # Ensure all selected columns actually exist after adding defaults
+    review_file_df = review_file_df[[col for col in final_columns if col in review_file_df.columns]]
+
 
     # 5. Final processing and sorting
-    # If colours are saved as list, convert to tuple
+    # Convert colours from list to tuple if necessary - apply is okay here unless lists are vast
     if 'color' in review_file_df.columns:
-        review_file_df["color"] = review_file_df["color"].apply(lambda x: tuple(x) if isinstance(x, list) else x)
+         # Check if the column actually contains lists before applying lambda
+         if review_file_df['color'].apply(lambda x: isinstance(x, list)).any():
+            review_file_df["color"] = review_file_df["color"].apply(lambda x: tuple(x) if isinstance(x, list) else x)
 
     # Sort the results
-    sort_columns = ['page', 'ymin', 'xmin', 'label']
     # Ensure sort columns exist before sorting
+    sort_columns = ['page', 'ymin', 'xmin', 'label']
     valid_sort_columns = [col for col in sort_columns if col in review_file_df.columns]
-    if valid_sort_columns:
-        review_file_df = review_file_df.sort_values(valid_sort_columns)
-
+    if valid_sort_columns and not review_file_df.empty: # Only sort non-empty df
+         # Convert potential numeric sort columns to appropriate types if necessary
+         # (e.g., 'page', 'ymin', 'xmin') to ensure correct sorting.
+         # dropna(subset=[...], inplace=True) earlier should handle NaNs in coords.
+         # page conversion already done before proximity match.
+         try:
+             review_file_df = review_file_df.sort_values(valid_sort_columns)
+         except TypeError as e:
+              print(f"Warning: Could not sort DataFrame due to type error in sort columns: {e}")
+              # Proceed without sorting
     return review_file_df
 
 def fill_missing_box_ids(data_input: dict) -> dict:
@@ -1472,20 +1641,18 @@ def fill_missing_box_ids(data_input: dict) -> dict:
 
 def fill_missing_ids(df: pd.DataFrame, column_name: str = 'id', length: int = 12) -> pd.DataFrame:
     """
-    Generates unique alphanumeric IDs for rows in a DataFrame column
-    where the value is missing (NaN, None) or an empty string.
+    Optimized: Generates unique alphanumeric IDs for rows in a DataFrame column
+    where the value is missing (NaN, None) or an empty/whitespace string.
 
     Args:
         df (pd.DataFrame): The input Pandas DataFrame.
         column_name (str): The name of the column to check and fill (defaults to 'id').
                            This column will be added if it doesn't exist.
         length (int): The desired length of the generated IDs (defaults to 12).
-                      Cannot exceed the limits that guarantee uniqueness based
-                      on the number of IDs needed and character set size.
 
     Returns:
         pd.DataFrame: The DataFrame with missing/empty IDs filled in the specified column.
-                      Note: The function modifies the DataFrame in place.
+                      Note: The function modifies the DataFrame directly (in-place).
     """
 
     # --- Input Validation ---
@@ -1497,43 +1664,59 @@ def fill_missing_ids(df: pd.DataFrame, column_name: str = 'id', length: int = 12
         raise ValueError("'length' must be a positive integer.")
 
     # --- Ensure Column Exists ---
+    original_dtype = None
     if column_name not in df.columns:
         print(f"Column '{column_name}' not found. Adding it to the DataFrame.")
-        df[column_name] = np.nan # Initialize with NaN
+        # Initialize with None (which Pandas often treats as NaN but allows object dtype)
+        df[column_name] = None
+        # Set original_dtype to object so it likely becomes string later
+        original_dtype = object
+    else:
+        original_dtype = df[column_name].dtype
 
     # --- Identify Rows Needing IDs ---
-    # Check for NaN, None, or empty strings ('')
-    # Convert to string temporarily for robust empty string check, handle potential errors
-    try:
-        df[column_name] = df[column_name].astype(str) #handles NaN/None conversion, .str.strip() removes whitespace
-        is_missing_or_empty = (
-            df[column_name].isna()
-            #| (df[column_name].astype(str).str.strip() == '')
-            #| (df[column_name] == "nan")
-            | (df[column_name].astype(str).str.len() != length)
-        )
-    except Exception as e:
-         # Fallback if conversion to string fails (e.g., column contains complex objects)
-         print(f"Warning: Could not perform reliable empty string check on column '{column_name}' due to data type issues. Checking for NaN/None only. Error: {e}")
-         is_missing_or_empty = df[column_name].isna()
+    # 1. Check for actual null values (NaN, None, NaT)
+    is_null = df[column_name].isna()
+
+    # 2. Check for empty or whitespace-only strings AFTER converting potential values to string
+    #    Only apply string checks on rows that are *not* null to avoid errors/warnings
+    #    Fill NaN temporarily for string operations, then check length or equality
+    is_empty_str = pd.Series(False, index=df.index) # Default to False
+    if not is_null.all(): # Only check strings if there are non-null values
+         temp_str_col = df.loc[~is_null, column_name].astype(str).str.strip()
+         is_empty_str.loc[~is_null] = (temp_str_col == '')
+
+    # Combine the conditions
+    is_missing_or_empty = is_null | is_empty_str
 
     rows_to_fill_index = df.index[is_missing_or_empty]
     num_needed = len(rows_to_fill_index)
 
     if num_needed == 0:
-        #print(f"No missing or empty values found in column '{column_name}'.")
+        # Ensure final column type is consistent if nothing was done
+        if pd.api.types.is_object_dtype(original_dtype) or pd.api.types.is_string_dtype(original_dtype):
+             pass # Likely already object or string
+        else:
+             # If original was numeric/etc., but might contain strings now? Unlikely here.
+             pass # Or convert to object: df[column_name] = df[column_name].astype(object)
+        # print(f"No missing or empty values found requiring IDs in column '{column_name}'.")
         return df
 
     print(f"Found {num_needed} rows requiring a unique ID in column '{column_name}'.")
 
     # --- Get Existing IDs to Ensure Uniqueness ---
-    try:
-        # Get all non-missing, non-empty string values from the column
-        existing_ids = set(df.loc[~is_missing_or_empty, column_name].astype(str))
-    except Exception as e:
-        print(f"Warning: Could not reliably get all existing string IDs from column '{column_name}' due to data type issues. Uniqueness check might be less strict. Error: {e}")
-        # Fallback: Get only non-NaN IDs, potential type issues ignored
-        existing_ids = set(df.loc[df[column_name].notna(), column_name])
+    # Consider only rows that are *not* missing/empty
+    valid_rows = df.loc[~is_missing_or_empty, column_name]
+    # Drop any remaining nulls (shouldn't be any based on mask, but belts and braces)
+    valid_rows = valid_rows.dropna()
+    # Convert to string *only* if not already string/object, then filter out empty strings again
+    if not pd.api.types.is_object_dtype(valid_rows.dtype) and not pd.api.types.is_string_dtype(valid_rows.dtype):
+         existing_ids = set(valid_rows.astype(str).str.strip())
+    else: # Already string or object, just strip and convert to set
+         existing_ids = set(valid_rows.astype(str).str.strip()) # astype(str) handles mixed types in object column
+
+    # Remove empty string from existing IDs if it's there after stripping
+    existing_ids.discard('')
 
 
     # --- Generate Unique IDs ---
@@ -1543,93 +1726,230 @@ def fill_missing_ids(df: pd.DataFrame, column_name: str = 'id', length: int = 12
 
     max_possible_ids = len(character_set) ** length
     if num_needed > max_possible_ids:
-         raise ValueError(f"Cannot generate {num_needed} unique IDs with length {length}. Maximum possible is {max_possible_ids}.")
-    # Add a check for practical limits if needed, e.g., if num_needed is very close to max_possible_ids, generation could be slow.
+        raise ValueError(f"Cannot generate {num_needed} unique IDs with length {length}. Maximum possible is {max_possible_ids}.")
+
+    # Pre-calculate safety break limit
+    max_attempts_per_id = max(1000, num_needed * 10) # Adjust multiplier as needed
 
     #print(f"Generating {num_needed} unique IDs of length {length}...")
     for i in range(num_needed):
         attempts = 0
         while True:
             candidate_id = ''.join(random.choices(character_set, k=length))
-            # Check against *all* existing IDs and *newly* generated ones
+            # Check against *all* known existing IDs and *newly* generated ones
             if candidate_id not in existing_ids and candidate_id not in generated_ids_set:
                 generated_ids_set.add(candidate_id)
                 new_ids_list.append(candidate_id)
                 break # Found a unique ID
             attempts += 1
-            if attempts > num_needed * 100 and attempts > 1000 : # Safety break for unlikely infinite loop
-                 raise RuntimeError(f"Failed to generate a unique ID after {attempts} attempts. Check length and character set or existing IDs.")
+            if attempts > max_attempts_per_id : # Safety break
+                raise RuntimeError(f"Failed to generate a unique ID after {attempts} attempts. Check length, character set, or density of existing IDs.")
 
-        # Optional progress update for large numbers
-        if (i + 1) % 1000 == 0:
-            print(f"Generated {i+1}/{num_needed} IDs...")
+        # Optional progress update
+        # if (i + 1) % 1000 == 0:
+        #    print(f"Generated {i+1}/{num_needed} IDs...")
 
 
     # --- Assign New IDs ---
     # Use the previously identified index to assign the new IDs correctly
+    # Assigning string IDs might change the column's dtype to 'object'
+    if not pd.api.types.is_object_dtype(original_dtype) and not pd.api.types.is_string_dtype(original_dtype):
+         warnings.warn(f"Column '{column_name}' dtype might change from '{original_dtype}' to 'object' due to string ID assignment.", UserWarning)
+
     df.loc[rows_to_fill_index, column_name] = new_ids_list
-    #print(f"Successfully filled {len(new_ids_list)} missing values in column '{column_name}'.")
+    print(f"Successfully assigned {len(new_ids_list)} new unique IDs to column '{column_name}'.")
+
+    # Optional: Convert the entire column to string type at the end for consistency
+    # df[column_name] = df[column_name].astype(str)
 
-    # The DataFrame 'df' has been modified in place
     return df
 
-def convert_review_df_to_annotation_json(review_file_df:pd.DataFrame,
-                                         image_paths:List[Image.Image],
-                                         page_sizes:List[dict]=[]) -> List[dict]:
-    '''
-    Convert a review csv to a json file for use by the Gradio Annotation object.
-    '''
-    # Make sure all relevant cols are float
-    float_cols = ["page", "xmin", "xmax", "ymin", "ymax"]
-    for col in float_cols:
-        review_file_df.loc[:, col] = pd.to_numeric(review_file_df.loc[:, col], errors='coerce')
-    
-    # Convert relative co-ordinates into image coordinates for the image annotation output object
-    if page_sizes:        
-        page_sizes_df = pd.DataFrame(page_sizes)
-        page_sizes_df[["page"]] = page_sizes_df[["page"]].apply(pd.to_numeric, errors="coerce")
+def convert_review_df_to_annotation_json(
+    review_file_df: pd.DataFrame,
+    image_paths: List[str], # List of image file paths
+    page_sizes: List[Dict], # List of dicts like [{'page': 1, 'image_path': '...', 'image_width': W, 'image_height': H}, ...]
+    xmin="xmin", xmax="xmax", ymin="ymin", ymax="ymax" # Coordinate column names
+) -> List[Dict]:
+    """
+    Optimized function to convert review DataFrame to Gradio Annotation JSON format.
 
-        review_file_df = multiply_coordinates_by_page_sizes(review_file_df, page_sizes_df)
-    
-    review_file_df = fill_missing_ids(review_file_df)
+    Ensures absolute coordinates, handles missing IDs, deduplicates based on key fields,
+    selects final columns, and structures data per image/page based on page_sizes.
 
-    if 'id' not in review_file_df.columns:
-        review_file_df['id'] = ''
-        review_file_df['id'] = review_file_df['id'].astype(str)
-           
-    # Keep only necessary columns
-    review_file_df = review_file_df[["image", "page", "label", "color", "xmin", "ymin", "xmax", "ymax", "id", "text"]].drop_duplicates(subset=["image", "page", "xmin", "ymin", "xmax", "ymax", "label", "id"])
-
-    # If colours are saved as list, convert to tuple
-    review_file_df.loc[:, "color"] = review_file_df.loc[:,"color"].apply(lambda x: tuple(x) if isinstance(x, list) else x)
+    Args:
+        review_file_df: Input DataFrame with annotation data.
+        image_paths: List of image file paths (Note: currently unused if page_sizes provides paths).
+        page_sizes: REQUIRED list of dictionaries, each containing 'page',
+                    'image_path', 'image_width', and 'image_height'. Defines
+                    output structure and dimensions for coordinate conversion.
+        xmin, xmax, ymin, ymax: Names of the coordinate columns.
 
-    # Group the DataFrame by the 'image' column
-    grouped_csv_pages = review_file_df.groupby('page')
+    Returns:
+        List of dictionaries suitable for Gradio Annotation output, one dict per image/page.
+    """
+    if not page_sizes:
+        raise ValueError("page_sizes argument is required and cannot be empty.")
 
-    # Create a list to hold the JSON data
-    json_data = []
+    # --- Prepare Page Sizes DataFrame ---
+    try:
+        page_sizes_df = pd.DataFrame(page_sizes)
+        required_ps_cols = {'page', 'image_path', 'image_width', 'image_height'}
+        if not required_ps_cols.issubset(page_sizes_df.columns):
+            missing = required_ps_cols - set(page_sizes_df.columns)
+            raise ValueError(f"page_sizes is missing required keys: {missing}")
+        # Convert page sizes columns to appropriate numeric types early
+        page_sizes_df['page'] = pd.to_numeric(page_sizes_df['page'], errors='coerce')
+        page_sizes_df['image_width'] = pd.to_numeric(page_sizes_df['image_width'], errors='coerce')
+        page_sizes_df['image_height'] = pd.to_numeric(page_sizes_df['image_height'], errors='coerce')
+        # Use nullable Int64 for page number consistency
+        page_sizes_df['page'] = page_sizes_df['page'].astype('Int64')
 
-    for page_no, pdf_image_path in enumerate(page_sizes_df["image_path"]):
-        
-        reported_page_number = int(page_no + 1)
+    except Exception as e:
+        raise ValueError(f"Error processing page_sizes: {e}") from e
 
-        if reported_page_number in review_file_df["page"].values:
 
-            # Convert each relevant group to a list of box dictionaries
-            selected_csv_pages = grouped_csv_pages.get_group(reported_page_number)
-            annotation_boxes = selected_csv_pages.drop(columns=['image', 'page']).to_dict(orient='records')
-            
-            annotation = {
-                "image": pdf_image_path,
-                "boxes": annotation_boxes
-            }
+    # Handle empty input DataFrame gracefully
+    if review_file_df.empty:
+        print("Input review_file_df is empty. Proceeding to generate JSON structure with empty boxes.")
+        # Ensure essential columns exist even if empty for later steps
+        for col in [xmin, xmax, ymin, ymax, "page", "label", "color", "id", "text"]:
+             if col not in review_file_df.columns:
+                 review_file_df[col] = pd.NA
+    else:
+        # --- Coordinate Conversion (if needed) ---
+        coord_cols_to_check = [c for c in [xmin, xmax, ymin, ymax] if c in review_file_df.columns]
+        needs_multiplication = False
+        if coord_cols_to_check:
+            temp_df_numeric = review_file_df[coord_cols_to_check].apply(pd.to_numeric, errors='coerce')
+            if temp_df_numeric.le(1).any().any(): # Check if any numeric coord <= 1 exists
+                 needs_multiplication = True
+
+        if needs_multiplication:
+            #print("Relative coordinates detected or suspected, running multiplication...")
+            review_file_df = multiply_coordinates_by_page_sizes(
+                review_file_df.copy(), # Pass a copy to avoid modifying original outside function
+                page_sizes_df,
+                xmin, xmax, ymin, ymax
+            )
+        else:
+            #print("No relative coordinates detected or required columns missing, skipping multiplication.")
+            # Still ensure essential coordinate/page columns are numeric if they exist
+            cols_to_convert = [c for c in [xmin, xmax, ymin, ymax, "page"] if c in review_file_df.columns]
+            for col in cols_to_convert:
+                review_file_df[col] = pd.to_numeric(review_file_df[col], errors='coerce')
 
+        # Handle potential case where multiplication returns an empty DF
+        if review_file_df.empty:
+            print("DataFrame became empty after coordinate processing.")
+            # Re-add essential columns if they were lost
+            for col in [xmin, xmax, ymin, ymax, "page", "label", "color", "id", "text"]:
+                if col not in review_file_df.columns:
+                    review_file_df[col] = pd.NA
+
+        # --- Fill Missing IDs ---
+        review_file_df = fill_missing_ids(review_file_df.copy()) # Pass a copy
+
+        # --- Deduplicate Based on Key Fields ---
+        base_dedupe_cols = ["page", xmin, ymin, xmax, ymax, "label", "id"]
+        # Identify which deduplication columns actually exist in the DataFrame
+        cols_for_dedupe = [col for col in base_dedupe_cols if col in review_file_df.columns]
+        # Add 'image' column for deduplication IF it exists (matches original logic intent)
+        if "image" in review_file_df.columns:
+            cols_for_dedupe.append("image")
+
+        # Ensure placeholder columns exist if they are needed for deduplication
+        # (e.g., 'label', 'id' should be present after fill_missing_ids)
+        for col in ['label', 'id']:
+            if col in cols_for_dedupe and col not in review_file_df.columns:
+                 # This might indicate an issue in fill_missing_ids or prior steps
+                 print(f"Warning: Column '{col}' needed for dedupe but not found. Adding NA.")
+                 review_file_df[col] = "" # Add default empty string
+
+        if cols_for_dedupe: # Only attempt dedupe if we have columns to check
+            #print(f"Deduplicating based on columns: {cols_for_dedupe}")
+            # Convert relevant columns to string before dedupe to avoid type issues with mixed data (optional, depends on data)
+            # for col in cols_for_dedupe:
+            #    review_file_df[col] = review_file_df[col].astype(str)
+            review_file_df = review_file_df.drop_duplicates(subset=cols_for_dedupe)
         else:
-            annotation = {}
-            annotation["image"] = pdf_image_path
-            annotation["boxes"] = []
+            print("Skipping deduplication: No valid columns found to deduplicate by.")
+
+
+    # --- Select and Prepare Final Output Columns ---
+    required_final_cols = ["page", "label", "color", xmin, ymin, xmax, ymax, "id", "text"]
+    # Identify which of the desired final columns exist in the (now potentially deduplicated) DataFrame
+    available_final_cols = [col for col in required_final_cols if col in review_file_df.columns]
+
+    # Ensure essential output columns exist, adding defaults if missing AFTER deduplication
+    for col in required_final_cols:
+         if col not in review_file_df.columns:
+             print(f"Adding missing final column '{col}' with default value.")
+             if col in ['label', 'id', 'text']:
+                 review_file_df[col] = "" # Default empty string
+             elif col == 'color':
+                 review_file_df[col] = None # Default None or a default color tuple
+             else: # page, coordinates
+                 review_file_df[col] = pd.NA # Default NA for numeric/page
+             available_final_cols.append(col) # Add to list of available columns
+
+    # Select only the final desired columns in the correct order
+    review_file_df = review_file_df[available_final_cols]
+
+    # --- Final Formatting ---
+    if not review_file_df.empty:
+        # Convert list colors to tuples (important for some downstream uses)
+        if 'color' in review_file_df.columns:
+            review_file_df['color'] = review_file_df['color'].apply(
+                lambda x: tuple(x) if isinstance(x, list) else x
+            )
+        # Ensure page column is nullable integer type for reliable grouping
+        if 'page' in review_file_df.columns:
+             review_file_df['page'] = review_file_df['page'].astype('Int64')
+
+    # --- Group Annotations by Page ---
+    if 'page' in review_file_df.columns:
+        grouped_annotations = review_file_df.groupby('page')
+        group_keys = set(grouped_annotations.groups.keys()) # Use set for faster lookups
+    else:
+        # Cannot group if page column is missing
+        print("Error: 'page' column missing, cannot group annotations.")
+        grouped_annotations = None
+        group_keys = set()
+
 
-        # Append the structured data to the json_data list
-        json_data.append(annotation)
+    # --- Build JSON Structure ---
+    json_data = []
+    output_cols_for_boxes = [col for col in ["label", "color", xmin, ymin, xmax, ymax, "id", "text"] if col in review_file_df.columns]
 
-    return json_data
+    # Iterate through page_sizes_df to define the structure (one entry per image path)
+    for _, row in page_sizes_df.iterrows():
+        page_num = row['page'] # Already Int64
+        pdf_image_path = row['image_path']
+        annotation_boxes = [] # Default to empty list
+
+        # Check if the page exists in the grouped annotations (using the faster set lookup)
+        # Check pd.notna because page_num could be <NA> if conversion failed
+        if pd.notna(page_num) and page_num in group_keys and grouped_annotations:
+            try:
+                page_group_df = grouped_annotations.get_group(page_num)
+                # Convert the group to list of dicts, selecting only needed box properties
+                # Handle potential NaN coordinates before conversion to JSON
+                annotation_boxes = page_group_df[output_cols_for_boxes].replace({np.nan: None}).to_dict(orient='records')
+
+                # Optional: Round coordinates here if needed AFTER potential multiplication
+                # for box in annotation_boxes:
+                #     for coord in [xmin, ymin, xmax, ymax]:
+                #         if coord in box and box[coord] is not None:
+                #             box[coord] = round(float(box[coord]), 2) # Example: round to 2 decimals
+
+            except KeyError:
+                 print(f"Warning: Group key {page_num} not found despite being in group_keys (should not happen).")
+                 annotation_boxes = [] # Keep empty
+
+        # Append the structured data for this image/page
+        json_data.append({
+            "image": pdf_image_path,
+            "boxes": annotation_boxes
+        })
+
+    return json_data

tools/file_redaction.py

@@ -258,8 +258,7 @@ def choose_and_run_redactor(file_paths:List[str],
 
 
     # Call prepare_image_or_pdf only if needed
-    if prepare_images_flag is not None:# and first_loop_state==True:
-        #print("Calling preparation function. prepare_images_flag:", prepare_images_flag)
+    if prepare_images_flag is not None:
         out_message, prepared_pdf_file_paths, pdf_image_file_paths, annotate_max_pages, annotate_max_pages_bottom, pymupdf_doc, annotations_all_pages, review_file_state, document_cropboxes, page_sizes, textract_output_found, all_img_details_state, placeholder_ocr_results_df = prepare_image_or_pdf(
             file_paths_loop, text_extraction_method, 0, out_message, True, 
             annotate_max_pages, annotations_all_pages, document_cropboxes, redact_whole_page_list, 
@@ -333,7 +332,7 @@ def choose_and_run_redactor(file_paths:List[str],
     # Try to connect to AWS services directly only if RUN_AWS_FUNCTIONS environmental variable is 1, otherwise an environment variable or direct textbox input is needed.
     if pii_identification_method == aws_pii_detector:
         if aws_access_key_textbox and aws_secret_key_textbox:
-            print("Connecting to Comprehend using AWS access key and secret keys from textboxes.")
+            print("Connecting to Comprehend using AWS access key and secret keys from user input.")
             comprehend_client = boto3.client('comprehend', 
                 aws_access_key_id=aws_access_key_textbox, 
                 aws_secret_access_key=aws_secret_key_textbox, region_name=AWS_REGION)
@@ -356,7 +355,7 @@ def choose_and_run_redactor(file_paths:List[str],
     # Try to connect to AWS Textract Client if using that text extraction method
     if text_extraction_method == textract_option:   
         if aws_access_key_textbox and aws_secret_key_textbox:
-            print("Connecting to Textract using AWS access key and secret keys from textboxes.")
+            print("Connecting to Textract using AWS access key and secret keys from user input.")
             textract_client = boto3.client('textract', 
                 aws_access_key_id=aws_access_key_textbox, 
                 aws_secret_access_key=aws_secret_key_textbox, region_name=AWS_REGION)
@@ -401,7 +400,7 @@ def choose_and_run_redactor(file_paths:List[str],
             is_a_pdf = is_pdf(file_path) == True
             if is_a_pdf == False and text_extraction_method == text_ocr_option:
                 # If user has not submitted a pdf, assume it's an image
-                print("File is not a pdf, assuming that image analysis needs to be used.")
+                print("File is not a PDF, assuming that image analysis needs to be used.")
                 text_extraction_method = tesseract_ocr_option
         else:
             out_message = "No file selected"
@@ -862,17 +861,6 @@ def convert_pikepdf_annotations_to_result_annotation_box(page:Page, annot:dict,
 
     rect = Rect(pymupdf_x1, pymupdf_y1, pymupdf_x2, pymupdf_y2)
 
-    # if image or image_dimensions:
-    #     print("Dividing result by image coordinates")
-
-    #     image_x1, image_y1, image_x2, image_y2 = convert_pymupdf_to_image_coords(page, pymupdf_x1, pymupdf_y1, pymupdf_x2, pymupdf_y2, image, image_dimensions=image_dimensions)
-
-    #     img_annotation_box["xmin"] = image_x1
-    #     img_annotation_box["ymin"] = image_y1
-    #     img_annotation_box["xmax"] = image_x2
-    #     img_annotation_box["ymax"] = image_y2
-
-    # else:
     convert_df = pd.DataFrame({
                             "page": [page_no],
                             "xmin": [pymupdf_x1],
@@ -1016,9 +1004,6 @@ def redact_page_with_pymupdf(page:Page, page_annotations:dict, image:Image=None,
 
             img_annotation_box = fill_missing_box_ids(img_annotation_box)
 
-            #print("image_dimensions:", image_dimensions)
-            #print("annot:", annot)
-
         all_image_annotation_boxes.append(img_annotation_box)
 
         # Redact the annotations from the document
@@ -1285,8 +1270,6 @@ def redact_image_pdf(file_path:str,
         page_handwriting_recogniser_results = []
         page_break_return = False
         reported_page_number = str(page_no + 1)
-
-        #print("page_sizes_df for row:", page_sizes_df.loc[page_sizes_df["page"] == (page_no + 1)])
         
         # Try to find image location
         try:
@@ -1328,7 +1311,7 @@ def redact_image_pdf(file_path:str,
 
             # Step 1: Perform OCR. Either with Tesseract, or with AWS Textract
 
-            # If using Tesseract, need to check if we have page as image_path
+            # If using Tesseract
             if text_extraction_method == tesseract_ocr_option:
                 #print("image_path:", image_path)
                 #print("print(type(image_path)):", print(type(image_path)))
@@ -1449,7 +1432,6 @@ def redact_image_pdf(file_path:str,
                     # Assume image_path is an image
                     image = image_path
 
-                print("image:", image)
 
                 fill = (0, 0, 0)   # Fill colour for redactions
                 draw = ImageDraw.Draw(image)
@@ -1631,8 +1613,6 @@ def get_text_container_characters(text_container:LTTextContainer):
                     for line in text_container
                     if isinstance(line, LTTextLine) or isinstance(line, LTTextLineHorizontal)
                     for char in line]
-        
-        #print("Initial characters:", characters)
     
         return characters
     return []
@@ -1762,9 +1742,6 @@ def create_text_redaction_process_results(analyser_results, analysed_bounding_bo
         analysed_bounding_boxes_df_new = pd.concat([analysed_bounding_boxes_df_new, analysed_bounding_boxes_df_text], axis = 1)
         analysed_bounding_boxes_df_new['page'] = page_num + 1
 
-        #analysed_bounding_boxes_df_new = fill_missing_ids(analysed_bounding_boxes_df_new)
-        analysed_bounding_boxes_df_new.to_csv("output/analysed_bounding_boxes_df_new_with_ids.csv")
-
         decision_process_table = pd.concat([decision_process_table, analysed_bounding_boxes_df_new], axis = 0).drop('result', axis=1)
     
     return decision_process_table
@@ -1772,7 +1749,6 @@ def create_text_redaction_process_results(analyser_results, analysed_bounding_bo
 def create_pikepdf_annotations_for_bounding_boxes(analysed_bounding_boxes):
     pikepdf_redaction_annotations_on_page = []
     for analysed_bounding_box in analysed_bounding_boxes:
-        #print("analysed_bounding_box:", analysed_bounding_boxes)
 
         bounding_box = analysed_bounding_box["boundingBox"]
         annotation = Dictionary(
@@ -1997,7 +1973,6 @@ def redact_text_pdf(
                     pass
                     #print("Not redacting page:", page_no)
 
-                #print("page_image_annotations after page", reported_page_number, "are", page_image_annotations)
 
                 # Join extracted text outputs for all lines together
                 if not page_text_ocr_outputs.empty:

tools/redaction_review.py

@@ -6,12 +6,11 @@ import numpy as np
 from xml.etree.ElementTree import Element, SubElement, tostring, parse
 from xml.dom import minidom
 import uuid
-from typing import List
+from typing import List, Tuple
 from gradio_image_annotation import image_annotator
 from gradio_image_annotation.image_annotator import AnnotatedImageData
 from pymupdf import Document, Rect
 import pymupdf
-#from fitz 
 from PIL import ImageDraw, Image
 
 from tools.config import OUTPUT_FOLDER, CUSTOM_BOX_COLOUR, MAX_IMAGE_PIXELS, INPUT_FOLDER
@@ -55,7 +54,6 @@ def update_zoom(current_zoom_level:int, annotate_current_page:int, decrease:bool
         
     return current_zoom_level, annotate_current_page
 
-
 def update_dropdown_list_based_on_dataframe(df:pd.DataFrame, column:str) -> List["str"]:
     '''
     Gather unique elements from a string pandas Series, then append 'ALL' to the start and return the list.
@@ -166,49 +164,205 @@ def update_recogniser_dataframes(page_image_annotator_object:AnnotatedImageData,
 
     return recogniser_entities_list, recogniser_dataframe_out_gr, recogniser_dataframe_out, recogniser_entities_drop, text_entities_drop, page_entities_drop
 
-def undo_last_removal(backup_review_state, backup_image_annotations_state, backup_recogniser_entity_dataframe_base):
+def undo_last_removal(backup_review_state:pd.DataFrame, backup_image_annotations_state:list[dict], backup_recogniser_entity_dataframe_base:pd.DataFrame):
     return backup_review_state, backup_image_annotations_state, backup_recogniser_entity_dataframe_base
 
-def update_annotator_page_from_review_df(review_df: pd.DataFrame,
-                                          image_file_paths:List[str],
-                                          page_sizes:List[dict],
-                                          current_page:int,
-                                          previous_page:int,
-                                          current_image_annotations_state:List[str],
-                                          current_page_annotator:object):
+def update_annotator_page_from_review_df(
+    review_df: pd.DataFrame,
+    image_file_paths:List[str], # Note: This input doesn't seem used in the original logic flow after the first line was removed
+    page_sizes:List[dict],
+    current_image_annotations_state:List[str], # This should ideally be List[dict] based on its usage
+    current_page_annotator:object, # Should be dict or a custom annotation object for one page
+    selected_recogniser_entity_df_row:pd.DataFrame,
+    input_folder:str,
+    doc_full_file_name_textbox:str
+) -> Tuple[object, List[dict], int, List[dict], pd.DataFrame, int]: # Correcting return types based on usage
     '''
-    Update the visible annotation object with the latest review file information
+    Update the visible annotation object and related objects with the latest review file information,
+    optimizing by processing only the current page's data.
     '''
-    out_image_annotations_state = current_image_annotations_state
-    out_current_page_annotator = current_page_annotator
-    gradio_annotator_current_page_number = current_page
+    # Assume current_image_annotations_state is List[dict] and current_page_annotator is dict
+    out_image_annotations_state: List[dict] = list(current_image_annotations_state) # Make a copy to avoid modifying input in place
+    out_current_page_annotator: dict = current_page_annotator
+
+    # Get the target page number from the selected row
+    # Safely access the page number, handling potential errors or empty DataFrame
+    gradio_annotator_current_page_number: int = 0
+    annotate_previous_page: int = 0 # Renaming for clarity if needed, matches original output
+    if not selected_recogniser_entity_df_row.empty and 'page' in selected_recogniser_entity_df_row.columns:
+        try:
+            # Use .iloc[0] and .item() for robust scalar extraction
+            gradio_annotator_current_page_number = int(selected_recogniser_entity_df_row['page'].iloc[0])
+            annotate_previous_page = gradio_annotator_current_page_number # Store original page number
+        except (IndexError, ValueError, TypeError):
+            print("Warning: Could not extract valid page number from selected_recogniser_entity_df_row. Defaulting to page 0 (or 1).")
+            gradio_annotator_current_page_number = 1 # Or 0 depending on 1-based vs 0-based indexing elsewhere
+
+    # Ensure page number is valid and 1-based for external display/logic
+    if gradio_annotator_current_page_number <= 0:
+        gradio_annotator_current_page_number = 1
+
+    page_max_reported = len(out_image_annotations_state)
+    if gradio_annotator_current_page_number > page_max_reported:
+        gradio_annotator_current_page_number = page_max_reported # Cap at max pages
+
+    page_num_reported_zero_indexed = gradio_annotator_current_page_number - 1
+
+    # Process page sizes DataFrame early, as it's needed for image path handling and potentially coordinate multiplication
+    page_sizes_df = pd.DataFrame(page_sizes)
+    if not page_sizes_df.empty:
+        # Safely convert page column to numeric and then int
+        page_sizes_df["page"] = pd.to_numeric(page_sizes_df["page"], errors="coerce")
+        page_sizes_df.dropna(subset=["page"], inplace=True)
+        if not page_sizes_df.empty:
+            page_sizes_df["page"] = page_sizes_df["page"].astype(int)
+        else:
+            print("Warning: Page sizes DataFrame became empty after processing.")
 
+    # --- OPTIMIZATION: Process only the current page's data from review_df ---
     if not review_df.empty:
-        #print("review_df just before convert_review_df:", review_df)
-        # First, check that the image on the current page is valid, replace with what exists in page_sizes object if not
-        if not gradio_annotator_current_page_number: gradio_annotator_current_page_number = 0
+        # Filter review_df for the current page
+        # Ensure 'page' column in review_df is comparable to page_num_reported
+        if 'page' in review_df.columns:
+             review_df['page'] = pd.to_numeric(review_df['page'], errors='coerce').fillna(-1).astype(int)
 
-        # Check bounding values for current page and page max
-        if gradio_annotator_current_page_number > 0: page_num_reported = gradio_annotator_current_page_number
-        elif gradio_annotator_current_page_number == 0: page_num_reported = 1 # minimum possible reported page is 1
-        else: 
-            gradio_annotator_current_page_number = 0
-            page_num_reported = 1
+             current_image_path = out_image_annotations_state[page_num_reported_zero_indexed]['image']
 
-        # Ensure page displayed can't exceed number of pages in document
-        page_max_reported = len(out_image_annotations_state)
-        if page_num_reported > page_max_reported: page_num_reported = page_max_reported
+             replaced_image_path, page_sizes_df = replace_placeholder_image_with_real_image(doc_full_file_name_textbox, current_image_path, page_sizes_df, gradio_annotator_current_page_number, input_folder)
 
-        page_num_reported_zero_indexed = page_num_reported - 1        
-        out_image_annotations_state = convert_review_df_to_annotation_json(review_df, image_file_paths, page_sizes)
+             # page_sizes_df has been changed - save back to page_sizes_object
+             page_sizes = page_sizes_df.to_dict(orient='records')
+             review_df.loc[review_df["page"]==gradio_annotator_current_page_number, 'image'] = replaced_image_path
+             images_list = list(page_sizes_df["image_path"])
+             images_list[page_num_reported_zero_indexed] = replaced_image_path
+             out_image_annotations_state[page_num_reported_zero_indexed]['image'] = replaced_image_path
 
-        page_image_annotator_object, out_image_annotations_state = replace_images_in_image_annotation_object(out_image_annotations_state, out_image_annotations_state[page_num_reported_zero_indexed], page_sizes, page_num_reported)    
+             current_page_review_df = review_df[review_df['page'] == gradio_annotator_current_page_number].copy()          
+             current_page_review_df = multiply_coordinates_by_page_sizes(current_page_review_df, page_sizes_df)
 
-        out_image_annotations_state[page_num_reported_zero_indexed] = page_image_annotator_object
+        else:
+            print(f"Warning: 'page' column not found in review_df. Cannot filter for page {gradio_annotator_current_page_number}. Skipping update from review_df.")
+            current_page_review_df = pd.DataFrame() # Empty dataframe if filter fails
+
+        if not current_page_review_df.empty:
+            # Convert the current page's review data to annotation list format for *this page*
+
+            current_page_annotations_list = []
+            # Define expected annotation dict keys, including 'image', 'page', coords, 'label', 'text', 'color' etc.
+            # Assuming review_df has compatible columns
+            expected_annotation_keys = ['label', 'color', 'xmin', 'ymin', 'xmax', 'ymax', 'text', 'id'] # Add/remove as needed
+
+            # Ensure necessary columns exist in current_page_review_df before converting rows
+            for key in expected_annotation_keys:
+                 if key not in current_page_review_df.columns:
+                      # Add missing column with default value
+                      # Use np.nan for numeric, '' for string/object
+                      default_value = np.nan if key in ['xmin', 'ymin', 'xmax', 'ymax'] else ''
+                      current_page_review_df[key] = default_value
+
+            # Convert filtered DataFrame rows to list of dicts
+            # Using .to_dict(orient='records') is efficient for this
+            current_page_annotations_list_raw = current_page_review_df[expected_annotation_keys].to_dict(orient='records')
+
+            current_page_annotations_list = current_page_annotations_list_raw
+
+            # Update the annotations state for the current page
+            # Each entry in out_image_annotations_state seems to be a dict containing keys like 'image', 'page', 'annotations' (List[dict])
+            # Need to update the 'annotations' list for the specific page.
+            # Find the entry for the current page in the state
+            page_state_entry_found = False
+            for i, page_state_entry in enumerate(out_image_annotations_state):
+                # Assuming page_state_entry has a 'page' key (1-based)
+
+                match = re.search(r"(\d+)\.png$", page_state_entry['image'])
+                if match: page_no = int(match.group(1))
+                else: page_no = -1
+
+                if 'image' in page_state_entry and page_no == page_num_reported_zero_indexed:
+                    # Replace the annotations list for this page with the new list from review_df
+                    out_image_annotations_state[i]['boxes'] = current_page_annotations_list
+
+                    # Update the image path as well, based on review_df if available, or keep existing
+                    # Assuming review_df has an 'image' column for this page
+                    if 'image' in current_page_review_df.columns and not current_page_review_df.empty:
+                         # Use the image path from the first row of the filtered review_df
+                         out_image_annotations_state[i]['image'] = current_page_review_df['image'].iloc[0]
+                    page_state_entry_found = True
+                    break
+
+            if not page_state_entry_found:
+                 # This scenario might happen if the current_image_annotations_state didn't initially contain
+                 # an entry for this page number. Depending on the application logic, you might need to
+                 # add a new entry here, but based on the original code's structure, it seems
+                 # out_image_annotations_state is pre-populated for all pages.
+                 print(f"Warning: Entry for page {gradio_annotator_current_page_number} not found in current_image_annotations_state. Cannot update page annotations.")
+
+
+    # --- Image Path and Page Size Handling (already seems focused on current page, keep similar logic) ---
+    # Get the image path for the current page from the updated state
+    # Ensure the entry exists before accessing
+    current_image_path = None
+    if len(out_image_annotations_state) > page_num_reported_zero_indexed and 'image' in out_image_annotations_state[page_num_reported_zero_indexed]:
+         current_image_path = out_image_annotations_state[page_num_reported_zero_indexed]['image']
+    else:
+         print(f"Warning: Could not get image path from state for page index {page_num_reported_zero_indexed}.")
+
+
+    # Replace placeholder image with real image path if needed
+    if current_image_path and not page_sizes_df.empty:
+        try:
+            replaced_image_path, page_sizes_df = replace_placeholder_image_with_real_image(
+                doc_full_file_name_textbox, current_image_path, page_sizes_df,
+                gradio_annotator_current_page_number, input_folder # Use 1-based page number
+            )
 
-        out_current_page_annotator = out_image_annotations_state[page_num_reported_zero_indexed]
+            # Update state and review_df with the potentially replaced image path
+            if len(out_image_annotations_state) > page_num_reported_zero_indexed:
+                 out_image_annotations_state[page_num_reported_zero_indexed]['image'] = replaced_image_path
 
-    return out_current_page_annotator, out_image_annotations_state
+            if 'page' in review_df.columns and 'image' in review_df.columns:
+                 review_df.loc[review_df["page"]==gradio_annotator_current_page_number, 'image'] = replaced_image_path
+
+        except Exception as e:
+             print(f"Error during image path replacement for page {gradio_annotator_current_page_number}: {e}")
+
+
+    # Save back page_sizes_df to page_sizes list format
+    if not page_sizes_df.empty:
+        page_sizes = page_sizes_df.to_dict(orient='records')
+    else:
+        page_sizes = [] # Ensure page_sizes is a list if df is empty
+
+    # --- Re-evaluate Coordinate Multiplication and Duplicate Removal ---
+    # The original code multiplied coordinates for the *entire* document and removed duplicates
+    # across the *entire* document *after* converting the full review_df to state.
+    # With the optimized approach, we updated only one page's annotations in the state.
+
+    # Let's assume remove_duplicate_images_with_blank_boxes expects the raw list of dicts state format:
+    try:
+         out_image_annotations_state = remove_duplicate_images_with_blank_boxes(out_image_annotations_state)
+    except Exception as e:
+         print(f"Error during duplicate removal: {e}. Proceeding without duplicate removal.")
+
+
+    # Select the current page's annotation object from the (potentially updated) state
+    if len(out_image_annotations_state) > page_num_reported_zero_indexed:
+         out_current_page_annotator = out_image_annotations_state[page_num_reported_zero_indexed]
+    else:
+         print(f"Warning: Cannot select current page annotator object for index {page_num_reported_zero_indexed}.")
+         out_current_page_annotator = {} # Or None, depending on expected output type
+
+
+    # The original code returns gradio_annotator_current_page_number as the 3rd value,
+    # which was potentially updated by bounding checks. Keep this.
+    final_page_number_returned = gradio_annotator_current_page_number
+
+    return (out_current_page_annotator,
+            out_image_annotations_state,
+            final_page_number_returned,
+            page_sizes,
+            review_df, # review_df might have its 'page' column type changed, keep it as is or revert if necessary
+            annotate_previous_page) # The original page number from selected_recogniser_entity_df_row
 
 def exclude_selected_items_from_redaction(review_df: pd.DataFrame,
                                           selected_rows_df: pd.DataFrame,
@@ -216,7 +370,7 @@ def exclude_selected_items_from_redaction(review_df: pd.DataFrame,
                                           page_sizes:List[dict],
                                           image_annotations_state:dict,
                                           recogniser_entity_dataframe_base:pd.DataFrame):
-    '''
+    ''' 
     Remove selected items from the review dataframe from the annotation object and review dataframe.
     '''
 
@@ -253,149 +407,267 @@ def exclude_selected_items_from_redaction(review_df: pd.DataFrame,
 
     return out_review_df, out_image_annotations_state, out_recogniser_entity_dataframe_base, backup_review_state, backup_image_annotations_state, backup_recogniser_entity_dataframe_base
 
-def update_annotator_object_and_filter_df(
-                    all_image_annotations:List[AnnotatedImageData],
-                    gradio_annotator_current_page_number:int,
-                    recogniser_entities_dropdown_value:str="ALL",
-                    page_dropdown_value:str="ALL",
-                    text_dropdown_value:str="ALL",
-                    recogniser_dataframe_base:gr.Dataframe=gr.Dataframe(pd.DataFrame(data={"page":[], "label":[], "text":[], "id":[]}), type="pandas", headers=["page", "label", "text", "id"], show_fullscreen_button=True, wrap=True, show_search='filter', max_height=400, static_columns=[0,1,2,3]),
-                    zoom:int=100,
-                    review_df:pd.DataFrame=[],
-                    page_sizes:List[dict]=[],
-                    doc_full_file_name_textbox:str='',
-                    input_folder:str=INPUT_FOLDER):
-    '''
-    Update a gradio_image_annotation object with new annotation data.
-    '''
-    zoom_str = str(zoom) + '%'
-
-    #print("all_image_annotations at start of update_annotator_object_and_filter_df[-1]:", all_image_annotations[-1])
-    
-    if not gradio_annotator_current_page_number: gradio_annotator_current_page_number = 0
-
-    # Check bounding values for current page and page max
-    if gradio_annotator_current_page_number > 0: page_num_reported = gradio_annotator_current_page_number
-    elif gradio_annotator_current_page_number == 0: page_num_reported = 1 # minimum possible reported page is 1
-    else: 
-        gradio_annotator_current_page_number = 0
-        page_num_reported = 1
+def replace_annotator_object_img_np_array_with_page_sizes_image_path(
+        all_image_annotations:List[dict],
+        page_image_annotator_object:AnnotatedImageData,
+        page_sizes:List[dict],
+        page:int):
 
-    # Ensure page displayed can't exceed number of pages in document
-    page_max_reported = len(all_image_annotations)
-    if page_num_reported > page_max_reported: page_num_reported = page_max_reported
+        '''
+        Check if the image value in an AnnotatedImageData dict is a placeholder or np.array. If either of these, replace the value with the file path of the image that is hopefully already loaded into the app related to this page.
+        '''
 
-    page_num_reported_zero_indexed = page_num_reported - 1
+        page_zero_index = page - 1
+        
+        if isinstance(all_image_annotations[page_zero_index]["image"], np.ndarray) or "placeholder_image" in all_image_annotations[page_zero_index]["image"] or isinstance(page_image_annotator_object['image'], np.ndarray):
+            page_sizes_df = pd.DataFrame(page_sizes)
+            page_sizes_df[["page"]] = page_sizes_df[["page"]].apply(pd.to_numeric, errors="coerce")
 
-    # First, check that the image on the current page is valid, replace with what exists in page_sizes object if not
-    page_image_annotator_object, all_image_annotations = replace_images_in_image_annotation_object(all_image_annotations, all_image_annotations[page_num_reported_zero_indexed], page_sizes, page_num_reported)    
+            # Check for matching pages
+            matching_paths = page_sizes_df.loc[page_sizes_df['page'] == page, "image_path"].unique()
 
-    all_image_annotations[page_num_reported_zero_indexed] = page_image_annotator_object
-    
-    current_image_path = all_image_annotations[page_num_reported_zero_indexed]['image']
+            if matching_paths.size > 0:
+                image_path = matching_paths[0]
+                page_image_annotator_object['image'] = image_path
+                all_image_annotations[page_zero_index]["image"] = image_path
+            else:
+                print(f"No image path found for page {page}.")
 
-    # If image path is still not valid, load in a new image an overwrite it. Then replace all items in the image annotation object for all pages based on the updated information.
-    page_sizes_df = pd.DataFrame(page_sizes)
+        return page_image_annotator_object, all_image_annotations
 
-    if not os.path.exists(current_image_path):        
+def replace_placeholder_image_with_real_image(doc_full_file_name_textbox:str, current_image_path:str, page_sizes_df:pd.DataFrame, page_num_reported:int, input_folder:str):
+        ''' If image path is still not valid, load in a new image an overwrite it. Then replace all items in the image annotation object for all pages based on the updated information.'''
+        page_num_reported_zero_indexed = page_num_reported - 1
 
-        page_num, replaced_image_path, width, height = process_single_page_for_image_conversion(doc_full_file_name_textbox, page_num_reported_zero_indexed, input_folder=input_folder)
+        if not os.path.exists(current_image_path):        
 
-        # Overwrite page_sizes values 
-        page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_width"] = width
-        page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_height"] = height
-        page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_path"] = replaced_image_path
-    
-    else:
-        if not page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_width"].isnull().all():
-            width = page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_width"].max()
-            height = page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_height"].max()      
-        else:
-            image = Image.open(current_image_path)
-            width = image.width
-            height = image.height
+            page_num, replaced_image_path, width, height = process_single_page_for_image_conversion(doc_full_file_name_textbox, page_num_reported_zero_indexed, input_folder=input_folder)
 
+            # Overwrite page_sizes values 
             page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_width"] = width
             page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_height"] = height
+            page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_path"] = replaced_image_path
+        
+        else:
+            if not page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_width"].isnull().all():
+                width = page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_width"].max()
+                height = page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_height"].max()      
+            else:
+                image = Image.open(current_image_path)
+                width = image.width
+                height = image.height
 
-        page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_path"] = current_image_path
+                page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_width"] = width
+                page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_height"] = height
 
-        replaced_image_path = current_image_path
+            page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_path"] = current_image_path
 
-    if review_df.empty: review_df = pd.DataFrame(columns=["image", "page", "label", "color", "xmin", "ymin", "xmax", "ymax", "text", "id"])
-    review_df.loc[review_df["page"]==page_num_reported, 'image'] = replaced_image_path
+            replaced_image_path = current_image_path
+        
+        return replaced_image_path, page_sizes_df
 
-    # Update dropdowns and review selection dataframe with the updated annotator object
-    recogniser_entities_list, recogniser_dataframe_out_gr, recogniser_dataframe_modified, recogniser_entities_dropdown_value, text_entities_drop, page_entities_drop = update_recogniser_dataframes(all_image_annotations, recogniser_dataframe_base, recogniser_entities_dropdown_value, text_dropdown_value, page_dropdown_value, review_df.copy(), page_sizes)
-    
-    recogniser_colour_list = [(0, 0, 0) for _ in range(len(recogniser_entities_list))]
+def update_annotator_object_and_filter_df(
+    all_image_annotations:List[AnnotatedImageData],
+    gradio_annotator_current_page_number:int,
+    recogniser_entities_dropdown_value:str="ALL",
+    page_dropdown_value:str="ALL",
+    text_dropdown_value:str="ALL",
+    recogniser_dataframe_base:gr.Dataframe=None, # Simplified default
+    zoom:int=100,
+    review_df:pd.DataFrame=None, # Use None for default empty DataFrame
+    page_sizes:List[dict]=[],
+    doc_full_file_name_textbox:str='',
+    input_folder:str=INPUT_FOLDER
+) -> Tuple[image_annotator, gr.Number, gr.Number, int, str, gr.Dataframe, pd.DataFrame, List[str], List[str], List[dict], List[AnnotatedImageData]]:
+    '''
+    Update a gradio_image_annotation object with new annotation data for the current page
+    and update filter dataframes, optimizing by processing only the current page's data for display.
+    '''
+    zoom_str = str(zoom) + '%'
+
+    # Handle default empty review_df and recogniser_dataframe_base
+    if review_df is None or not isinstance(review_df, pd.DataFrame):
+         review_df = pd.DataFrame(columns=["image", "page", "label", "color", "xmin", "ymin", "xmax", "ymax", "text", "id"])
+    if recogniser_dataframe_base is None: # Create a simple default if None
+         recogniser_dataframe_base = gr.Dataframe(pd.DataFrame(data={"page":[], "label":[], "text":[], "id":[]}))
+
+
+    # Handle empty all_image_annotations state early
+    if not all_image_annotations:
+        print("No all_image_annotation object found")
+        # Return blank/default outputs
+        blank_annotator = gr.ImageAnnotator(
+            value = None, boxes_alpha=0.1, box_thickness=1, label_list=[], label_colors=[],
+            show_label=False, height=zoom_str, width=zoom_str, box_min_size=1,
+            box_selected_thickness=2, handle_size=4, sources=None,
+            show_clear_button=False, show_share_button=False, show_remove_button=False,
+            handles_cursor=True, interactive=True, use_default_label=True
+        )
+        blank_df_out_gr = gr.Dataframe(pd.DataFrame(columns=["page", "label", "text", "id"]))
+        blank_df_modified = pd.DataFrame(columns=["page", "label", "text", "id"])
 
-    # page_sizes_df has been changed - save back to page_sizes_object
-    page_sizes = page_sizes_df.to_dict(orient='records')
+        return (blank_annotator, gr.Number(value=1), gr.Number(value=1), 1,
+                recogniser_entities_dropdown_value, blank_df_out_gr, blank_df_modified,
+                [], [], [], []) # Return empty lists/defaults for other outputs
+
+    # Validate and bound the current page number (1-based logic)
+    page_num_reported = max(1, gradio_annotator_current_page_number) # Minimum page is 1
+    page_max_reported = len(all_image_annotations)
+    if page_num_reported > page_max_reported:
+        page_num_reported = page_max_reported
 
-    images_list = list(page_sizes_df["image_path"])
-    images_list[page_num_reported_zero_indexed] = replaced_image_path
+    page_num_reported_zero_indexed = page_num_reported - 1
+    annotate_previous_page = page_num_reported # Store the determined page number
 
-    all_image_annotations[page_num_reported_zero_indexed]['image'] = replaced_image_path
+    # --- Process page sizes DataFrame ---
+    page_sizes_df = pd.DataFrame(page_sizes)
+    if not page_sizes_df.empty:
+        page_sizes_df["page"] = pd.to_numeric(page_sizes_df["page"], errors="coerce")
+        page_sizes_df.dropna(subset=["page"], inplace=True)
+        if not page_sizes_df.empty:
+            page_sizes_df["page"] = page_sizes_df["page"].astype(int)
+        else:
+            print("Warning: Page sizes DataFrame became empty after processing.")
+
+    # --- Handle Image Path Replacement for the Current Page ---
+    # This modifies the specific page entry within all_image_annotations list
+    # Assuming replace_annotator_object_img_np_array_with_page_sizes_image_path
+    # correctly updates the image path within the list element.
+    if len(all_image_annotations) > page_num_reported_zero_indexed:
+        # Make a shallow copy of the list and deep copy the specific page dict before modification
+        # to avoid modifying the input list unexpectedly if it's used elsewhere.
+        # However, the original code modified the list in place, so we'll stick to that
+        # pattern but acknowledge it.
+        page_object_to_update = all_image_annotations[page_num_reported_zero_indexed]
+
+        # Use the helper function to replace the image path within the page object
+        # Note: This helper returns the potentially modified page_object and the full state.
+        # The full state return seems redundant if only page_object_to_update is modified.
+        # Let's call it and assume it correctly updates the item in the list.
+        updated_page_object, all_image_annotations_after_img_replace = replace_annotator_object_img_np_array_with_page_sizes_image_path(
+             all_image_annotations, page_object_to_update, page_sizes, page_num_reported)
+
+        # The original code immediately re-assigns all_image_annotations.
+        # We'll rely on the function modifying the list element in place or returning the updated list.
+        # Assuming it returns the updated list for robustness:
+        all_image_annotations = all_image_annotations_after_img_replace
+
+
+        # Now handle the actual image file path replacement using replace_placeholder_image_with_real_image
+        current_image_path = updated_page_object.get('image') # Get potentially updated image path
+
+        if current_image_path and not page_sizes_df.empty:
+            try:
+                replaced_image_path, page_sizes_df = replace_placeholder_image_with_real_image(
+                    doc_full_file_name_textbox, current_image_path, page_sizes_df,
+                    page_num_reported, input_folder=input_folder # Use 1-based page num
+                )
 
-    # Multiply out image_annotation coordinates from relative to absolute if necessary
-    all_image_annotations_df = convert_annotation_data_to_dataframe(all_image_annotations)
+                # Update the image path in the state and review_df for the current page
+                # Find the correct entry in all_image_annotations list again by index
+                if len(all_image_annotations) > page_num_reported_zero_indexed:
+                     all_image_annotations[page_num_reported_zero_indexed]['image'] = replaced_image_path
 
-    all_image_annotations_df = multiply_coordinates_by_page_sizes(all_image_annotations_df, page_sizes_df, xmin="xmin", xmax="xmax", ymin="ymin", ymax="ymax")   
+                # Update review_df's image path for this page
+                if 'page' in review_df.columns and 'image' in review_df.columns:
+                     # Ensure review_df page column is numeric for filtering
+                     review_df['page'] = pd.to_numeric(review_df['page'], errors='coerce').fillna(-1).astype(int)
+                     review_df.loc[review_df["page"]==page_num_reported, 'image'] = replaced_image_path
 
-    #print("all_image_annotations_df[-1] just before creating annotation dicts:", all_image_annotations_df.iloc[-1, :])
 
-    all_image_annotations = create_annotation_dicts_from_annotation_df(all_image_annotations_df, page_sizes)
+            except Exception as e:
+                 print(f"Error during image path replacement for page {page_num_reported}: {e}")
+    else:
+         print(f"Warning: Page index {page_num_reported_zero_indexed} out of bounds for all_image_annotations list.")
 
-    #print("all_image_annotations[-1] after creating annotation dicts:", all_image_annotations[-1])
 
+    # Save back page_sizes_df to page_sizes list format
+    if not page_sizes_df.empty:
+        page_sizes = page_sizes_df.to_dict(orient='records')
+    else:
+        page_sizes = [] # Ensure page_sizes is a list if df is empty
+
+    # --- OPTIMIZATION: Prepare data *only* for the current page for display ---
+    current_page_image_annotator_object = None
+    if len(all_image_annotations) > page_num_reported_zero_indexed:
+        page_data_for_display = all_image_annotations[page_num_reported_zero_indexed]
+
+        # Convert current page annotations list to DataFrame for coordinate multiplication IF needed
+        # Assuming coordinate multiplication IS needed for display if state stores relative coords
+        current_page_annotations_df = convert_annotation_data_to_dataframe([page_data_for_display])
+
+
+        if not current_page_annotations_df.empty and not page_sizes_df.empty:
+             # Multiply coordinates *only* for this page's DataFrame
+             try:
+                 # Need the specific page's size for multiplication
+                 page_size_row = page_sizes_df[page_sizes_df['page'] == page_num_reported]
+                 if not page_size_row.empty:
+                      current_page_annotations_df = multiply_coordinates_by_page_sizes(
+                          current_page_annotations_df, page_size_row, # Pass only the row for the current page
+                          xmin="xmin", xmax="xmax", ymin="ymin", ymax="ymax"
+                      )
+             
+             except Exception as e:
+                  print(f"Warning: Error during coordinate multiplication for page {page_num_reported}: {e}. Using original coordinates.")
+                  # If error, proceed with original coordinates or handle as needed
+
+        if "color" not in current_page_annotations_df.columns:
+            current_page_annotations_df['color'] = '(0, 0, 0)'
+
+        # Convert the processed DataFrame back to the list of dicts format for the annotator
+        processed_current_page_annotations_list = current_page_annotations_df[["xmin", "xmax", "ymin", "ymax", "label", "color", "text", "id"]].to_dict(orient='records')
+
+        # Construct the final object expected by the Gradio ImageAnnotator value parameter
+        current_page_image_annotator_object: AnnotatedImageData = {
+            'image': page_data_for_display.get('image'), # Use the (potentially updated) image path
+            'boxes': processed_current_page_annotations_list
+        }
 
+    # --- Update Dropdowns and Review DataFrame ---
+    # This external function still operates on potentially large DataFrames.
+    # It receives all_image_annotations and a copy of review_df.
+    try:
+        recogniser_entities_list, recogniser_dataframe_out_gr, recogniser_dataframe_modified, recogniser_entities_dropdown_value, text_entities_drop, page_entities_drop = update_recogniser_dataframes(
+             all_image_annotations, # Pass the updated full state
+             recogniser_dataframe_base,
+             recogniser_entities_dropdown_value,
+             text_dropdown_value,
+             page_dropdown_value,
+             review_df.copy(), # Keep the copy as per original function call
+             page_sizes # Pass updated page sizes
+        )
+        # Generate default black colors for labels if needed by image_annotator
+        recogniser_colour_list = [(0, 0, 0) for _ in range(len(recogniser_entities_list))]
 
-    # Remove blank duplicate entries
-    all_image_annotations = remove_duplicate_images_with_blank_boxes(all_image_annotations)
+    except Exception as e:
+        print(f"Error calling update_recogniser_dataframes: {e}. Returning empty/default filter data.")
+        recogniser_entities_list = []
+        recogniser_colour_list = []
+        recogniser_dataframe_out_gr = gr.Dataframe(pd.DataFrame(columns=["page", "label", "text", "id"]))
+        recogniser_dataframe_modified = pd.DataFrame(columns=["page", "label", "text", "id"])
+        text_entities_drop = []
+        page_entities_drop = []
 
-    current_page_image_annotator_object = all_image_annotations[page_num_reported_zero_indexed]
 
-    #print("current_page_image_annotator_object that goes into annotator object:", current_page_image_annotator_object)
+    # --- Final Output Components ---
+    page_number_reported_gradio_comp = gr.Number(label = "Current page", value=page_num_reported, precision=0)
 
-    page_number_reported_gradio = gr.Number(label = "Current page", value=page_num_reported, precision=0)
     
-    ###
-    # If no data, present a blank page
-    if not all_image_annotations:
-        print("No all_image_annotation object found")
-        page_num_reported = 1
 
-        out_image_annotator = image_annotator(
-        value = None,
-        boxes_alpha=0.1,
-        box_thickness=1,
-        label_list=recogniser_entities_list,
-        label_colors=recogniser_colour_list,
-        show_label=False,
-        height=zoom_str,
-        width=zoom_str,
-        box_min_size=1,
-        box_selected_thickness=2,
-        handle_size=4,
-        sources=None,#["upload"],
-        show_clear_button=False,
-        show_share_button=False,
-        show_remove_button=False,
-        handles_cursor=True,
-        interactive=True,
-        use_default_label=True
-        )        
-
-        return out_image_annotator, page_number_reported_gradio, page_number_reported_gradio, page_num_reported, recogniser_entities_dropdown_value, recogniser_dataframe_out_gr, recogniser_dataframe_modified, text_entities_drop, page_entities_drop, page_sizes, all_image_annotations
-    
+    ### Present image_annotator outputs
+    # Handle the case where current_page_image_annotator_object couldn't be prepared
+    if current_page_image_annotator_object is None:
+        # This should ideally be covered by the initial empty check for all_image_annotations,
+        # but as a safeguard:
+        print("Warning: Could not prepare annotator object for the current page.")
+        out_image_annotator = image_annotator(value=None, interactive=False) # Present blank/non-interactive
     else:
-        ### Present image_annotator outputs
         out_image_annotator = image_annotator(
             value = current_page_image_annotator_object,
             boxes_alpha=0.1,
             box_thickness=1,
-            label_list=recogniser_entities_list,
+            label_list=recogniser_entities_list, # Use labels from update_recogniser_dataframes
             label_colors=recogniser_colour_list,
             show_label=False,
             height=zoom_str,
@@ -408,41 +680,23 @@ def update_annotator_object_and_filter_df(
             show_share_button=False,
             show_remove_button=False,
             handles_cursor=True,
-            interactive=True
+            interactive=True # Keep interactive if data is present
         )
 
-    #print("all_image_annotations at end of update_annotator...:", all_image_annotations)
-    #print("review_df at end of update_annotator_object:", review_df)
-
-    return out_image_annotator, page_number_reported_gradio, page_number_reported_gradio, page_num_reported, recogniser_entities_dropdown_value, recogniser_dataframe_out_gr, recogniser_dataframe_modified, text_entities_drop, page_entities_drop, page_sizes, all_image_annotations
-
-def replace_images_in_image_annotation_object(
-        all_image_annotations:List[dict],
-        page_image_annotator_object:AnnotatedImageData,
-        page_sizes:List[dict],
-        page:int):
-
-        '''
-        Check if the image value in an AnnotatedImageData dict is a placeholder or np.array. If either of these, replace the value with the file path of the image that is hopefully already loaded into the app related to this page.
-        '''
-
-        page_zero_index = page - 1
-        
-        if isinstance(all_image_annotations[page_zero_index]["image"], np.ndarray) or "placeholder_image" in all_image_annotations[page_zero_index]["image"] or isinstance(page_image_annotator_object['image'], np.ndarray):
-            page_sizes_df = pd.DataFrame(page_sizes)
-            page_sizes_df[["page"]] = page_sizes_df[["page"]].apply(pd.to_numeric, errors="coerce")
-
-            # Check for matching pages
-            matching_paths = page_sizes_df.loc[page_sizes_df['page'] == page, "image_path"].unique()
-
-            if matching_paths.size > 0:
-                image_path = matching_paths[0]
-                page_image_annotator_object['image'] = image_path
-                all_image_annotations[page_zero_index]["image"] = image_path
-            else:
-                print(f"No image path found for page {page}.")
-
-        return page_image_annotator_object, all_image_annotations
+    # The original code returned page_number_reported_gradio twice;
+    # returning the Gradio component and the plain integer value.
+    # Let's match the output signature.
+    return (out_image_annotator,
+            page_number_reported_gradio_comp,
+            page_number_reported_gradio_comp, # Redundant, but matches original return signature
+            page_num_reported, # Plain integer value
+            recogniser_entities_dropdown_value,
+            recogniser_dataframe_out_gr,
+            recogniser_dataframe_modified,
+            text_entities_drop, # List of text entities for dropdown
+            page_entities_drop, # List of page numbers for dropdown
+            page_sizes, # Updated page_sizes list
+            all_image_annotations) # Return the updated full state
 
 def update_all_page_annotation_object_based_on_previous_page(
                                     page_image_annotator_object:AnnotatedImageData,
@@ -459,12 +713,9 @@ def update_all_page_annotation_object_based_on_previous_page(
     previous_page_zero_index = previous_page -1
  
     if not current_page: current_page = 1
-
-    #print("page_image_annotator_object at start of update_all_page_annotation_object:", page_image_annotator_object)
-      
-    page_image_annotator_object, all_image_annotations = replace_images_in_image_annotation_object(all_image_annotations, page_image_annotator_object, page_sizes, previous_page)
-
-    #print("page_image_annotator_object after replace_images in update_all_page_annotation_object:", page_image_annotator_object)
+    
+    # This replaces the numpy array image object with the image file path
+    page_image_annotator_object, all_image_annotations = replace_annotator_object_img_np_array_with_page_sizes_image_path(all_image_annotations, page_image_annotator_object, page_sizes, previous_page)
 
     if clear_all == False: all_image_annotations[previous_page_zero_index] = page_image_annotator_object
     else: all_image_annotations[previous_page_zero_index]["boxes"] = []
@@ -493,7 +744,7 @@ def apply_redactions_to_review_df_and_files(page_image_annotator_object:Annotate
     page_image_annotator_object = all_image_annotations[current_page - 1]   
 
     # This replaces the numpy array image object with the image file path
-    page_image_annotator_object, all_image_annotations = replace_images_in_image_annotation_object(all_image_annotations, page_image_annotator_object, page_sizes, current_page)
+    page_image_annotator_object, all_image_annotations = replace_annotator_object_img_np_array_with_page_sizes_image_path(all_image_annotations, page_image_annotator_object, page_sizes, current_page)
     page_image_annotator_object['image'] = all_image_annotations[current_page - 1]["image"]
 
     if not page_image_annotator_object:
@@ -529,7 +780,7 @@ def apply_redactions_to_review_df_and_files(page_image_annotator_object:Annotate
                         # Check if all elements are integers in the range 0-255
                         if all(isinstance(c, int) and 0 <= c <= 255 for c in fill):
                             pass
-                            #print("fill:", fill)
+
                         else:
                             print(f"Invalid color values: {fill}. Defaulting to black.")
                             fill = (0, 0, 0)  # Default to black if invalid
@@ -553,7 +804,6 @@ def apply_redactions_to_review_df_and_files(page_image_annotator_object:Annotate
                 doc = [image]
 
             elif file_extension in '.csv':
-                #print("This is a csv")
                 pdf_doc = []
 
             # If working with pdfs
@@ -797,11 +1047,9 @@ def df_select_callback(df: pd.DataFrame, evt: gr.SelectData):
 
         row_value_df = pd.DataFrame(data={"page":[row_value_page], "label":[row_value_label], "text":[row_value_text], "id":[row_value_id]})
 
-        return row_value_page, row_value_df
+        return row_value_df
 
 def df_select_callback_textract_api(df: pd.DataFrame, evt: gr.SelectData):
-        
-        #print("evt.data:", evt._data)
 
         row_value_job_id = evt.row_value[0] # This is the page number value
         # row_value_label = evt.row_value[1] # This is the label number value
@@ -829,59 +1077,108 @@ def df_select_callback_ocr(df: pd.DataFrame, evt: gr.SelectData):
 
         return row_value_page, row_value_df
 
-def update_selected_review_df_row_colour(redaction_row_selection:pd.DataFrame, review_df:pd.DataFrame, previous_id:str="", previous_colour:str='(0, 0, 0)', page_sizes:List[dict]=[], output_folder:str=OUTPUT_FOLDER, colour:str='(1, 0, 255)'):
+def update_selected_review_df_row_colour(
+    redaction_row_selection: pd.DataFrame,
+    review_df: pd.DataFrame,
+    previous_id: str = "",
+    previous_colour: str = '(0, 0, 0)',
+    colour: str = '(1, 0, 255)'
+) -> tuple[pd.DataFrame, str, str]:
     '''
     Update the colour of a single redaction box based on the values in a selection row
+    (Optimized Version)
     '''
-    colour_tuple =  str(tuple(colour))
 
-    if "color" not in review_df.columns: review_df["color"] = '(0, 0, 0)'
+    # Ensure 'color' column exists, default to previous_colour if previous_id is provided
+    if "color" not in review_df.columns:
+        review_df["color"] = previous_colour if previous_id else '(0, 0, 0)'
+
+    # Ensure 'id' column exists
     if "id" not in review_df.columns:
-        review_df = fill_missing_ids(review_df)
+         # Assuming fill_missing_ids is a defined function that returns a DataFrame
+         # It's more efficient if this is handled outside if possible,
+         # or optimized internally.
+         print("Warning: 'id' column not found. Calling fill_missing_ids.")
+         review_df = fill_missing_ids(review_df) # Keep this if necessary, but note it can be slow
+
+    # --- Optimization 1 & 2: Reset existing highlight colours using vectorized assignment ---
+    # Reset the color of the previously highlighted row
+    if previous_id and previous_id in review_df["id"].values:
+         review_df.loc[review_df["id"] == previous_id, "color"] = previous_colour
+
+    # Reset the color of any row that currently has the highlight colour (handle cases where previous_id might not have been tracked correctly)
+    # Convert to string for comparison only if the dtype might be mixed or not purely string
+    # If 'color' is consistently string, the .astype(str) might be avoidable.
+    # Assuming color is consistently string format like '(R, G, B)'
+    review_df.loc[review_df["color"] == colour, "color"] = '(0, 0, 0)'
 
-    # Reset existing highlight colours
-    review_df.loc[review_df["id"]==previous_id, "color"] = review_df.loc[review_df["id"]==previous_id, "color"].apply(lambda _: previous_colour)
-    review_df.loc[review_df["color"].astype(str)==colour, "color"] = review_df.loc[review_df["color"].astype(str)==colour, "color"].apply(lambda _: '(0, 0, 0)')
 
     if not redaction_row_selection.empty and not review_df.empty:
         use_id = (
-            "id" in redaction_row_selection.columns 
-            and "id" in review_df.columns 
-            and not redaction_row_selection["id"].isnull().all() 
+            "id" in redaction_row_selection.columns
+            and "id" in review_df.columns
+            and not redaction_row_selection["id"].isnull().all()
             and not review_df["id"].isnull().all()
         )
 
-        selected_merge_cols = ["id"] if use_id else ["label", "page", "text"]        
+        selected_merge_cols = ["id"] if use_id else ["label", "page", "text"]
 
-        review_df = review_df.merge(redaction_row_selection[selected_merge_cols], on=selected_merge_cols, indicator=True, how="left")
+        # --- Optimization 3: Use inner merge directly ---
+        # Merge to find rows in review_df that match redaction_row_selection
+        merged_reviews = review_df.merge(
+            redaction_row_selection[selected_merge_cols],
+            on=selected_merge_cols,
+            how="inner" # Use inner join as we only care about matches
+        )
 
-    if "_merge" in review_df.columns:
-        filtered_reviews = review_df.loc[review_df["_merge"]=="both"]
-    else:
-        filtered_reviews = pd.DataFrame()
+        if not merged_reviews.empty:
+             # Assuming we only expect one match for highlighting a single row
+             # If multiple matches are possible and you want to highlight all,
+             # the logic for previous_id and previous_colour needs adjustment.
+            new_previous_colour = str(merged_reviews["color"].iloc[0])
+            new_previous_id = merged_reviews["id"].iloc[0]
+
+            # --- Optimization 1 & 2: Update color of the matched row using vectorized assignment ---
+
+            if use_id:
+                 # Faster update if using unique 'id' as merge key
+                 review_df.loc[review_df["id"].isin(merged_reviews["id"]), "color"] = colour
+            else:
+                 # More general case using multiple columns - might be slower
+                 # Create a temporary key for comparison
+                 def create_merge_key(df, cols):
+                     return df[cols].astype(str).agg('_'.join, axis=1)
+
+                 review_df_key = create_merge_key(review_df, selected_merge_cols)
+                 merged_reviews_key = create_merge_key(merged_reviews, selected_merge_cols)
+
+                 review_df.loc[review_df_key.isin(merged_reviews_key), "color"] = colour
+
+            previous_colour = new_previous_colour
+            previous_id = new_previous_id
+        else:
+             # No rows matched the selection
+             print("No reviews found matching selection criteria")
+             # The reset logic at the beginning already handles setting color to (0, 0, 0)
+             # if it was the highlight colour and didn't match.
+             # No specific action needed here for color reset beyond what's done initially.
+             previous_colour = '(0, 0, 0)' # Reset previous_colour as no row was highlighted
+             previous_id = '' # Reset previous_id
 
-    if not filtered_reviews.empty:
-        previous_colour = str(filtered_reviews["color"].values[0])
-        previous_id = filtered_reviews["id"].values[0]
-        review_df.loc[review_df["_merge"]=="both", "color"] =  review_df.loc[review_df["_merge"] == "both", "color"].apply(lambda _: colour)
     else:
-        # Handle the case where no rows match the condition
-        print("No reviews found with _merge == 'both'")
-        previous_colour = '(0, 0, 0)'
-        review_df.loc[review_df["color"]==colour, "color"] = previous_colour
-        previous_id =''
+         # If selection is empty, reset any existing highlights
+         review_df.loc[review_df["color"] == colour, "color"] = '(0, 0, 0)'
+         previous_colour = '(0, 0, 0)'
+         previous_id = ''
 
-    review_df.drop("_merge", axis=1, inplace=True)
 
-    # Ensure that all output coordinates are in proportional size
-    #page_sizes_df = pd.DataFrame(page_sizes)
-    #page_sizes_df .loc[:, "page"] = pd.to_numeric(page_sizes_df["page"], errors="coerce")
-    #print("review_df before divide:", review_df)
-    #print("page_sizes_df before divide:", page_sizes_df)
-    #review_df = divide_coordinates_by_page_sizes(review_df, page_sizes_df)
-    #print("review_df after divide:", review_df)
+    # Ensure column order is maintained if necessary, though pandas generally preserves order
+    # Creating a new DataFrame here might involve copying data, consider if this is strictly needed.
+    if set(["image", "page", "label", "color", "xmin","ymin", "xmax", "ymax", "text", "id"]).issubset(review_df.columns):
+        review_df = review_df[["image", "page", "label", "color", "xmin","ymin", "xmax", "ymax", "text", "id"]]
+    else:
+         print("Warning: Not all expected columns are present in review_df for reordering.")
 
-    review_df = review_df[["image", "page", "label", "color", "xmin","ymin", "xmax", "ymax", "text", "id"]]
 
     return review_df, previous_id, previous_colour
 
@@ -988,8 +1285,6 @@ def create_xfdf(review_file_df:pd.DataFrame, pdf_path:str, pymupdf_doc:object, i
         page_sizes_df = pd.DataFrame(page_sizes)
 
         # If there are no image coordinates, then convert coordinates to pymupdf coordinates prior to export
-        #print("Using pymupdf coordinates for conversion.")
-
         pages_are_images = False
 
         if "mediabox_width" not in review_file_df.columns:            
@@ -1041,33 +1336,9 @@ def create_xfdf(review_file_df:pd.DataFrame, pdf_path:str, pymupdf_doc:object, i
                 raise ValueError(f"Invalid cropbox format: {document_cropboxes[page_python_format]}")
         else:
             print("Document cropboxes not found.")
-
         
         pdf_page_height = pymupdf_page.mediabox.height
-        pdf_page_width = pymupdf_page.mediabox.width
-
-        # Check if image dimensions for page exist in page_sizes_df
-        # image_dimensions = {}
-
-        # image_dimensions['image_width'] = page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_width"].max()
-        # image_dimensions['image_height'] = page_sizes_df.loc[page_sizes_df['page']==page_num_reported, "image_height"].max()
-
-        # if pd.isna(image_dimensions['image_width']):
-        #     image_dimensions = {}
-
-        # image = image_paths[page_python_format]
-
-        # if image_dimensions:
-        #     image_page_width, image_page_height = image_dimensions["image_width"], image_dimensions["image_height"]
-        # if isinstance(image, str) and 'placeholder' not in image:
-        #     image = Image.open(image)
-        #     image_page_width, image_page_height = image.size
-        # else:
-        #     try:
-        #         image = Image.open(image)
-        #         image_page_width, image_page_height = image.size
-        #     except Exception as e:
-        #         print("Could not get image sizes due to:", e)        
+        pdf_page_width = pymupdf_page.mediabox.width     
 
         # Create redaction annotation
         redact_annot = SubElement(annots, 'redact')
@@ -1345,8 +1616,6 @@ def convert_xfdf_to_dataframe(file_paths_list:List[str], pymupdf_doc, image_path
                 # Optionally, you can add the image path or other relevant information
                 df.loc[_, 'image'] = image_path
 
-                #print('row:', row)
-
     out_file_path = output_folder + file_path_name + "_review_file.csv"
     df.to_csv(out_file_path, index=None)