move parts of the code to tabs

README.md

@@ -1,5 +1,5 @@
 ---
-title: HTR Pipeline
+title: HTRFLOW
 emoji: 🏢
 colorFrom: purple
 colorTo: green

app.py

@@ -1,89 +1,44 @@
 import os
 
 import gradio as gr
-from apscheduler.schedulers.background import BackgroundScheduler
 
 from helper.gradio_config import css, theme
-from helper.text.text_about import TextAbout
 from helper.text.text_app import TextApp
-from helper.text.text_howto import TextHowTo
-from helper.text.text_roadmap import TextRoadmap
-from helper.utils import add_ip_data, backup_db
+from helper.utils import TrafficDataHandler
+from tabs.about_tab import about_tab
+from tabs.help_tab import help_tab
 from tabs.htr_tool import htr_tool_tab
 from tabs.stepwise_htr_tool import stepwise_htr_tool_tab
 
-SECRET_KEY = os.environ.get("AM_I_IN_A_DOCKER_CONTAINER", False)
+handler = TrafficDataHandler()
 
-if SECRET_KEY:
-    scheduler = BackgroundScheduler()
-    scheduler.add_job(func=backup_db, trigger="interval", seconds=60)
-    scheduler.start()
+VERSION = "Demo version 0.0.2"
 
-
-with gr.Blocks(title="HTR Riksarkivet", theme=theme, css=css) as demo:
+with gr.Blocks(title="Riksarkivet", theme=theme, css=css) as demo:
     with gr.Row():
         with gr.Column(scale=1):
-            text_ip_output = gr.Markdown()
+            text_ip_output = gr.Markdown(VERSION)
         with gr.Column(scale=1):
             gr.Markdown(TextApp.title_markdown)
         with gr.Column(scale=1):
             gr.Markdown(TextApp.title_markdown_img)
 
     with gr.Tabs():
-        with gr.Tab("HTR Tool"):
+        with gr.Tab("Fast track"):
             htr_tool_tab.render()
 
-        with gr.Tab("Stepwise HTR Tool"):
+        with gr.Tab("Stepwise"):
             stepwise_htr_tool_tab.render()
 
-        with gr.Tab("About"):
-            with gr.Tabs():
-                with gr.Tab("Project"):
-                    with gr.Row():
-                        with gr.Column():
-                            gr.Markdown(TextAbout.intro_text)
-                        with gr.Column():
-                            gr.Markdown(TextAbout.text_src_code_data_models)
-                    with gr.Row():
-                        gr.Markdown(TextAbout.pipeline_overview_text)
-                    with gr.Row():
-                        with gr.Tabs():
-                            with gr.Tab("I. Binarization"):
-                                gr.Markdown(TextAbout.binarization)
-                            with gr.Tab("II. Region Segmentation"):
-                                gr.Markdown(TextAbout.text_region_segment)
-                            with gr.Tab("III. Line Segmentation"):
-                                gr.Markdown(TextAbout.text_line_segmentation)
-                            with gr.Tab("IV. Transcriber"):
-                                gr.Markdown(TextAbout.text_htr)
-
-                with gr.Tab("Contribution"):
-                    with gr.Row():
-                        gr.Markdown(TextRoadmap.text_contribution)
+        with gr.Tab("Help"):
+            help_tab.render()
 
-                with gr.Tab("API & Duplicate for Privat use"):
-                    with gr.Row():
-                        with gr.Column():
-                            gr.Markdown(TextHowTo.htr_tool_api_text)
-                            gr.Code(
-                                value=TextHowTo.code_for_api,
-                                language="python",
-                                interactive=False,
-                                show_label=False,
-                            )
-                        with gr.Column():
-                            gr.Markdown(TextHowTo.duplicatin_space_htr_text)
-                            gr.Markdown(TextHowTo.figure_htr_hardware)
-                            gr.Markdown(TextHowTo.duplicatin_for_privat)
-
-                with gr.Tab("Roadmap"):
-                    with gr.Row():
-                        with gr.Column():
-                            gr.Markdown(TextRoadmap.roadmap)
-                        with gr.Column():
-                            gr.Markdown(TextRoadmap.discussion)
+        with gr.Tab("About"):
+            about_tab.render()
 
-    demo.load(add_ip_data)
+    SECRET_KEY = os.environ.get("AM_I_IN_A_DOCKER_CONTAINER", False)
+    if SECRET_KEY:
+        demo.load(handler.onload_store_metric_data)
 
 
 demo.queue(concurrency_count=2, max_size=2)

helper/text/text_app.py

@@ -2,7 +2,7 @@ class TextApp:
     title_markdown = """
 
 
-    <h1><center> HTRflow - Demo </center></h1>
+    <h1><center> HTRFLOW </center></h1>
 
     <h3><center> Swedish National Archives - Riksarkivet </center></h3>"""
 

helper/utils.py

@@ -2,68 +2,89 @@ import hashlib
 import os
 import shutil
 import sqlite3
+import uuid
 from datetime import datetime
 
 import gradio as gr
 import huggingface_hub
 import pandas as pd
 import pytz
-
-
-def hash_ip(ip_address):
-    return hashlib.sha256(ip_address.encode()).hexdigest()
-
-
-def current_time_sw():
-    swedish_tz = pytz.timezone("Europe/Stockholm")
-    return datetime.now(swedish_tz).strftime("%Y-%m-%d %H:%M:%S")
-
-
-def add_ip_data(request: gr.Request):
-    host = request.client.host
-    hashed_ip = hash_ip(host)
-
-    db = sqlite3.connect(DB_FILE)
-    cursor = db.cursor()
-    cursor.execute("INSERT INTO ip_data(current_time, hashed_ip) VALUES(?,?)", [current_time_sw(), hashed_ip])
-    db.commit()
-    db.close()
-
-
-def backup_db():
-    shutil.copyfile(DB_FILE, "./data/traffic_data.db")
-    db = sqlite3.connect(DB_FILE)
-    ip_data = db.execute("SELECT * FROM ip_data").fetchall()
-    pd.DataFrame(ip_data, columns=["id", "current_time", "hashed_ip"]).to_csv("./data/ip_data.csv", index=False)
-
-    print("updating traffic_data")
-    repo.push_to_hub(blocking=False, commit_message=f"Updating data at {datetime.now()}")
-
-
-DB_FILE = "./traffic_data.db"
-
-TOKEN = os.environ.get("HUB_TOKEN")
-repo = huggingface_hub.Repository(
-    local_dir="data", repo_type="dataset", clone_from="Riksarkivet/traffic_demo_data", use_auth_token=TOKEN
-)
-repo.git_pull()
-
-# Set db to latest
-shutil.copyfile("./data/traffic_data.db", DB_FILE)
-
-
-# Create table if it doesn't already exist
-db = sqlite3.connect(DB_FILE)
-try:
-    db.execute("SELECT * FROM ip_data").fetchall()
-    db.close()
-except sqlite3.OperationalError:
-    db.execute(
-        """
-        CREATE TABLE ip_data (id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
-                              current_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL,
-                              hashed_ip TEXT)
-        """
-    )
-    db.commit()
-    db.close()
+from apscheduler.schedulers.background import BackgroundScheduler
+
+
+class TrafficDataHandler:
+    _DB_FILE_PATH = "./traffic_data.db"
+    _DB_TEMP_PATH = "./data/traffic_data.db"
+    _TOKEN = os.environ.get("HUB_TOKEN")
+    _TZ = "Europe/Stockholm"
+    _INTERVAL_MIN_UPDATE = 30
+
+    def __init__(self, dataset_repo="Riksarkivet/traffic_demo_data"):
+        self._repo = huggingface_hub.Repository(
+            local_dir="data", repo_type="dataset", clone_from=dataset_repo, use_auth_token=self._TOKEN
+        )
+        self._pull_repo_data()
+        self._setup_database()
+
+    def _pull_repo_data(self):
+        self._repo.git_pull()
+        shutil.copyfile(self._DB_TEMP_PATH, self._DB_FILE_PATH)
+
+    def _hash_ip(self, ip_address):
+        return hashlib.sha256(ip_address.encode()).hexdigest()
+
+    def _current_time_in_sweden(self):
+        swedish_tz = pytz.timezone(self._TZ)
+        return datetime.now(swedish_tz).strftime("%Y-%m-%d %H:%M:%S")
+
+    def onload_store_metric_data(self, request: gr.Request):
+        self._session_uuid = str(uuid.uuid1())
+        hashed_host = self._hash_ip(request.client.host)
+        self._backup_and_update_database(hashed_host, "load")
+
+    def store_metric_data(self, action, request: gr.Request):
+        self._session_uuid = str(uuid.uuid1())
+        hashed_host = self._hash_ip(request.client.host)
+        self._backup_and_update_database(hashed_host, action)
+
+    def _commit_host_to_database(self, hashed_host, action):
+        with sqlite3.connect(self._DB_FILE_PATH) as db:
+            db.execute(
+                "INSERT INTO ip_data(current_time, hashed_ip, session_uuid, action) VALUES(?,?,?,?)",
+                [self._current_time_in_sweden(), hashed_host, self._session_uuid, action],
+            )
+
+    def _setup_database(self):
+        with sqlite3.connect(self._DB_FILE_PATH) as db:
+            try:
+                db.execute("SELECT * FROM ip_data").fetchall()
+            except sqlite3.OperationalError:
+                db.execute(
+                    """
+                    CREATE TABLE ip_data (id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
+                                          current_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL,
+                                          hashed_ip TEXT,
+                                          session_uuid TEXT,
+                                          action TEXT)
+                    """
+                )
+
+    def _backup_and_update_database(self, hashed_host, action):
+        self._commit_host_to_database(hashed_host, action)
+        shutil.copyfile(self._DB_FILE_PATH, self._DB_TEMP_PATH)
+
+        with sqlite3.connect(self._DB_FILE_PATH) as db:
+            ip_data = db.execute("SELECT * FROM ip_data").fetchall()
+            pd.DataFrame(ip_data, columns=["id", "current_time", "hashed_ip", "session_uuid", "action"]).to_csv(
+                "./data/ip_data.csv", index=False
+            )
+
+        self._repo.push_to_hub(blocking=False, commit_message=f"Updating data at {datetime.now()}")
+
+    def _initialize_and_schedule_backup(self, hashed_host, action):
+        self._backup_and_update_database(hashed_host, action)
+        scheduler = BackgroundScheduler()
+        scheduler.add_job(
+            self._backup_and_update_database, "interval", minutes=self._INTERVAL_MIN_UPDATE, args=(hashed_host, action)
+        )
+        scheduler.start()

requirements.txt

@@ -7,7 +7,7 @@ numpy==1.25.0
 opencv-python-headless==4.7.0.72
 Jinja2==3.1.2
 transformers==4.30.2
-huggingface-hub
+huggingface-hub==0.15.1
 datasets==2.14.5
 requests==2.31.0
 apscheduler

tabs/about_tab.py

@@ -0,0 +1,54 @@
+import gradio as gr
+
+from helper.text.text_about import TextAbout
+from helper.text.text_roadmap import TextRoadmap
+
+with gr.Blocks() as about_tab:
+    with gr.Tabs():
+        with gr.Tab("HTRFLOW"):
+            gr.Markdown(
+                "update... todo.. here we should talk about the pipline and the app as seperate things... pipline overview perhaps be moved?"
+            )
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown(TextAbout.intro_text)
+                with gr.Column():
+                    gr.Markdown(TextAbout.text_src_code_data_models)
+            with gr.Row():
+                gr.Markdown(TextAbout.pipeline_overview_text)
+            with gr.Row():
+                with gr.Tabs():
+                    with gr.Tab("1. Binarization"):
+                        gr.Markdown(TextAbout.binarization)
+                    with gr.Tab("2. Region Segmentation"):
+                        gr.Markdown(TextAbout.text_region_segment)
+                    with gr.Tab("3. Line Segmentation"):
+                        gr.Markdown(TextAbout.text_line_segmentation)
+                    with gr.Tab("4. Transcriber"):
+                        gr.Markdown(TextAbout.text_htr)
+
+        with gr.Tab("Contributions"):
+            with gr.Row():
+                gr.Markdown(TextRoadmap.text_contribution)
+
+        # with gr.Tab("API & Duplicate for own use"):
+        #     with gr.Row():
+        #         with gr.Column():
+        #             gr.Markdown(TextHowTo.htr_tool_api_text)
+        #             gr.Code(
+        #                 value=TextHowTo.code_for_api,
+        #                 language="python",
+        #                 interactive=False,
+        #                 show_label=False,
+        #             )
+        #         with gr.Column():
+        #             gr.Markdown(TextHowTo.duplicatin_space_htr_text)
+        #             gr.Markdown(TextHowTo.figure_htr_hardware)
+        #             gr.Markdown(TextHowTo.duplicatin_for_privat)
+
+        with gr.Tab("Changelog & Roadmap"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown(TextRoadmap.roadmap)
+                with gr.Column():
+                    gr.Markdown(TextRoadmap.discussion)

tabs/help_tab.py

@@ -0,0 +1,35 @@
+import gradio as gr
+
+from helper.text.text_howto import TextHowTo
+
+with gr.Blocks() as help_tab:
+    gr.Markdown("lorem ipsum...")
+    with gr.Tabs():
+        with gr.Tab("Discussion & FAQ"):
+            pass
+
+        with gr.Tab("Fast track"):
+            pass
+        with gr.Tab("Stepwise"):
+            with gr.Row():
+                with gr.Accordion("Info", open=False) as example_accord:
+                    with gr.Row(equal_height=False):
+                        gr.Markdown(TextHowTo.stepwise_htr_tool)
+                    with gr.Row():
+                        gr.Markdown(TextHowTo.stepwise_htr_tool_tab_intro)
+                    with gr.Row():
+                        with gr.Tabs():
+                            with gr.Tab("1. Region Segmentation"):
+                                gr.Markdown(TextHowTo.stepwise_htr_tool_tab1)
+                            with gr.Tab("2. Line Segmentation"):
+                                gr.Markdown(TextHowTo.stepwise_htr_tool_tab2)
+                            with gr.Tab("3. Transcribe Text"):
+                                gr.Markdown(TextHowTo.stepwise_htr_tool_tab3)
+                            with gr.Tab("4. Explore Results"):
+                                gr.Markdown(TextHowTo.stepwise_htr_tool_tab4)
+                    gr.Markdown(TextHowTo.stepwise_htr_tool_end)
+
+        with gr.Tab("API"):
+            pass
+        with gr.Tab("Duplicating for own use"):
+            pass

tabs/htr_tool.py

@@ -4,9 +4,7 @@ from helper.examples.examples import DemoImages
 from src.htr_pipeline.gradio_backend import FastTrack, SingletonModelLoader
 
 model_loader = SingletonModelLoader()
-
 fast_track = FastTrack(model_loader)
-
 images_for_demo = DemoImages()
 
 terminate = False
@@ -21,7 +19,7 @@ with gr.Blocks() as htr_tool_tab:
                 )
 
             with gr.Row():
-                with gr.Tab("Output and Settings") as tab_output_and_setting_selector:
+                with gr.Tab("Run & Settings") as tab_output_and_setting_selector:
                     with gr.Row():
                         stop_htr_button = gr.Button(
                             value="Stop HTR",
@@ -41,10 +39,10 @@ with gr.Blocks() as htr_tool_tab:
                         label="Download output file", visible=True, scale=1, height=100, elem_id="download_file"
                     )
 
-                with gr.Tab("Image Viewer") as tab_image_viewer_selector:
+                with gr.Tab("Visualize results") as tab_image_viewer_selector:
                     with gr.Row():
                         gr.Button(
-                            value="External Image Viewer",
+                            value="Image Viewer (demo)",
                             variant="secondary",
                             link="https://huggingface.co/spaces/Riksarkivet/Viewer_demo",
                             interactive=True,
@@ -58,17 +56,34 @@ with gr.Blocks() as htr_tool_tab:
                         interactive=False, label="Text Selector", info="Select a mask on Image Viewer to return text"
                     )
 
+                with gr.Tab("(WIP) Compare runs") as tab_model_compare_selector:
+                    with gr.Box():
+                        gr.Markdown(
+                            "Compare different runs with uploaded Ground Truth and calculate CER. You will also be able to upload output format files"
+                        )
+
+                        calc_cer_button_fast = gr.Button("Calculate CER", variant="primary", visible=True)
+
         with gr.Column(scale=4):
             with gr.Box():
                 with gr.Row(visible=True) as output_and_setting_tab:
+                    with gr.Column(scale=2):
+                        fast_name_files_placeholder = gr.Markdown(visible=False)
+                        gr.Examples(
+                            examples=images_for_demo.examples_list,
+                            inputs=[fast_name_files_placeholder, fast_track_input_region_image],
+                            label="Example images",
+                            examples_per_page=5,
+                        )
+
                     with gr.Column(scale=3):
                         with gr.Row():
                             with gr.Group():
                                 gr.Markdown("    ⚙️ Settings ")
                                 with gr.Row():
                                     radio_file_input = gr.CheckboxGroup(
-                                        choices=["Txt", "XML"],
-                                        value=["Txt", "XML"],
+                                        choices=["Txt", "Page XML"],
+                                        value=["Txt", "Page XML"],
                                         label="Output file extension",
                                         # info="Only txt and page xml is supported for now!",
                                         scale=1,
@@ -84,54 +99,59 @@ with gr.Blocks() as htr_tool_tab:
                                         label="Output prediction threshold",
                                         info="Output XML with prediction score",
                                     )
-                                with gr.Row():
-                                    gr.Slider(
-                                        value=0.6,
-                                        minimum=0.5,
-                                        maximum=1,
-                                        label="HTR threshold",
-                                        info="Prediction score threshold for transcribed lines",
-                                        scale=1,
-                                    )
-                                    gr.Slider(
-                                        value=0.7,
-                                        minimum=0.6,
-                                        maximum=1,
-                                        label="Avg threshold",
-                                        info="Average prediction score for a region",
-                                        scale=1,
-                                    )
-
-                                htr_tool_region_segment_model_dropdown = gr.Dropdown(
-                                    choices=["Riksarkivet/rtmdet_region"],
-                                    value="Riksarkivet/rtmdet_region",
-                                    label="Region Segment models",
-                                    info="Will add more models later!",
-                                )
-
-                                # with gr.Accordion("Transcribe settings:", open=False):
-                                htr_tool_line_segment_model_dropdown = gr.Dropdown(
-                                    choices=["Riksarkivet/rtmdet_lines"],
-                                    value="Riksarkivet/rtmdet_lines",
-                                    label="Line Segment models",
-                                    info="Will add more models later!",
-                                )
-
-                                htr_tool_transcriber_model_dropdown = gr.Dropdown(
-                                    choices=["Riksarkivet/satrn_htr", "microsoft/trocr-base-handwritten"],
-                                    value="Riksarkivet/satrn_htr",
-                                    label="Transcriber models",
-                                    info="Models will be continuously  updated with future additions for specific cases.",
-                                )
 
-                    with gr.Column(scale=2):
-                        fast_name_files_placeholder = gr.Markdown(visible=False)
-                        gr.Examples(
-                            examples=images_for_demo.examples_list,
-                            inputs=[fast_name_files_placeholder, fast_track_input_region_image],
-                            label="Example images",
-                            examples_per_page=5,
-                        )
+                                with gr.Accordion("Models", open=False):
+                                    with gr.Group():
+                                        with gr.Row():
+                                            htr_tool_region_segment_model_dropdown = gr.Dropdown(
+                                                choices=["Riksarkivet/rtmdet_region"],
+                                                value="Riksarkivet/rtmdet_region",
+                                                label="Region segmentation models",
+                                                info="More models will be added",
+                                            )
+
+                                            gr.Slider(
+                                                minimum=0.4,
+                                                maximum=1,
+                                                value=0.5,
+                                                step=0.05,
+                                                label="P-threshold",
+                                                info="""Filter confidence score for a prediction score to be considered""",
+                                            )
+
+                                        with gr.Row():
+                                            htr_tool_line_segment_model_dropdown = gr.Dropdown(
+                                                choices=["Riksarkivet/rtmdet_lines"],
+                                                value="Riksarkivet/rtmdet_lines",
+                                                label="Line segmentation models",
+                                                info="More models will be added",
+                                            )
+
+                                            gr.Slider(
+                                                minimum=0.4,
+                                                maximum=1,
+                                                value=0.5,
+                                                step=0.05,
+                                                label="P-threshold",
+                                                info="""Filter confidence score for a prediction score to be considered""",
+                                            )
+
+                                        with gr.Row():
+                                            htr_tool_transcriber_model_dropdown = gr.Dropdown(
+                                                choices=["Riksarkivet/satrn_htr", "microsoft/trocr-base-handwritten"],
+                                                value="Riksarkivet/satrn_htr",
+                                                label="Text recognition models",
+                                                info="More models will be added",
+                                            )
+
+                                            gr.Slider(
+                                                value=0.6,
+                                                minimum=0.5,
+                                                maximum=1,
+                                                label="HTR threshold",
+                                                info="Prediction score threshold for transcribed lines",
+                                                scale=1,
+                                            )
 
                 with gr.Row(visible=False) as image_viewer_tab:
                     text_polygon_dict = gr.Variable()
@@ -140,6 +160,43 @@ with gr.Blocks() as htr_tool_tab:
                         label="Image Viewer", type="numpy", height=600, interactive=False
                     )
 
+                with gr.Column(visible=False) as model_compare_selector:
+                    with gr.Row():
+                        gr.Radio(
+                            choices=["Compare Page XML", "Compare different runs"],
+                            value="Compare Page XML",
+                            info="Compare different runs from HTRFLOW or with external runs, e.g with Transkibus ",
+                        )
+                    with gr.Row():
+                        gr.UploadButton(label="Run A")
+
+                        gr.UploadButton(label="Run B")
+
+                        gr.UploadButton(label="Ground Truth")
+
+                    with gr.Row():
+                        gr.HighlightedText(
+                            label="Text diff runs",
+                            combine_adjacent=True,
+                            show_legend=True,
+                            color_map={"+": "red", "-": "green"},
+                        )
+
+                    with gr.Row():
+                        gr.HighlightedText(
+                            label="Text diff ground truth",
+                            combine_adjacent=True,
+                            show_legend=True,
+                            color_map={"+": "red", "-": "green"},
+                        )
+
+                    with gr.Row():
+                        with gr.Column(scale=1):
+                            with gr.Row(equal_height=False):
+                                cer_output_fast = gr.Textbox(label="CER:")
+                        with gr.Column(scale=2):
+                            pass
+
     xml_rendered_placeholder_for_api = gr.Textbox(visible=False)
 
     htr_event_click_event = htr_pipeline_button.click(
@@ -165,17 +222,25 @@ with gr.Blocks() as htr_tool_tab:
     )
 
     def update_selected_tab_output_and_setting():
-        return gr.update(visible=True), gr.update(visible=False)
+        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
 
     def update_selected_tab_image_viewer():
-        return gr.update(visible=False), gr.update(visible=True)
+        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
+
+    def update_selected_tab_model_compare():
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
 
     tab_output_and_setting_selector.select(
-        fn=update_selected_tab_output_and_setting, outputs=[output_and_setting_tab, image_viewer_tab]
+        fn=update_selected_tab_output_and_setting,
+        outputs=[output_and_setting_tab, image_viewer_tab, model_compare_selector],
     )
 
     tab_image_viewer_selector.select(
-        fn=update_selected_tab_image_viewer, outputs=[output_and_setting_tab, image_viewer_tab]
+        fn=update_selected_tab_image_viewer, outputs=[output_and_setting_tab, image_viewer_tab, model_compare_selector]
+    )
+
+    tab_model_compare_selector.select(
+        fn=update_selected_tab_model_compare, outputs=[output_and_setting_tab, image_viewer_tab, model_compare_selector]
     )
 
     def stop_function():
@@ -195,3 +260,5 @@ with gr.Blocks() as htr_tool_tab:
     fast_track_output_image.select(
         fast_track.get_text_from_coords, inputs=text_polygon_dict, outputs=selection_text_from_image_viewer
     )
+
+    htr_pipeline_button.click(fn=handler.store_metric_data, inputs="htr_pipeline_button")

tabs/stepwise_htr_tool.py

@@ -1,11 +1,11 @@
 import os
 import shutil
+from difflib import Differ
 
 import evaluate
 import gradio as gr
 
 from helper.examples.examples import DemoImages
-from helper.text.text_howto import TextHowTo
 from src.htr_pipeline.gradio_backend import CustomTrack, SingletonModelLoader
 
 model_loader = SingletonModelLoader()
@@ -19,57 +19,28 @@ cer_metric = evaluate.load("cer")
 
 with gr.Blocks() as stepwise_htr_tool_tab:
     with gr.Tabs():
-        with gr.Tab("1. Region Segmentation"):
+        with gr.Tab("1. Region segmentation"):
             with gr.Row():
-                with gr.Accordion("Info", open=False) as example_accord:
-                    with gr.Row(equal_height=False):
-                        gr.Markdown(TextHowTo.stepwise_htr_tool)
-                    with gr.Row():
-                        gr.Markdown(TextHowTo.stepwise_htr_tool_tab_intro)
-                    with gr.Row():
-                        with gr.Tabs():
-                            with gr.Tab("1. Region Segmentation"):
-                                gr.Markdown(TextHowTo.stepwise_htr_tool_tab1)
-                            with gr.Tab("2. Line Segmentation"):
-                                gr.Markdown(TextHowTo.stepwise_htr_tool_tab2)
-                            with gr.Tab("3. Transcribe Text"):
-                                gr.Markdown(TextHowTo.stepwise_htr_tool_tab3)
-                            with gr.Tab("4. Explore Results"):
-                                gr.Markdown(TextHowTo.stepwise_htr_tool_tab4)
-                    gr.Markdown(TextHowTo.stepwise_htr_tool_end)
-            with gr.Row():
-                with gr.Column(scale=2):
+                with gr.Column(scale=1):
                     vis_data_folder_placeholder = gr.Markdown(visible=False)
                     name_files_placeholder = gr.Markdown(visible=False)
 
-                    with gr.Row():
+                    with gr.Group():
                         input_region_image = gr.Image(
-                            label="Image to Region segment",
+                            label="Image to region segment",
                             # type="numpy",
                             tool="editor",
-                            height=400,
+                            height=500,
                         )
-
-                    with gr.Row():
-                        clear_button = gr.Button("Clear", variant="secondary", elem_id="clear_button")
-
-                        region_segment_button = gr.Button(
-                            "Segment Region",
-                            variant="primary",
-                            elem_id="region_segment_button",
-                        )
-
-                    with gr.Group():
-                        with gr.Accordion("Region segment settings:", open=False):
-                            with gr.Row():
+                        with gr.Accordion("Settings", open=False):
+                            with gr.Group():
                                 reg_pred_score_threshold_slider = gr.Slider(
                                     minimum=0.4,
                                     maximum=1,
                                     value=0.5,
                                     step=0.05,
                                     label="P-threshold",
-                                    info="""Filter and determine the confidence score 
-                                                    required for a prediction score to be considered""",
+                                    info="""Filter the confidence score for a prediction score to be considered""",
                                 )
                                 reg_containments_threshold_slider = gr.Slider(
                                     minimum=0,
@@ -81,27 +52,37 @@ with gr.Blocks() as stepwise_htr_tool_tab:
                                                     for a detected region or object to be considered valid""",
                                 )
 
-                            with gr.Row():
                                 region_segment_model_dropdown = gr.Dropdown(
-                                    choices=["Riksarkivet/RmtDet_region"],
-                                    value="Riksarkivet/RmtDet_region",
-                                    label="Region segment model",
-                                    info="Will add more models later!",
+                                    choices=["Riksarkivet/rtm_region"],
+                                    value="Riksarkivet/rtm_region",
+                                    label="Region segmentation model",
+                                    info="More models will be added",
                                 )
 
-                        with gr.Accordion("Example images to use:", open=False) as example_accord:
-                            gr.Examples(
-                                examples=images_for_demo.examples_list,
-                                inputs=[name_files_placeholder, input_region_image],
-                                label="Example images",
-                                examples_per_page=5,
-                            )
+                    with gr.Row():
+                        clear_button = gr.Button("Clear", variant="secondary", elem_id="clear_button")
 
-                with gr.Column(scale=3):
-                    output_region_image = gr.Image(label="Segmented regions", type="numpy", height=550)
+                        region_segment_button = gr.Button(
+                            "Run",
+                            variant="primary",
+                            elem_id="region_segment_button",
+                        )
+
+                with gr.Column(scale=2):
+                    with gr.Box():
+                        with gr.Row():
+                            with gr.Column(scale=2):
+                                gr.Examples(
+                                    examples=images_for_demo.examples_list,
+                                    inputs=[name_files_placeholder, input_region_image],
+                                    label="Example images",
+                                    examples_per_page=5,
+                                )
+                            with gr.Column(scale=3):
+                                output_region_image = gr.Image(label="Segmented regions", type="numpy", height=600)
 
         ##############################################
-        with gr.Tab("2. Line Segmentation"):
+        with gr.Tab("2. Line segmentation"):
             image_placeholder_lines = gr.Image(
                 label="Segmented lines",
                 # type="numpy",
@@ -112,50 +93,50 @@ with gr.Blocks() as stepwise_htr_tool_tab:
 
             with gr.Row(visible=False) as control_line_segment:
                 with gr.Column(scale=2):
-                    with gr.Box():
-                        regions_cropped_gallery = gr.Gallery(
-                            label="Segmented regions",
-                            elem_id="gallery",
-                            columns=[2],
-                            rows=[2],
-                            # object_fit="contain",
-                            height=450,
-                            preview=True,
-                            container=False,
-                        )
+                    with gr.Group():
+                        with gr.Box():
+                            regions_cropped_gallery = gr.Gallery(
+                                label="Segmented regions",
+                                elem_id="gallery",
+                                columns=[2],
+                                rows=[2],
+                                # object_fit="contain",
+                                height=450,
+                                preview=True,
+                                container=False,
+                            )
 
-                    input_region_from_gallery = gr.Image(
-                        label="Region segmentation to line segment", interactive="False", visible=False, height=400
-                    )
+                        input_region_from_gallery = gr.Image(
+                            label="Region segmentation to line segment", interactive="False", visible=False, height=400
+                        )
 
-                    with gr.Row():
-                        with gr.Accordion("Line segment settings:", open=False):
-                            with gr.Row():
-                                line_pred_score_threshold_slider = gr.Slider(
-                                    minimum=0.3,
-                                    maximum=1,
-                                    value=0.4,
-                                    step=0.05,
-                                    label="Pred_score threshold",
-                                    info="""Filter and determine the confidence score 
-                                                    required for a prediction score to be considered""",
-                                )
-                                line_containments_threshold_slider = gr.Slider(
-                                    minimum=0,
-                                    maximum=1,
-                                    value=0.5,
-                                    step=0.05,
-                                    label="Containments threshold",
-                                    info="""The minimum required overlap or similarity 
-                                                    for a detected region or object to be considered valid""",
-                                )
-                            with gr.Row(equal_height=False):
-                                line_segment_model_dropdown = gr.Dropdown(
-                                    choices=["Riksarkivet/RmtDet_lines"],
-                                    value="Riksarkivet/RmtDet_lines",
-                                    label="Line segment model",
-                                    info="Will add more models later!",
-                                )
+                        with gr.Row():
+                            with gr.Accordion("Settings", open=False):
+                                with gr.Row():
+                                    line_pred_score_threshold_slider = gr.Slider(
+                                        minimum=0.3,
+                                        maximum=1,
+                                        value=0.4,
+                                        step=0.05,
+                                        label="Pred_score threshold",
+                                        info="""Filter the confidence score for a prediction score to be considered""",
+                                    )
+                                    line_containments_threshold_slider = gr.Slider(
+                                        minimum=0,
+                                        maximum=1,
+                                        value=0.5,
+                                        step=0.05,
+                                        label="Containments threshold",
+                                        info="""The minimum required overlap or similarity 
+                                                        for a detected region or object to be considered valid""",
+                                    )
+                                with gr.Row(equal_height=False):
+                                    line_segment_model_dropdown = gr.Dropdown(
+                                        choices=["Riksarkivet/rtmdet_lines"],
+                                        value="Riksarkivet/rtmdet_lines",
+                                        label="Line segment model",
+                                        info="More models will be added",
+                                    )
                     with gr.Row():
                         clear_line_segment_button = gr.Button(
                             " ",
@@ -165,7 +146,7 @@ with gr.Blocks() as stepwise_htr_tool_tab:
                         )
 
                         line_segment_button = gr.Button(
-                            "Segment Lines",
+                            "Run",
                             variant="primary",
                             # elem_id="center_button",
                             scale=1,
@@ -179,7 +160,7 @@ with gr.Blocks() as stepwise_htr_tool_tab:
                     )
 
         ###############################################
-        with gr.Tab("3. Transcribe Text"):
+        with gr.Tab("3. Text recognition"):
             image_placeholder_htr = gr.Image(
                 label="Transcribed lines",
                 # type="numpy",
@@ -192,33 +173,45 @@ with gr.Blocks() as stepwise_htr_tool_tab:
                 inputs_lines_to_transcribe = gr.Variable()
 
                 with gr.Column(scale=2):
-                    image_inputs_lines_to_transcribe = gr.Image(
-                        label="Transcribed lines", type="numpy", interactive="False", visible=False, height=470
-                    )
-                    with gr.Row():
-                        with gr.Accordion("Transcribe settings:", open=False):
-                            transcriber_model = gr.Dropdown(
-                                choices=["Riksarkivet/SATRN_transcriber", "microsoft/trocr-base-handwritten"],
-                                value="Riksarkivet/SATRN_transcriber",
-                                label="Transcriber model",
-                                info="Will add more models later!",
-                            )
+                    with gr.Group():
+                        image_inputs_lines_to_transcribe = gr.Image(
+                            label="Transcribed lines", type="numpy", interactive="False", visible=False, height=470
+                        )
+                        with gr.Row():
+                            with gr.Accordion("Settings", open=False):
+                                transcriber_model = gr.Dropdown(
+                                    choices=["Riksarkivet/satrn_htr", "microsoft/trocr-base-handwritten"],
+                                    value="Riksarkivet/satrn_htr",
+                                    label="Text recognition model",
+                                    info="More models will be added",
+                                )
+
+                                gr.Slider(
+                                    value=0.6,
+                                    minimum=0.5,
+                                    maximum=1,
+                                    label="HTR threshold",
+                                    info="Prediction score threshold for transcribed lines",
+                                    scale=1,
+                                )
+
                     with gr.Row():
                         clear_transcribe_button = gr.Button(" ", variant="Secondary", visible=True, scale=1)
 
-                        transcribe_button = gr.Button("Transcribe Lines", variant="primary", visible=True, scale=1)
+                        transcribe_button = gr.Button("Run", variant="primary", visible=True, scale=1)
 
                 with gr.Column(scale=3):
                     with gr.Row():
                         transcribed_text = gr.Textbox(
                             label="Transcribed text",
-                            info="Transcribed text is being streamed back from the HTR-model",
-                            lines=25,
+                            info="Transcribed text is being streamed back from the Text recognition model",
+                            lines=26,
                             value="",
+                            show_copy_button=True,
                         )
 
         #####################################
-        with gr.Tab("4. Explore Results"):
+        with gr.Tab("4. Explore results"):
             image_placeholder_explore_results = gr.Image(
                 label="Cropped transcribed lines",
                 # type="numpy",
@@ -229,40 +222,48 @@ with gr.Blocks() as stepwise_htr_tool_tab:
 
             with gr.Row(visible=False, equal_height=False) as control_results_transcribe:
                 with gr.Column(scale=1, visible=True):
-                    with gr.Box():
-                        temp_gallery_input = gr.Variable()
-
-                        gallery_inputs_lines_to_transcribe = gr.Gallery(
-                            label="Cropped transcribed lines",
-                            elem_id="gallery_lines",
-                            columns=[3],
-                            rows=[3],
-                            # object_fit="contain",
-                            height=250,
-                            preview=True,
-                            container=False,
+                    with gr.Group():
+                        with gr.Box():
+                            temp_gallery_input = gr.Variable()
+
+                            gallery_inputs_lines_to_transcribe = gr.Gallery(
+                                label="Cropped transcribed lines",
+                                elem_id="gallery_lines",
+                                columns=[3],
+                                rows=[3],
+                                # object_fit="contain",
+                                height=150,
+                                preview=True,
+                                container=False,
+                            )
+                        with gr.Row():
+                            dataframe_text_index = gr.Textbox(
+                                label="Text from DataFrame selection",
+                                placeholder="Select row from the DataFrame.",
+                                interactive=False,
+                            )
+                        with gr.Row():
+                            gt_text_index = gr.Textbox(
+                                label="Ground Truth",
+                                placeholder="Provide the ground truth, if available.",
+                                interactive=True,
+                            )
+                    with gr.Row():
+                        diff_token_output = gr.HighlightedText(
+                            label="Text diff",
+                            combine_adjacent=True,
+                            show_legend=True,
+                            color_map={"+": "red", "-": "green"},
                         )
 
-                    dataframe_text_index = gr.Textbox(
-                        label="Text from DataFrame selection",
-                        placeholder="Select row from the DataFrame.",
-                        interactive=False,
-                    )
-
-                    gt_text_index = gr.Textbox(
-                        label="Ground Truth",
-                        placeholder="Provide the ground truth, if available.",
-                        interactive=True,
-                    )
                     with gr.Row(equal_height=False):
-                        calc_cer_button = gr.Button("Calculate CER", variant="primary", visible=True)
-
                         cer_output = gr.Textbox(label="CER:")
+                        calc_cer_button = gr.Button("Calculate CER", variant="primary", visible=True)
 
                 with gr.Column(scale=1, visible=True):
                     mapping_dict = gr.Variable()
                     transcribed_text_df_finish = gr.Dataframe(
-                        headers=["Transcribed text", "pred score"],
+                        headers=["Transcribed text", "Pred score"],
                         max_rows=14,
                         col_count=(2, "fixed"),
                         wrap=True,
@@ -272,6 +273,21 @@ with gr.Blocks() as stepwise_htr_tool_tab:
                     )
 
     # custom track
+
+    def diff_texts(text1, text2):
+        d = Differ()
+        return [(token[2:], token[0] if token[0] != " " else None) for token in d.compare(text1, text2)]
+
+    def compute_cer(dataframe_text_index, gt_text_index):
+        if gt_text_index is not None and gt_text_index.strip() != "":
+            return cer_metric.compute(predictions=[dataframe_text_index], references=[gt_text_index])
+        else:
+            return "Ground truth not provided"
+
+    calc_cer_button.click(compute_cer, inputs=[dataframe_text_index, gt_text_index], outputs=cer_output)
+
+    calc_cer_button.click(diff_texts, inputs=[dataframe_text_index, gt_text_index], outputs=[diff_token_output])
+
     region_segment_button.click(
         custom_track.region_segment,
         inputs=[input_region_image, reg_pred_score_threshold_slider, reg_containments_threshold_slider],
@@ -318,14 +334,6 @@ with gr.Blocks() as stepwise_htr_tool_tab:
         ],
     )
 
-    def compute_cer(dataframe_text_index, gt_text_index):
-        if gt_text_index is not None and gt_text_index.strip() != "":
-            return cer_metric.compute(predictions=[dataframe_text_index], references=[gt_text_index])
-        else:
-            return "Ground truth not provided"
-
-    calc_cer_button.click(compute_cer, inputs=[dataframe_text_index, gt_text_index], outputs=cer_output)
-
     clear_button.click(
         lambda: (
             (shutil.rmtree("./vis_data") if os.path.exists("./vis_data") else None, None)[1],