+
+
+ HTML visualization showing metrics and comparison results from the OCR pipeline
+
+ Side-by-side comparison of OCR results across different models
+
{selected_model_display} OCR'
+ st.markdown(logo_html, unsafe_allow_html=True)
+ else:
+ st.title(f"{selected_model_display} OCR")
+
+
+def render_sidebar():
+ """Render the sidebar with model selection and settings."""
+ with st.sidebar:
+ st.header("Settings")
+
+ comparison_mode = st.radio(
+ "Mode",
+ ["Single Model", "Compare Models"],
+ horizontal=True,
+ )
+
+ # Get unique models (avoid duplicates from shorthand keys)
+ unique_models = {
+ model_id: model_config
+ for model_id, model_config in MODEL_CONFIGS.items()
+ if model_id == model_config.name
+ }
+
+ if comparison_mode == "Single Model":
+ model_options = [
+ model_config.display for model_id, model_config in unique_models.items()
+ ]
+
+ # Find default model index
+ default_index = 0
+ for i, option in enumerate(model_options):
+ if option == DEFAULT_MODEL.display:
+ default_index = i
+ break
+
+ model_choice = st.selectbox(
+ "Select OCR Model",
+ options=model_options,
+ index=default_index,
+ key="model_selector_single",
+ )
+
+ # Map display name back to model ID
+ selected_model_id = None
+ for model_id, model_config in unique_models.items():
+ if model_config.display == model_choice:
+ selected_model_id = model_id
+ break
+
+ st.session_state["comparison_mode"] = False
+ st.session_state["selected_model_id"] = selected_model_id
+ else:
+ # Multiselect for model comparison
+ st.subheader("Models to Compare")
+ model_display_to_id = {
+ model_config.display: model_id for model_id, model_config in unique_models.items()
+ }
+ default_models = [DEFAULT_MODEL.display]
+ previously_selected = st.session_state.get("selected_model_displays", default_models)
+ multiselect_key = f"model_multiselect_{len(previously_selected)}"
+
+ selected_model_displays = st.multiselect(
+ "Select models to compare",
+ options=list(model_display_to_id.keys()),
+ default=previously_selected,
+ key=multiselect_key,
+ )
+
+ if not selected_model_displays:
+ st.warning("At least one model must be selected. Using default model.")
+ selected_model_displays = default_models
+
+ if selected_model_displays != previously_selected:
+ st.session_state["selected_model_displays"] = selected_model_displays
+ st.rerun()
+
+ comparison_models = [
+ model_display_to_id[display] for display in selected_model_displays
+ ]
+
+ st.session_state["comparison_mode"] = True
+ st.session_state["comparison_models"] = comparison_models
+
+ st.header("Custom Prompt")
+ use_custom_prompt = st.checkbox("Use custom prompt", value=False)
+ custom_prompt = (
+ st.text_area(
+ "Enter your custom prompt",
+ "Extract and analyze all text from the image and identify any objects present.",
+ height=100,
+ )
+ if use_custom_prompt
+ else ""
+ )
+
+ st.header("Upload Image")
+ uploaded_file = st.file_uploader(
+ "Choose an image...",
+ type=["png", "jpg", "jpeg", "webp", "tiff"],
+ key=st.session_state.get("file_uploader_key", "default_uploader"),
+ )
+
+ # Store uploaded file in session state
+ if uploaded_file is not None:
+ st.session_state["uploaded_file"] = uploaded_file
+
+ st.markdown("---")
+ display_provider_info()
+
+ return st.session_state.get(
+ "uploaded_file", uploaded_file
+ ), custom_prompt if use_custom_prompt else None
+
+
+def display_provider_info():
+ """Display information about which providers are being used."""
+ providers = set()
+ for model_id, model_config in MODEL_CONFIGS.items():
+ if model_id == model_config.name: # Skip shorthand duplicates
+ providers.add(model_config.ocr_processor)
+
+ provider_names = {"ollama": "Ollama", "openai": "OpenAI", "mistral": "MistralAI"}
+ provider_text = " + ".join([provider_names.get(p, p.capitalize()) for p in providers])
+ st.caption(f"Powered by {provider_text}")
+
+
+def add_clear_button():
+ """Add a clear button to reset the app state."""
+ col1, col2 = st.columns([6, 1])
+ with col2:
+ if st.button("Clear 🗑️"):
+ keys_to_clear = ["ocr_result", "uploaded_file", "processed_image"]
+ for model_id in MODEL_CONFIGS:
+ result_key = f"{model_id}_result"
+ keys_to_clear.append(result_key)
+
+ for key in keys_to_clear:
+ if key in st.session_state:
+ del st.session_state[key]
+
+ # Force file uploader to reset
+ st.session_state["file_uploader_key"] = str(time.time())
+ st.rerun()
+
+
+def process_uploaded_image(image):
+ """Process and display the uploaded image."""
+ img_width, img_height = image.size
+ max_width, max_height = 600, 500
+
+ # Calculate scaling factor to maintain aspect ratio
+ width_ratio = max_width / img_width
+ height_ratio = max_height / img_height
+ scale_factor = min(width_ratio, height_ratio)
+
+ # Only resize if needed
+ if img_width > max_width or img_height > max_height:
+ new_width = int(img_width * scale_factor)
+ new_height = int(img_height * scale_factor)
+ image = image.resize((new_width, new_height))
+
+ # Store the processed image in session state
+ st.session_state["processed_image"] = image
+ st.image(image, caption="Uploaded Image", use_container_width=False)
+ return image
+
+
+def check_for_error(result):
+ """Check if the OCR model has an error."""
+ if "error" in result:
+ return True
+
+ raw_text = result.get("raw_text", "")
+ if not isinstance(raw_text, str):
+ raw_text = str(raw_text)
+
+ return raw_text.startswith("Error:") or ("success" in result and result["success"] is False)
+
+
+def has_no_text(result):
+ """Check if the OCR model has no text."""
+ raw_text = result.get("raw_text", "")
+ if not isinstance(raw_text, str):
+ raw_text = str(raw_text)
+ return raw_text == "No text found"
+
+
+def display_result(label, result, proc_time, model_id, model_logos):
+ """Display the results of the OCR model."""
+ model_config = MODEL_CONFIGS[model_id]
+ model_display = model_config.display
+
+ # Display header with logo if available
+ if model_display in model_logos:
+ header_html = f'
Extract structured text from images using your chosen OCR model!
', + unsafe_allow_html=True, + ) + + processed_image = None + + if uploaded_file is not None: + image = Image.open(uploaded_file) + processed_image = process_uploaded_image(image) + elif "processed_image" in st.session_state: + processed_image = st.session_state["processed_image"] + st.image(processed_image, caption="Uploaded Image", use_container_width=False) + + if processed_image is not None: + if st.button("Extract Text 🔍", type="primary"): + comparison_mode = st.session_state.get("comparison_mode", False) + + if comparison_mode: + comparison_models = st.session_state.get("comparison_models", [DEFAULT_MODEL.name]) + run_multiple_models(processed_image, comparison_models, custom_prompt) + else: + selected_model_id = st.session_state.get("selected_model_id", DEFAULT_MODEL.name) + run_single_model(processed_image, selected_model_id, custom_prompt) + else: + st.info("Upload an image and click 'Extract Text' to see the results here.") + + # Footer + st.markdown("---") + st.caption("ZenOCR - Comparing LLM OCR capabilities") + + +if __name__ == "__main__": + main() diff --git a/omni-reader/assets/docs/metrics.png b/omni-reader/assets/docs/metrics.png new file mode 100644 index 00000000..7f59b482 Binary files /dev/null and b/omni-reader/assets/docs/metrics.png differ diff --git a/omni-reader/assets/docs/pipeline_dags.png b/omni-reader/assets/docs/pipeline_dags.png new file mode 100644 index 00000000..90ca4cbc Binary files /dev/null and b/omni-reader/assets/docs/pipeline_dags.png differ diff --git a/omni-reader/assets/docs/streamlit.png b/omni-reader/assets/docs/streamlit.png new file mode 100644 index 00000000..ce18d03b Binary files /dev/null and b/omni-reader/assets/docs/streamlit.png differ diff --git a/omni-reader/assets/docs/visualization.png b/omni-reader/assets/docs/visualization.png new file mode 100644 index 00000000..d0390bfc Binary files /dev/null and b/omni-reader/assets/docs/visualization.png differ diff --git a/omni-reader/assets/logos/default.svg b/omni-reader/assets/logos/default.svg new file mode 100644 index 00000000..b709fab6 --- /dev/null +++ b/omni-reader/assets/logos/default.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/omni-reader/assets/logos/gemma.svg b/omni-reader/assets/logos/gemma.svg new file mode 100644 index 00000000..116ec367 --- /dev/null +++ b/omni-reader/assets/logos/gemma.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/omni-reader/assets/logos/microsoft.svg b/omni-reader/assets/logos/microsoft.svg new file mode 100644 index 00000000..5334aa7c --- /dev/null +++ b/omni-reader/assets/logos/microsoft.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/omni-reader/assets/logos/mistral.svg b/omni-reader/assets/logos/mistral.svg new file mode 100644 index 00000000..4e97da2a --- /dev/null +++ b/omni-reader/assets/logos/mistral.svg @@ -0,0 +1,18 @@ + \ No newline at end of file diff --git a/omni-reader/assets/logos/ollama.svg b/omni-reader/assets/logos/ollama.svg new file mode 100644 index 00000000..cc887e3d --- /dev/null +++ b/omni-reader/assets/logos/ollama.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/omni-reader/assets/logos/openai.svg b/omni-reader/assets/logos/openai.svg new file mode 100644 index 00000000..3b4eff96 --- /dev/null +++ b/omni-reader/assets/logos/openai.svg @@ -0,0 +1,2 @@ + + \ No newline at end of file diff --git a/omni-reader/assets/samples_for_ocr/easy_example.jpeg b/omni-reader/assets/samples_for_ocr/easy_example.jpeg new file mode 100644 index 00000000..7a158f02 Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/easy_example.jpeg differ diff --git a/omni-reader/assets/samples_for_ocr/education_article_excerpt.webp b/omni-reader/assets/samples_for_ocr/education_article_excerpt.webp new file mode 100644 index 00000000..f5e8b3eb Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/education_article_excerpt.webp differ diff --git a/omni-reader/assets/samples_for_ocr/incomplete_sentence.png b/omni-reader/assets/samples_for_ocr/incomplete_sentence.png new file mode 100644 index 00000000..6510a1f4 Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/incomplete_sentence.png differ diff --git a/omni-reader/assets/samples_for_ocr/lexus_vin_number.webp b/omni-reader/assets/samples_for_ocr/lexus_vin_number.webp new file mode 100644 index 00000000..cce02b1f Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/lexus_vin_number.webp differ diff --git a/omni-reader/assets/samples_for_ocr/montreal_signs.jpg b/omni-reader/assets/samples_for_ocr/montreal_signs.jpg new file mode 100644 index 00000000..24f014d6 Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/montreal_signs.jpg differ diff --git a/omni-reader/assets/samples_for_ocr/paris_signs.jpg b/omni-reader/assets/samples_for_ocr/paris_signs.jpg new file mode 100644 index 00000000..da8936ae Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/paris_signs.jpg differ diff --git a/omni-reader/assets/samples_for_ocr/reporter_notes.png b/omni-reader/assets/samples_for_ocr/reporter_notes.png new file mode 100644 index 00000000..27d320b1 Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/reporter_notes.png differ diff --git a/omni-reader/assets/samples_for_ocr/rx_prescription_clear.jpg b/omni-reader/assets/samples_for_ocr/rx_prescription_clear.jpg new file mode 100644 index 00000000..b7e8c077 Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/rx_prescription_clear.jpg differ diff --git a/omni-reader/assets/samples_for_ocr/rx_prescription_unclear.png b/omni-reader/assets/samples_for_ocr/rx_prescription_unclear.png new file mode 100644 index 00000000..75440102 Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/rx_prescription_unclear.png differ diff --git a/omni-reader/assets/samples_for_ocr/tire_serial_number.jpg b/omni-reader/assets/samples_for_ocr/tire_serial_number.jpg new file mode 100644 index 00000000..2f2a0e6a Binary files /dev/null and b/omni-reader/assets/samples_for_ocr/tire_serial_number.jpg differ diff --git a/omni-reader/configs/batch_pipeline.yaml b/omni-reader/configs/batch_pipeline.yaml new file mode 100644 index 00000000..3400d2c4 --- /dev/null +++ b/omni-reader/configs/batch_pipeline.yaml @@ -0,0 +1,84 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# pipeline configuration +build: batch-ocr-pipeline-20254007 +run_name: run_ocr + +# environment configuration +settings: + docker: + requirements: requirements.txt + required_integrations: + - aws + - s3 + python_package_installer: uv + environment: + OPENAI_API_KEY: ${{ env.OPENAI_API_KEY }} + MISTRAL_API_KEY: ${{ env.MISTRAL_API_KEY }} + +# enable flags +enable_artifact_metadata: True +enable_artifact_visualization: True +enable_cache: False +enable_step_logs: True + +# step configuration +steps: + load_images: + parameters: + image_folder: ./assets/samples_for_ocr + image_paths: [] + enable_cache: False + + run_ocr: + parameters: + custom_prompt: null + models: # can be model names or shorthands + - pixtral + - gemma3 + - llava-phi3 + - gpt4o + - granite + enable_cache: False + +# vision models configuration +models_registry: + - name: mistral/pixtral-12b-2409 + shorthand: pixtral + ocr_processor: litellm + provider: mistral + + - name: gpt-4o-mini + shorthand: gpt4o + ocr_processor: openai + + - name: gemma3:27b + shorthand: gemma3 + ocr_processor: ollama + + - name: llava:34b + shorthand: llava34b + ocr_processor: ollama + + - name: llava-phi3 + shorthand: llava-phi3 + ocr_processor: ollama + + - name: granite3.2-vision + shorthand: granite + ocr_processor: ollama diff --git a/omni-reader/configs/evaluation_pipeline.yaml b/omni-reader/configs/evaluation_pipeline.yaml new file mode 100644 index 00000000..f3f1f241 --- /dev/null +++ b/omni-reader/configs/evaluation_pipeline.yaml @@ -0,0 +1,55 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# pipeline configuration +build: ocr-evaluation-pipeline-20254007 +run_name: ocr_evaluation_run + +# environment configuration +settings: + docker: + requirements: requirements.txt + required_integrations: + - aws + - s3 + python_package_installer: uv + environment: + OPENAI_API_KEY: ${{ env.OPENAI_API_KEY }} + MISTRAL_API_KEY: ${{ env.MISTRAL_API_KEY }} + +# enable flags +enable_artifact_metadata: True +enable_artifact_visualization: True +enable_cache: False +enable_step_logs: True + +# steps configuration +steps: + load_ocr_results: # loads OCR results from batch pipeline runs + enable_cache: False + parameters: + artifact_name: ocr_results + version: null + + load_ground_truth_texts: + enable_cache: False + parameters: + ground_truth_folder: ground_truth_texts + ground_truth_files: [] + + evaluate_models: + enable_cache: False diff --git a/omni-reader/ground_truth_texts/education_article_excerpt.txt b/omni-reader/ground_truth_texts/education_article_excerpt.txt new file mode 100644 index 00000000..9ebe4c8f --- /dev/null +++ b/omni-reader/ground_truth_texts/education_article_excerpt.txt @@ -0,0 +1,13 @@ +Only a matter of style? + +For educational purposes we analyse the opening pages of an 11-page article that appeared in The American Mathematical Monthly, Volume 102 Number 2 / February 1995. We have added line numbers in the right margin. + +line 4: Since in this article, squares don’t get alternating colours, it could be argued that the term “chessboard” is misplaced. + +line 4: The introduction of the name “B” seems unnecessary; it is used in the combination “the board B” in the text for Figure 1 and in line 7; in both cases just “the board” would have done fine. In line 77 occurs the last use of B, viz. in “X⊂B”, which is dubious since B was a board and not a set; in line 77, I would have preferred “Given a set X of cells”. + +line 7/8: The first move, being a move like any other, does not deserve a separate description. The term “step” is redundant. + +line 8: Why not “a move consists of”? + +line 10/11: At this stage the italics are puzzling, since a move is possible if, diff --git a/omni-reader/ground_truth_texts/incomplete_sentence.txt b/omni-reader/ground_truth_texts/incomplete_sentence.txt new file mode 100644 index 00000000..a5924119 --- /dev/null +++ b/omni-reader/ground_truth_texts/incomplete_sentence.txt @@ -0,0 +1 @@ +In mid-April Anglesey moved his family and entourage from Rome to Naples, there to await the arrival of diff --git a/omni-reader/ground_truth_texts/lexus_vin_number.txt b/omni-reader/ground_truth_texts/lexus_vin_number.txt new file mode 100644 index 00000000..767a51f4 --- /dev/null +++ b/omni-reader/ground_truth_texts/lexus_vin_number.txt @@ -0,0 +1 @@ +JTHBH5D2405012812 \ No newline at end of file diff --git a/omni-reader/ground_truth_texts/montreal_signs.txt b/omni-reader/ground_truth_texts/montreal_signs.txt new file mode 100644 index 00000000..c31b8cf6 --- /dev/null +++ b/omni-reader/ground_truth_texts/montreal_signs.txt @@ -0,0 +1,3 @@ +Basilique Notre-Dame +Place Royale +Place d’Armes \ No newline at end of file diff --git a/omni-reader/ground_truth_texts/paris_signs.txt b/omni-reader/ground_truth_texts/paris_signs.txt new file mode 100644 index 00000000..56b98b30 --- /dev/null +++ b/omni-reader/ground_truth_texts/paris_signs.txt @@ -0,0 +1,5 @@ +Palais Royal +Les Arts Décoratifs +Musée du LOUVRE +Église ST GERMAIN l’AUXERROIS +Musée Eugène DELACROIX \ No newline at end of file diff --git a/omni-reader/ground_truth_texts/reporter_notes.txt b/omni-reader/ground_truth_texts/reporter_notes.txt new file mode 100644 index 00000000..03596fac --- /dev/null +++ b/omni-reader/ground_truth_texts/reporter_notes.txt @@ -0,0 +1 @@ +An attempt to get more information about the Admiralty House meeting will be made in the House of Commons this afternoon. Labour M.P.s already have many questions to the Prime Minister asking for a statement. President Kennedy flew from London Airport last night to arrive in Washington this morning. He is to make a 30-minute nation-wide broadcast and television report on his talks with Mr. Khrushchev this evening. diff --git a/omni-reader/ground_truth_texts/rx_prescription_clear.txt b/omni-reader/ground_truth_texts/rx_prescription_clear.txt new file mode 100644 index 00000000..d0cc4230 --- /dev/null +++ b/omni-reader/ground_truth_texts/rx_prescription_clear.txt @@ -0,0 +1,6 @@ +Rx +Amoxicillin + Clavulanic acid (Co-Amoxiclav) 500/125 mg/tab #21 +Sig: Take one with food every 8 hours for 7 days + +Paracetamol 500 mg/tab #5 +Sig: Take one with food every 4 hours as needed for fever (temp. ≥ 37.8°C) diff --git a/omni-reader/ground_truth_texts/rx_prescription_unclear.txt b/omni-reader/ground_truth_texts/rx_prescription_unclear.txt new file mode 100644 index 00000000..dc53da0a --- /dev/null +++ b/omni-reader/ground_truth_texts/rx_prescription_unclear.txt @@ -0,0 +1,10 @@ +Rx + +Tab. Pansec 20 +1 + 0 + 1 + +Tab. Apitez 160 +0 + 1 + 0 + +Tab. Linagliptin / Linita 5mg +2 + 0 + 0 \ No newline at end of file diff --git a/omni-reader/ground_truth_texts/tire_serial_number.txt b/omni-reader/ground_truth_texts/tire_serial_number.txt new file mode 100644 index 00000000..8fc900d6 --- /dev/null +++ b/omni-reader/ground_truth_texts/tire_serial_number.txt @@ -0,0 +1 @@ +3702692432 \ No newline at end of file diff --git a/omni-reader/pipelines/__init__.py b/omni-reader/pipelines/__init__.py new file mode 100644 index 00000000..f90f1c83 --- /dev/null +++ b/omni-reader/pipelines/__init__.py @@ -0,0 +1,19 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""OCR pipelines""" + +from pipelines.batch_pipeline import batch_ocr_pipeline, run_batch_ocr_pipeline +from pipelines.evaluation_pipeline import ocr_evaluation_pipeline, run_ocr_evaluation_pipeline diff --git a/omni-reader/pipelines/batch_pipeline.py b/omni-reader/pipelines/batch_pipeline.py new file mode 100644 index 00000000..913ffe14 --- /dev/null +++ b/omni-reader/pipelines/batch_pipeline.py @@ -0,0 +1,105 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""OCR Batch Pipeline implementation for processing images with multiple models.""" + +import os +from typing import Any, Dict, List, Optional + +from dotenv import load_dotenv +from zenml import pipeline +from zenml.config import DockerSettings +from zenml.logger import get_logger + +from steps import ( + load_images, + run_ocr, +) + +load_dotenv() + +logger = get_logger(__name__) + +docker_settings = DockerSettings( + required_integrations=["s3", "aws"], + python_package_installer="uv", + requirements="requirements.txt", + environment={ + "OPENAI_API_KEY": os.getenv("OPENAI_API_KEY"), + "MISTRAL_API_KEY": os.getenv("MISTRAL_API_KEY"), + }, +) + + +@pipeline(settings={"docker": docker_settings}) +def batch_ocr_pipeline( + image_paths: Optional[List[str]] = None, + image_folder: Optional[str] = None, + custom_prompt: Optional[str] = None, + models: List[str] = None, +) -> None: + """Run OCR batch processing pipeline with multiple models. + + Args: + image_paths: Optional list of specific image paths to process + image_folder: Optional folder to search for images + custom_prompt: Optional custom prompt to use for the models + models: List of model names to use for OCR + """ + images = load_images( + image_paths=image_paths, + image_folder=image_folder, + ) + + run_ocr( + images=images, + models=models, + custom_prompt=custom_prompt, + ) + + +def run_batch_ocr_pipeline(config: Dict[str, Any]) -> None: + """Run the OCR batch pipeline from a configuration dictionary. + + Args: + config: Dictionary containing configuration + + Returns: + None + """ + pipeline_instance = batch_ocr_pipeline.with_options( + enable_cache=config.get("enable_cache", False), + ) + + load_images_params = config.get("steps", {}).get("load_images", {}).get("parameters", {}) + image_folder = load_images_params.get("image_folder") + image_paths = load_images_params.get("image_paths", []) + if not image_folder and len(image_paths) == 0: + raise ValueError("Either image_folder or image_paths must be provided") + + run_ocr_params = config.get("steps", {}).get("run_ocr", {}).get("parameters", {}) + custom_prompt = run_ocr_params.get("custom_prompt") + selected_models = run_ocr_params.get("models", []) + if not selected_models or len(selected_models) == 0: + raise ValueError( + "No models found in the run_ocr step of the batch_ocr_pipeline config file. At least one model must be specified in the 'models' parameter." + ) + + pipeline_instance( + image_paths=image_paths, + image_folder=image_folder, + custom_prompt=custom_prompt, + models=selected_models, + ) diff --git a/omni-reader/pipelines/evaluation_pipeline.py b/omni-reader/pipelines/evaluation_pipeline.py new file mode 100644 index 00000000..a1b7a21c --- /dev/null +++ b/omni-reader/pipelines/evaluation_pipeline.py @@ -0,0 +1,99 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""OCR Evaluation Pipeline implementation for comparing models using existing results.""" + +import os +from typing import Any, Dict, List, Optional + +from dotenv import load_dotenv +from zenml import pipeline +from zenml.config import DockerSettings +from zenml.logger import get_logger + +from steps import ( + evaluate_models, + load_ground_truth_texts, + load_ocr_results, +) + +load_dotenv() + +logger = get_logger(__name__) + +docker_settings = DockerSettings( + requirements="requirements.txt", + required_integrations=["s3", "aws"], + python_package_installer="uv", + environment={ + "OPENAI_API_KEY": os.getenv("OPENAI_API_KEY"), + "MISTRAL_API_KEY": os.getenv("MISTRAL_API_KEY"), + }, +) + + +@pipeline(settings={"docker": docker_settings}) +def ocr_evaluation_pipeline( + ground_truth_folder: Optional[str] = None, + ground_truth_files: Optional[List[str]] = None, +) -> None: + """Run OCR evaluation pipeline comparing existing model results.""" + if not ground_truth_folder and not ground_truth_files: + raise ValueError( + "Either ground_truth_folder or ground_truth_files must be provided for evaluation" + ) + + model_results = load_ocr_results(artifact_name="ocr_results") + + ground_truth_df = load_ground_truth_texts( + model_results=model_results, + ground_truth_folder=ground_truth_folder, + ground_truth_files=ground_truth_files, + ) + + evaluate_models( + model_results=model_results, + ground_truth_df=ground_truth_df, + ) + + +def run_ocr_evaluation_pipeline(config: Dict[str, Any]) -> None: + """Run the OCR evaluation pipeline from a configuration dictionary. + + Args: + config: Dictionary containing configuration + + Returns: + None + """ + mode = config.get("parameters", {}).get("mode", "evaluation") + if mode != "evaluation": + logger.warning(f"Expected mode 'evaluation', but got '{mode}'. Proceeding anyway.") + + pipeline_instance = ocr_evaluation_pipeline.with_options( + enable_artifact_metadata=config.get("enable_artifact_metadata", True), + enable_artifact_visualization=config.get("enable_artifact_visualization", True), + enable_cache=config.get("enable_cache", False), + enable_step_logs=config.get("enable_step_logs", True), + ) + + load_ground_truth_texts_params = ( + config.get("steps", {}).get("load_ground_truth_texts", {}).get("parameters", {}) + ) + + pipeline_instance( + ground_truth_folder=load_ground_truth_texts_params.get("ground_truth_folder"), + ground_truth_files=load_ground_truth_texts_params.get("ground_truth_files", []), + ) diff --git a/omni-reader/requirements.txt b/omni-reader/requirements.txt new file mode 100644 index 00000000..0c69b01d --- /dev/null +++ b/omni-reader/requirements.txt @@ -0,0 +1,16 @@ +instructor +jiwer +jiter +importlib-metadata<7.0,>=1.4.0 +litellm +mistralai==1.0.3 +numpy<2.0,>=1.9.0 +openai==1.69.0 +Pillow==11.1.0 +polars==1.26.0 +pyarrow>=7.0.0 +python-dotenv +streamlit==1.44.0 +pydantic>=2.8.2,<2.9.0 +tqdm==4.66.4 +zenml>=0.80.0 diff --git a/omni-reader/run.py b/omni-reader/run.py new file mode 100644 index 00000000..4c2eaa28 --- /dev/null +++ b/omni-reader/run.py @@ -0,0 +1,419 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Run OCR pipeline with or without ZenML tracking. + +This module provides two modes of operation: +1. UI Mode: Direct OCR with no metadata/artifact tracking (for Streamlit) +2. Pipeline Mode: Full ZenML pipeline with tracking +""" + +import argparse +import os +import time +from typing import Any, Dict, List, Optional, Union + +from dotenv import load_dotenv +from PIL import Image + +from pipelines.batch_pipeline import run_batch_ocr_pipeline +from pipelines.evaluation_pipeline import run_ocr_evaluation_pipeline +from utils.config import ( + get_image_paths, + list_available_ground_truth_files, + load_config, + override_batch_config, + override_evaluation_config, + print_config_summary, + select_config_path, + validate_batch_config, + validate_evaluation_config, +) +from utils.model_configs import DEFAULT_MODEL, MODEL_CONFIGS +from utils.ocr_processing import run_ocr + +load_dotenv() + + +def run_ocr_from_ui( + image: Union[str, Image.Image], + model: str, + custom_prompt: Optional[str] = None, +) -> Dict[str, Any]: + """Extract text directly using OCR model without ZenML tracking.""" + start_time = time.time() + + # Get model configuration based on model ID + if model not in MODEL_CONFIGS: + return { + "raw_text": f"Error: Model '{model}' not found in MODEL_CONFIGS", + "error": f"Invalid model: {model}", + "processing_time": 0, + "model": model, + } + + try: + # Use the unified run_ocr function with track_metadata=False for UI + result = run_ocr( + image_input=image, + model_ids=model, + custom_prompt=custom_prompt, + track_metadata=False, + ) + + # Ensure processing_time is properly set + if "processing_time" not in result: + result["processing_time"] = time.time() - start_time + + # Ensure model info is set + result["model"] = model + result["display_name"] = MODEL_CONFIGS[model].display + result["ocr_processor"] = MODEL_CONFIGS[model].ocr_processor + + return result + except Exception as e: + processing_time = time.time() - start_time + return { + "raw_text": f"Error: Failed to extract text - {str(e)}", + "error": str(e), + "processing_time": processing_time, + "model": model, + "display_name": MODEL_CONFIGS[model].display, + "ocr_processor": MODEL_CONFIGS[model].ocr_processor, + } + + +def run_models_in_parallel( + image_path: Union[str, Image.Image], + model_ids: List[str], + custom_prompt: Optional[str] = None, +) -> Dict[str, Dict[str, Any]]: + """Process an image with multiple models in parallel.""" + try: + # Use the unified run_ocr function for parallel processing + results = run_ocr( + image_input=image_path, + model_ids=model_ids, + custom_prompt=custom_prompt, + track_metadata=False, + ) + + # Display progress in CLI mode + print(f"Processed image with {len(model_ids)} models") + return results + except Exception as e: + print(f"Error processing models in parallel: {str(e)}") + results = {} + + # Create error results for each model + for model_id in model_ids: + results[model_id] = { + "raw_text": f"Error: {str(e)}", + "error": str(e), + "processing_time": 0, + "model": model_id, + } + + return results + + +def list_supported_models(): + """List all supported models.""" + print("\nSupported models:") + print("-" * 70) + print(f"{'Model ID':<25} {'Display Name':<30} {'OCR Processor':<15}") + print("-" * 70) + + for model_id, config in MODEL_CONFIGS.items(): + print(f"{model_id:<25} {config.display:<30} {config.ocr_processor:<15}") + if config.provider: + print(f"{'Provider':<25} {config.provider:<30}") + + print("\nDefault model:", DEFAULT_MODEL.name) + print("-" * 70) + + +def format_model_results(model_id, result): + """Format results for a specific model.""" + model_config = MODEL_CONFIGS.get(model_id, None) + model_display = model_config.display if model_config else model_id + + output = f"\n{model_display} results:" + + if "error" in result: + output += f"\n❌ Error: {result.get('error', 'Unknown error')}" + else: + text = result["raw_text"] + if len(text) > 150: + text = f"{text[:150]}..." + output += f"\n✅ Text: {text}" + + output += f"\n⏱️ Processing time: {result.get('processing_time', 0):.2f}s" + + if "confidence" in result and result["confidence"] is not None: + output += f"\n🎯 Confidence: {result['confidence']:.2%}" + + return output + + +def run_ui_mode(args, parser): + """Run the application in streamlit UI mode without ZenML tracking.""" + if args.list_models: + list_supported_models() + return + + if args.image_paths: + image_path = args.image_paths[0] # Take the first image for UI mode + elif args.image_folder: + image_paths = get_image_paths(args.image_folder) + if not image_paths: + print(f"No images found in directory: {args.image_folder}") + return + image_path = image_paths[0] # Take the first image for UI mode + else: + parser.error("Error: Please provide an image path or folder") + return + + if not os.path.exists(image_path): + parser.error(f"Error: Image file '{image_path}' not found.") + return + + start_time = time.time() + + if args.models == "all": + # Run all models in parallel + print(f"Processing image with all {len(MODEL_CONFIGS)} models in parallel...") + results = run_models_in_parallel( + image_path, + list(MODEL_CONFIGS.keys()), + args.custom_prompt, + ) + + successful_models = sum(1 for result in results.values() if "error" not in result) + failed_models = len(results) - successful_models + + print("\n" + "=" * 50) + print(f"OCR COMPARISON RESULTS ({successful_models} successful, {failed_models} failed)") + print("=" * 50) + + # individual model results + for model_id, result in results.items(): + print(format_model_results(model_id, result)) + + print(f"\n⏱️ Total time: {time.time() - start_time:.2f}s") + print("=" * 50) + + elif "," in args.models: + # Run specific models in parallel + model_ids = [model_id.strip() for model_id in args.models.split(",")] + + invalid_models = [model_id for model_id in model_ids if model_id not in MODEL_CONFIGS] + if invalid_models: + print(f"Error: The following models are not supported: {', '.join(invalid_models)}") + print("Use --list-models to see all supported models.") + return + + print(f"Processing image with {len(model_ids)} selected models in parallel...") + results = run_models_in_parallel( + image_path, + model_ids, + args.custom_prompt, + ) + + successful_models = sum(1 for result in results.values() if "error" not in result) + failed_models = len(results) - successful_models + + print("\n" + "=" * 50) + print(f"OCR COMPARISON RESULTS ({successful_models} successful, {failed_models} failed)") + print("=" * 50) + + # individual model results + for model_id, result in results.items(): + print(format_model_results(model_id, result)) + + print(f"\n⏱️ Total time: {time.time() - start_time:.2f}s") + print("=" * 50) + + else: + # Run a single model + model_id = args.models + if model_id not in MODEL_CONFIGS: + print(f"Error: Model '{model_id}' not supported.") + print("Use --list-models to see all supported models.") + return + + print(f"\nProcessing with {model_id} model...") + result = run_ocr_from_ui( + image_path, + model_id, + args.custom_prompt, + ) + + print("\n" + "=" * 50) + print(f"OCR RESULT FOR {model_id}") + print("=" * 50) + print(format_model_results(model_id, result)) + print("=" * 50) + + +def run_pipeline_mode(args, parser): + """Run the application in full pipeline mode with ZenML tracking.""" + # List available ground truth files if requested + if args.list_ground_truth_files: + gt_files = list_available_ground_truth_files(directory=args.ground_truth_dir) + if gt_files: + print("Available ground truth files:") + for i, file in enumerate(gt_files): + print(f" {i + 1}. {file}") + else: + print(f"No ground truth files found in '{args.ground_truth_dir}'") + return + + # Determine pipeline mode and select config path + evaluation_mode = args.eval + + if args.config: + config_path = args.config + else: + config_path = select_config_path(evaluation_mode) + print(f"Auto-selecting config file: {config_path}") + + if not os.path.exists(config_path): + parser.error(f"Config file not found: {config_path}") + return + + # Load the configuration + try: + config = load_config(config_path) + except (ValueError, FileNotFoundError) as e: + parser.error(f"Error loading configuration: {str(e)}") + return + + cli_args = { + "image_paths": args.image_paths, + "image_folder": args.image_folder, + "custom_prompt": args.custom_prompt, + "ground_truth_dir": args.ground_truth_dir, + } + + # Override configuration with CLI arguments if provided + try: + if evaluation_mode: + config = override_evaluation_config(config, cli_args) + validate_evaluation_config(config) + else: + config = override_batch_config(config, cli_args) + validate_batch_config(config) + except ValueError as e: + parser.error(f"Configuration error: {str(e)}") + return + + print_config_summary(config, is_evaluation_config=evaluation_mode) + + try: + if evaluation_mode: + print("Running OCR Evaluation Pipeline...") + run_ocr_evaluation_pipeline(config) + else: + print("Running OCR Batch Pipeline...") + run_batch_ocr_pipeline(config) + except Exception as e: + print(f"Error running pipeline: {str(e)}") + return + + +def main(): + """Main entry point for the OCR tool.""" + parser = argparse.ArgumentParser( + description="Run OCR between vision models with or without ZenML tracking" + ) + + # Mode selection + parser.add_argument( + "--ui_mode", + action="store_true", + help="Run in UI mode without ZenML tracking (for Streamlit)", + ) + + # Config file options (pipeline mode) + config_group = parser.add_argument_group("Pipeline Mode Configuration") + config_group.add_argument( + "--config", + type=str, + help="Path to YAML configuration file (for pipeline mode)", + ) + config_group.add_argument( + "--eval", + action="store_true", + help="Run in evaluation pipeline mode (defaults to batch pipeline if not specified)", + ) + + # Ground truth utilities (pipeline mode) + gt_group = parser.add_argument_group("Ground truth utilities (pipeline mode)") + gt_group.add_argument( + "--list-ground-truth-files", + action="store_true", + help="List available ground truth files and exit", + ) + gt_group.add_argument( + "--ground-truth-dir", + type=str, + default="ground_truth_texts", + help="Directory to look for ground truth files (for --list-ground-truth-files)", + ) + + # Quick access options (shared between modes) + input_group = parser.add_argument_group("Input options (shared)") + input_group.add_argument( + "--image-paths", + nargs="+", + help="Paths to images to process", + ) + input_group.add_argument( + "--image-folder", + type=str, + help="Folder containing images to process", + ) + input_group.add_argument( + "--custom-prompt", + type=str, + dest="custom_prompt", + help="Custom prompt to use for OCR models", + ) + + # UI mode specific options + ui_group = parser.add_argument_group("UI Mode Options") + ui_group.add_argument( + "--models", + type=str, + default=DEFAULT_MODEL.name, + help="Model(s) to use: a specific model ID, 'all' to compare all, or a comma-separated list (for UI mode)", + ) + ui_group.add_argument( + "--list-models", + action="store_true", + help="List all supported models and exit (for UI mode)", + ) + + args = parser.parse_args() + + if args.ui_mode: + run_ui_mode(args, parser) + else: + run_pipeline_mode(args, parser) + + +if __name__ == "__main__": + main() diff --git a/omni-reader/schemas/__init__.py b/omni-reader/schemas/__init__.py new file mode 100644 index 00000000..0805615d --- /dev/null +++ b/omni-reader/schemas/__init__.py @@ -0,0 +1,20 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from .image_description import ImageDescription +from .ocr_result import ( + OCRResult, + OCRResultMapping, +) diff --git a/omni-reader/schemas/image_description.py b/omni-reader/schemas/image_description.py new file mode 100644 index 00000000..f7a24e4a --- /dev/null +++ b/omni-reader/schemas/image_description.py @@ -0,0 +1,27 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Schemas for LLM responses.""" + +from typing import Optional + +from pydantic import BaseModel + + +class ImageDescription(BaseModel): + """Base model for OCR results.""" + + raw_text: str + confidence: Optional[float] = None diff --git a/omni-reader/schemas/ocr_result.py b/omni-reader/schemas/ocr_result.py new file mode 100644 index 00000000..43aa74bd --- /dev/null +++ b/omni-reader/schemas/ocr_result.py @@ -0,0 +1,36 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Schemas for OCR results.""" + +from typing import Dict, List + +from pydantic import BaseModel, RootModel + + +class OCRResult(BaseModel): + """OCR result for a single image.""" + + id: int + image_name: str + raw_text: str + processing_time: float + confidence: float + + +class OCRResultMapping(RootModel): + """Each model name maps to a list of OCRResult entries.""" + + root: Dict[str, List[OCRResult]] diff --git a/omni-reader/steps/__init__.py b/omni-reader/steps/__init__.py new file mode 100644 index 00000000..ca6c1d0f --- /dev/null +++ b/omni-reader/steps/__init__.py @@ -0,0 +1,22 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from .evaluate_models import evaluate_models +from .loaders import ( + load_ground_truth_texts, + load_images, + load_ocr_results, +) +from .run_ocr import run_ocr \ No newline at end of file diff --git a/omni-reader/steps/evaluate_models.py b/omni-reader/steps/evaluate_models.py new file mode 100644 index 00000000..a7d564c0 --- /dev/null +++ b/omni-reader/steps/evaluate_models.py @@ -0,0 +1,235 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2025. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""This module contains the steps for evaluating the OCR models.""" + +from typing import Dict + +import polars as pl +from typing_extensions import Annotated +from zenml import log_metadata, step +from zenml.types import HTMLString + +from utils import ( + calculate_custom_metrics, + calculate_model_similarities, + compare_multi_model, + create_model_comparison_card, + create_summary_visualization, + get_model_info, +) + + +@step(enable_cache=False) +def evaluate_models( + model_results: pl.DataFrame, + ground_truth_df: pl.DataFrame, +) -> Annotated[HTMLString, "ocr_visualization"]: + """Compare the performance of multiple configurable models with visualization. + + Args: + model_results: Dictionary containing single or multiple model ocr results + ground_truth_df: DataFrame containing ground truth texts + + Returns: + HTML visualization of the evaluation results + """ + if model_results is None or len(model_results.columns) == 0: + raise ValueError("At least one model is required for evaluation") + + if ground_truth_df is None or ground_truth_df.is_empty(): + raise ValueError("Ground truth data is required for evaluation") + + gt_df = ground_truth_df + + # --- 2. Build model info for evaluation models --- + model_keys = list(model_results.columns) + model_info = {} + model_displays = [] + model_prefixes = {} + for model_name in model_keys: + display, prefix = get_model_info(model_name) + model_info[model_name] = (display, prefix) + model_displays.append(display) + model_prefixes[display] = prefix + + # --- 3. Convert DataFrame rows to dictionaries --- + model_results_dict = {} + for model_name in model_keys: + model_data = model_results[model_name].to_dicts() + model_results_dict[model_name] = pl.DataFrame(model_data) + + # --- 4. Merge evaluation models' results --- + base_model = model_keys[0] + base_display, base_prefix = model_info[base_model] + merged_results = model_results_dict[base_model].clone() + for i, model_name in enumerate(model_keys[1:], start=1): + disp, pref = model_info[model_name] + suffix = f"_{pref}" if i > 1 else "_right" + merged_results = merged_results.join( + model_results_dict[model_name], + on=["id", "image_name"], + how="inner", + suffix=suffix, + ) + + # --- 5. Join ground truth data if available --- + if gt_df is not None: + merged_results = merged_results.join( + gt_df, on=["id", "image_name"], how="inner", suffix="_gt" + ) + + # --- 6. Calculate processing times for evaluation models --- + all_model_times = {} + for model_name, df in model_results.items(): + disp, pref = model_info[model_name] + time_key = f"avg_{pref}_time" + all_model_times[time_key] = df.select("processing_time").to_series().mean() + all_model_times[f"{pref}_display"] = disp + + fastest_model_time, fastest_key = min( + [(time, key) for key, time in all_model_times.items() if not key.endswith("_display")], + key=lambda x: x[0], + ) + fastest_prefix = fastest_key.replace("avg_", "").replace("_time", "") + fastest_display = all_model_times.get(f"{fastest_prefix}_display", fastest_prefix) + + # --- 7. Per-image evaluation: compute metrics, error analysis, and build per-image cards --- + evaluation_metrics = [] + image_cards_html = "" + gt_text_col = "ground_truth_text" + + # Check if we have ground truth data in our joined dataset + if gt_text_col not in merged_results.columns and "raw_text_gt" in merged_results.columns: + gt_text_col = "raw_text_gt" # Fall back to legacy ground truth model format + + for row in merged_results.iter_rows(named=True): + if gt_text_col not in row: + continue + ground_truth = row[gt_text_col] + model_texts = {} + model_texts[base_display] = row["raw_text"] + for i, mkey in enumerate(model_keys[1:], start=1): + disp, pref = model_info[mkey] + col = "raw_text_right" if i == 1 else f"raw_text_{pref}" + if col in row: + model_texts[disp] = row[col] + row_metrics = calculate_custom_metrics(ground_truth, model_texts, list(model_texts.keys())) + error_analysis = compare_multi_model(ground_truth, model_texts) + result_metrics = {"id": row["id"], "image_name": row["image_name"]} + for disp in model_texts.keys(): + if disp in row_metrics: + for met_name, val in row_metrics[disp].items(): + result_metrics[f"{disp} {met_name}"] = val + for disp, errs in error_analysis.items(): + for met_name, val in errs.items(): + result_metrics[f"{disp} {met_name}"] = val + evaluation_metrics.append(result_metrics) + + # Merge GT Similarity values from row_metrics into error_analysis + for disp in model_texts.keys(): + if disp in row_metrics and "GT Similarity" in row_metrics[disp]: + if disp not in error_analysis: + error_analysis[disp] = {} + error_analysis[disp]["GT Similarity"] = row_metrics[disp]["GT Similarity"] + + comparison_card = create_model_comparison_card( + image_name=row["image_name"], + ground_truth=ground_truth, + model_texts=model_texts, + model_metrics=error_analysis, + ) + image_cards_html += comparison_card + + # --- 8. Compute average metrics for evaluation models --- + model_metric_averages = {d: {} for d in model_displays} + if evaluation_metrics: + df_eval = pl.DataFrame(evaluation_metrics) + for disp in model_displays: + cer_col = f"{disp} CER" + wer_col = f"{disp} WER" + sim_col = f"{disp} GT Similarity" + if cer_col in df_eval.columns: + model_metric_averages[disp]["CER"] = df_eval[cer_col].mean() + if wer_col in df_eval.columns: + model_metric_averages[disp]["WER"] = df_eval[wer_col].mean() + if sim_col in df_eval.columns: + model_metric_averages[disp]["GT Similarity"] = df_eval[sim_col].mean() + for disp in model_displays: + pref = model_prefixes[disp] + tkey = f"avg_{pref}_time" + if tkey in all_model_times: + model_metric_averages[disp]["Proc. Time"] = all_model_times[tkey] + + # --- 9. Calculate similarity matrix for evaluation models only --- + sim_results = [] + for row in merged_results.iter_rows(named=True): + texts_map = {} + texts_map[base_display] = row.get("raw_text", "") + for i, mkey in enumerate(model_keys[1:], start=1): + disp, pref = model_info[mkey] + col = "raw_text_right" if i == 1 else f"raw_text_{pref}" + texts_map[disp] = row.get(col, "") + sim_results.append(texts_map) + similarities = {} + if len(model_displays) > 1: + similarities = calculate_model_similarities( + results=[ + { + f"raw_text_{disp.lower().replace(' ', '_')}": texts_map[disp] + for disp in model_displays + } + for texts_map in sim_results + ], + model_displays=model_displays, + ) + + # --- 10. Build time comparison info --- + time_comparison = { + **all_model_times, + "fastest_model": fastest_display, + "model_count": len(model_keys), + } + if len(model_keys) >= 2: + d1, p1 = model_info[model_keys[0]] + d2, p2 = model_info[model_keys[1]] + tk1, tk2 = f"avg_{p1}_time", f"avg_{p2}_time" + if tk1 in all_model_times and tk2 in all_model_times: + time_comparison["time_difference"] = abs(all_model_times[tk1] - all_model_times[tk2]) + + # Log metadata (customize the metadata_dict as needed) + log_metadata( + metadata={ + "fastest_model": fastest_display, + "model_count": len(model_keys), + } + ) + + summary_html = create_summary_visualization( + model_metrics=model_metric_averages, + time_comparison=time_comparison, + similarities=similarities, + ) + + # --- 11. Combine summary and per-image details --- + final_html = f""" + {summary_html} +| Metric | +Value | +
|---|
| Metric | + {headers} +
|---|
'
+ metrics_rows = ""
+ metric_keys = ["CER", "WER", "GT Similarity", "Proc. Time"]
+ for key in metric_keys:
+ if key in metrics:
+ metrics_rows += f"""
+ '
+ model_sections += f"""
+ Higher values indicate more similar outputs between models
+| + {headers} + |
|---|
Pipeline Run: {pipeline_run_name}
+'
+ except Exception as e:
+ logger.warning(f"Error loading logo for {model}: {e}")
+ # Default to just the model name without logo if there's an issue
+ pass
+
+ model_cards += f"""
+ | Image | + {model_header} +Chars | +Confidence | +Time (s) | +Text Preview | +
|---|
Processed {total_results} images with {total_chars} total characters in {total_proc_time:.2f}s
+