refactored redid the ui of fast track

app.py

@@ -118,7 +118,7 @@ print(job.result())
     # demo.load(None, None, None, _js=js)
 
 
-demo.queue(concurrency_count=1, max_size=1)
+demo.queue(concurrency_count=2, max_size=2)
 
 
 if __name__ == "__main__":

helper/gradio_config.py

@@ -22,6 +22,7 @@ class GradioConfig:
         #gallery {height: 400px} 
         .fixed-height.svelte-g4rw9.svelte-g4rw9 {min-height: 400px;}
 
+        #download_file > div.empty.svelte-lk9eg8.large.unpadded_box  {min-height: 100px;}
         #gallery_lines > div.preview.svelte-1b19cri > div.thumbnails.scroll-hide.svelte-1b19cri {display: none;}
 
                 """

src/htr_pipeline/gradio_backend.py

@@ -1,6 +1,8 @@
 import os
 
+import cv2
 import gradio as gr
+import numpy as np
 import pandas as pd
 
 from src.htr_pipeline.inferencer import Inferencer, InferencerInterface
@@ -23,13 +25,21 @@ class SingletonModelLoader:
             self.pipeline = Pipeline(self.inferencer)
 
 
+def handling_callback_stop_inferencer():
+    from src.htr_pipeline.utils import pipeline_inferencer
+
+    pipeline_inferencer.terminate = False
+
+
 # fast track
 class FastTrack:
     def __init__(self, model_loader):
         self.pipeline: PipelineInterface = model_loader.pipeline
 
     def segment_to_xml(self, image, radio_button_choices):
-        gr.Info("Running HTR-pipeline")
+        handling_callback_stop_inferencer()
+
+        gr.Info("Excuting HTR on image")
         xml_xml = "page_xml.xml"
         xml_txt = "page_txt.txt"
 
@@ -41,10 +51,18 @@ class FastTrack:
         with open(xml_xml, "w") as f:
             f.write(rendered_xml)
 
-        xml_img = self.visualize_xml_and_return_txt(image, xml_txt)
+        if os.path.exists(f"./{xml_txt}"):
+            os.remove(f"./{xml_txt}")
+
+        self.pipeline.parse_xml_to_txt()
+
         returned_file_extension = self.file_extenstion_to_return(radio_button_choices, xml_xml, xml_txt)
 
-        return xml_img, returned_file_extension, gr.update(visible=True)
+        return returned_file_extension, gr.update(visible=True)
+
+    def visualize_image_viewer(self, image):
+        xml_img, text_polygon_dict = self.pipeline.visualize_xml(image)
+        return xml_img, text_polygon_dict
 
     def file_extenstion_to_return(self, radio_button_choices, xml_xml, xml_txt):
         if len(radio_button_choices) < 2:
@@ -56,20 +74,19 @@ class FastTrack:
             returned_file_extension = [xml_txt, xml_xml]
         return returned_file_extension
 
+    def get_text_from_coords(self, text_polygon_dict, evt: gr.SelectData):
+        x, y = evt.index[0], evt.index[1]
+
+        for text, polygon_coords in text_polygon_dict.items():
+            if (
+                cv2.pointPolygonTest(np.array(polygon_coords), (x, y), False) >= 0
+            ):  # >= 0 means on the polygon or inside
+                return text
+
     def segment_to_xml_api(self, image):
         rendered_xml = self.pipeline.running_htr_pipeline(image)
         return rendered_xml
 
-    def visualize_xml_and_return_txt(self, img, xml_txt):
-        xml_img = self.pipeline.visualize_xml(img)
-
-        if os.path.exists(f"./{xml_txt}"):
-            os.remove(f"./{xml_txt}")
-
-        self.pipeline.parse_xml_to_txt()
-
-        return xml_img
-
 
 # Custom track
 class CustomTrack:

src/htr_pipeline/pipeline.py

@@ -40,8 +40,8 @@ class Pipeline:
     def visualize_xml(self, input_image: np.ndarray) -> np.ndarray:
         xml_viz = XmlViz()
         bin_input_image = self.preprocess_img.binarize_img(input_image)
-        xml_image = xml_viz.visualize_xml(bin_input_image)
-        return xml_image
+        xml_image, text_polygon_dict = xml_viz.visualize_xml(bin_input_image)
+        return xml_image, text_polygon_dict
 
     @timer_func
     def parse_xml_to_txt(self) -> None:

src/htr_pipeline/utils/pipeline_inferencer.py

@@ -4,6 +4,8 @@ from tqdm import tqdm
 from src.htr_pipeline.utils.process_segmask import SegMaskHelper
 from src.htr_pipeline.utils.xml_helper import XMLHelper
 
+terminate = False
+
 
 class PipelineInferencer:
     def __init__(self, process_seg_mask: SegMaskHelper, xml_helper: XMLHelper):
@@ -29,6 +31,8 @@ class PipelineInferencer:
         containments_threshold,
         htr_threshold=0.7,
     ):
+        global terminate
+
         _, regions_cropped_ordered, reg_polygons_ordered, reg_masks_ordered = inferencer.predict_regions(
             image,
             pred_score_threshold=pred_score_threshold_regions,
@@ -38,6 +42,8 @@ class PipelineInferencer:
         gr.Info(f"Found {len(regions_cropped_ordered)} Regions to parse")
         region_data_list = []
         for i, data in tqdm(enumerate(zip(regions_cropped_ordered, reg_polygons_ordered, reg_masks_ordered))):
+            if terminate:
+                break
             region_data = self._create_region_data(
                 data, i, inferencer, pred_score_threshold_lines, containments_threshold, htr_threshold
             )
@@ -68,7 +74,7 @@ class PipelineInferencer:
         region_data["textLines"] = text_lines
         mean_htr_score = sum(htr_scores) / len(htr_scores) if htr_scores else 0
 
-        return region_data if mean_htr_score > htr_threshold else None
+        return region_data if mean_htr_score > htr_threshold + 0.1 else None
 
     def _process_lines(
         self, text_region, inferencer, pred_score_threshold, containments_threshold, mask, region_id, htr_threshold=0.7
@@ -90,7 +96,11 @@ class PipelineInferencer:
 
         gr.Info(f" Region {id_number}, found {total_lines_len} lines to parse and transcribe.")
 
+        global terminate
+
         for index, (line, line_pol) in enumerate(zip(lines_cropped_ordered, line_polygons_ordered_trans)):
+            if terminate:
+                break
             line_data, htr_score = self._create_line_data(line, line_pol, index, region_id, inferencer, htr_threshold)
 
             if line_data:

src/htr_pipeline/utils/visualize_xml.py

@@ -1,8 +1,8 @@
 import random
 import xml.etree.ElementTree as ET
 
-import gradio as gr
-from PIL import Image, ImageDraw, ImageFont
+import cv2
+import numpy as np
 
 
 class XmlViz:
@@ -11,58 +11,24 @@ class XmlViz:
         self.root = self.tree.getroot()
         self.namespace = "{http://schema.primaresearch.org/PAGE/gts/pagecontent/2013-07-15}"
 
-    def visualize_xml(
-        self,
-        background_image,
-        font_size=9,
-        text_offset=10,
-        font_path_tff="./src/htr_pipeline/utils/templates/arial.ttf",
-    ):
-        image = Image.fromarray(background_image).convert("RGBA")
-
-        text_offset = -text_offset
-        base_font_size = font_size
-        font_path = font_path_tff
-
-        max_bbox_width = 0  # Initialize maximum bounding box width
-        gr.Info("Parsing XML to visualize the data.")
-        for textregion in self.root.findall(f".//{self.namespace}TextRegion"):
-            coords = textregion.find(f"{self.namespace}Coords").attrib["points"].split()
-            points = [tuple(map(int, point.split(","))) for point in coords]
-            x_coords, y_coords = zip(*points)
-            min_x, max_x = min(x_coords), max(x_coords)
-            bbox_width = max_x - min_x  # Width of the current bounding box
-            max_bbox_width = max(max_bbox_width, bbox_width)  # Update maximum bounding box width
-
-        scaling_factor = max_bbox_width / 400.0  # Use maximum bounding box width for scaling
-        font_size_scaled = int(base_font_size * scaling_factor)
-        font = ImageFont.truetype(font_path, font_size_scaled)
+    def visualize_xml(self, background_image):
+        overlay = background_image.copy()
+        text_polygon_dict = {}
 
         for textregion in self.root.findall(f".//{self.namespace}TextRegion"):
-            fill_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 100)
+            fill_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
             for textline in textregion.findall(f".//{self.namespace}TextLine"):
                 coords = textline.find(f"{self.namespace}Coords").attrib["points"].split()
                 points = [tuple(map(int, point.split(","))) for point in coords]
-
-                poly_image = Image.new("RGBA", image.size)
-                poly_draw = ImageDraw.Draw(poly_image)
-                poly_draw.polygon(points, fill=fill_color)
+                cv2.fillPoly(overlay, [np.array(points)], fill_color)
 
                 text = textline.find(f"{self.namespace}TextEquiv").find(f"{self.namespace}Unicode").text
+                text_polygon_dict[text] = points
 
-                x_coords, y_coords = zip(*points)
-                min_x, max_x = min(x_coords), max(x_coords)
-                min_y = min(y_coords)
-                text_width, text_height = poly_draw.textsize(text, font=font)  # Get text size
-                text_position = (
-                    (min_x + max_x) // 2 - text_width // 2,
-                    min_y + text_offset,
-                )  # Center text horizontally
-
-                poly_draw.text(text_position, text, fill=(0, 0, 0), font=font)
-                image = Image.alpha_composite(image, poly_image)
+        # Blend the overlay with the original image
+        cv2.addWeighted(overlay, 0.5, background_image, 0.5, 0, background_image)
 
-        return image
+        return background_image, text_polygon_dict
 
 
 if __name__ == "__main__":

tabs/htr_tool.py

@@ -9,6 +9,8 @@ fast_track = FastTrack(model_loader)
 
 images_for_demo = DemoImages()
 
+terminate = False
+
 
 with gr.Blocks() as htr_tool_tab:
     with gr.Row(equal_height=True):
@@ -19,54 +21,131 @@ with gr.Blocks() as htr_tool_tab:
                 )
 
             with gr.Row():
-                radio_file_input = gr.CheckboxGroup(
-                    choices=["Txt", "XML"],
-                    value=["XML"],
-                    label="Output file extension",
-                    # info="Only txt and page xml is supported for now!",
-                    scale=1,
-                )
+                with gr.Tab("Output and Settings") as tab_output_and_setting_selector:
+                    with gr.Row():
+                        stop_htr_button = gr.Button(
+                            value="Stop HTR",
+                            variant="stop",
+                        )
 
-                htr_pipeline_button = gr.Button(
-                    "Run HTR", variant="primary", visible=True, elem_id="run_pipeline_button", scale=1
-                )
-            with gr.Group():
-                with gr.Row():
-                    fast_file_downlod = gr.File(label="Download output file", visible=False)
-                with gr.Row():
-                    with gr.Accordion("Example images to use:", open=False) as fast_example_accord:
-                        fast_name_files_placeholder = gr.Markdown(visible=False)
+                        htr_pipeline_button = gr.Button(
+                            "Run HTR",
+                            variant="primary",
+                            visible=True,
+                            elem_id="run_pipeline_button",
+                        )
+
+                    htr_pipeline_button_api = gr.Button("Run pipeline", variant="primary", visible=False, scale=1)
+
+                    fast_file_downlod = gr.File(
+                        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.Row():
+                        gr.Button(
+                            value="External Image Viewer",
+                            variant="secondary",
+                            link="https://huggingface.co/spaces/Riksarkivet/Viewer_demo",
+                            interactive=True,
+                        )
+
+                        run_image_visualizer_button = gr.Button(
+                            value="Visualize results", variant="primary", interactive=True
+                        )
+
+                    selection_text_from_image_viewer = gr.Textbox(
+                        interactive=False, label="Text Selector", info="Select a mask on Image Viewer to return text"
+                    )
+
+        with gr.Column(scale=4):
+            with gr.Box():
+                with gr.Row(visible=True) as output_and_setting_tab:
+                    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=["XML"],
+                                        label="Output file extension",
+                                        # info="Only txt and page xml is supported for now!",
+                                        scale=1,
+                                    )
+                                with gr.Row():
+                                    gr.Checkbox(
+                                        value=True,
+                                        label="Binarize image",
+                                        info="Binarize image to reduce background noise",
+                                    )
+                                    gr.Checkbox(
+                                        value=True,
+                                        label="Output prediction threshold",
+                                        info="Output XML with prediction score",
+                                    )
+                                with gr.Row():
+                                    gr.Slider(
+                                        value=0.7,
+                                        minimum=0.5,
+                                        maximum=1,
+                                        label="HTR threshold",
+                                        info="Prediction score threshold for transcribed lines",
+                                        scale=1,
+                                    )
+                                    gr.Slider(
+                                        value=0.8,
+                                        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/RmtDet_region"],
+                                    value="Riksarkivet/RmtDet_region",
+                                    label="Region segment model",
+                                    info="Will add more models later!",
+                                )
+
+                                # with gr.Accordion("Transcribe settings:", open=False):
+                                htr_tool_line_segment_model_dropdown = gr.Dropdown(
+                                    choices=["Riksarkivet/RmtDet_lines"],
+                                    value="Riksarkivet/RmtDet_lines",
+                                    label="Line segment model",
+                                    info="Will add more models later!",
+                                )
+
+                                htr_tool_transcriber_model_dropdown = gr.Dropdown(
+                                    choices=["Riksarkivet/SATRN_transcriber", "microsoft/trocr-base-handwritten"],
+                                    value="Riksarkivet/SATRN_transcriber",
+                                    label="Transcribe model",
+                                    info="Will add more models later!",
+                                )
 
+                    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.Row():
-                gr.Markdown(
-                    """
-                    Image viewer for xml output: 
-                    <p align="center">
-                        <a href="https://huggingface.co/spaces/Riksarkivet/Viewer_demo">
-                            <img src="https://huggingface.co/datasets/huggingface/badges/raw/main/open-in-hf-spaces-xl-dark.svg" alt="Badge 1">
-                        </a>
-                    </p>
-                    
-                    """
-                )
 
-        with gr.Column(scale=4):
-            fast_track_output_image = gr.Image(label="HTR results visualizer", type="numpy", tool="editor", height=650)
+                with gr.Row(visible=False) as image_viewer_tab:
+                    text_polygon_dict = gr.Variable()
 
-        with gr.Row(visible=False) as api_placeholder:
-            htr_pipeline_button_api = gr.Button("Run pipeline", variant="primary", visible=False, scale=1)
+                    fast_track_output_image = gr.Image(
+                        label="Image Viewer", type="numpy", height=600, interactive=False
+                    )
 
     xml_rendered_placeholder_for_api = gr.Textbox(visible=False)
-    htr_pipeline_button.click(
+
+    htr_event_click_event = htr_pipeline_button.click(
         fast_track.segment_to_xml,
         inputs=[fast_track_input_region_image, radio_file_input],
-        outputs=[fast_track_output_image, fast_file_downlod, fast_file_downlod],
+        outputs=[fast_file_downlod, fast_file_downlod],
     )
 
     htr_pipeline_button_api.click(
@@ -75,3 +154,35 @@ with gr.Blocks() as htr_tool_tab:
         outputs=[xml_rendered_placeholder_for_api],
         api_name="predict",
     )
+
+    def update_selected_tab_output_and_setting():
+        return gr.update(visible=True), gr.update(visible=False)
+
+    def update_selected_tab_image_viewer():
+        return 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]
+    )
+
+    tab_image_viewer_selector.select(
+        fn=update_selected_tab_image_viewer, outputs=[output_and_setting_tab, image_viewer_tab]
+    )
+
+    def stop_function():
+        from src.htr_pipeline.utils import pipeline_inferencer
+
+        pipeline_inferencer.terminate = True
+        gr.Info("The HTR execution was halted")
+
+    stop_htr_button.click(fn=stop_function, inputs=None, outputs=None, cancels=[htr_event_click_event])
+
+    run_image_visualizer_button.click(
+        fn=fast_track.visualize_image_viewer,
+        inputs=fast_track_input_region_image,
+        outputs=[fast_track_output_image, text_polygon_dict],
+    )
+
+    fast_track_output_image.select(
+        fast_track.get_text_from_coords, inputs=text_polygon_dict, outputs=selection_text_from_image_viewer
+    )