fixed some minior bugs in visuazliser

app/main.py

@@ -1,3 +1,4 @@
+import shutil
 import gradio as gr
 import os
 from app.gradio_config import css, theme
@@ -5,7 +6,6 @@ from app.tabs.submit import (
     submit,
     custom_template_yaml,
     collection_submit_state,
-    batch_image_gallery,
 )
 from app.tabs.visualizer import visualizer, collection_viz_state, viz_image_gallery
 from app.tabs.templating import (
@@ -53,6 +53,12 @@ with gr.Blocks(title="HTRflow", theme=theme, css=css) as demo:
         with gr.Tab(label="Visualize Result") as tab_visualizer:
             visualizer.render()
 
+    @demo.load()
+    def inital_yaml_code():
+        tmp_dir = "tmp/"
+        if os.path.exists(tmp_dir) and os.path.isdir(tmp_dir):
+            shutil.rmtree(tmp_dir)
+
     @demo.load(
         inputs=[template_output_yaml_code],
         outputs=[template_output_yaml_code],
@@ -85,12 +91,6 @@ with gr.Blocks(title="HTRflow", theme=theme, css=css) as demo:
         fn=sync_gradio_objects,
     )
 
-    # tab_visualizer.select(
-    #     inputs=[batch_image_gallery, viz_image_gallery],
-    #     outputs=[viz_image_gallery],
-    #     fn=sync_gradio_objects,
-    # )
-
     tab_visualizer.select(
         inputs=[collection_submit_state, collection_viz_state],
         outputs=[collection_viz_state],

app/tabs/submit.py

@@ -93,6 +93,8 @@ def run_htrflow(custom_template_yaml, batch_image_gallery, progress=gr.Progress(
     progress(0, desc="HTRflow: Starting")
     time.sleep(0.3)
 
+    print(temp_config)
+
     if batch_image_gallery is None:
         gr.Warning("HTRflow: You must upload atleast 1 image or more")
 
@@ -135,6 +137,8 @@ def tracking_exported_files(tmp_output_paths):
 
     exported_files = set()
 
+    print(tmp_output_paths)
+
     for tmp_folder in tmp_output_paths:
         for ext in accepted_extensions:
             search_pattern = os.path.join(tmp_folder, "**", f"*{ext}")

app/tabs/visualizer.py

@@ -4,10 +4,65 @@ import numpy as np
 from htrflow.volume.volume import Collection
 from htrflow.utils.draw import draw_polygons
 from htrflow.utils import imgproc
-
+import time
 from htrflow.results import Segment
 
 
+def load_visualize_state_from_submit(col: Collection, progress):
+    results = []
+
+    time.sleep(1)
+
+    total_steps = len(col.pages)
+
+    for page_idx, page_node in enumerate(col):
+        page_image = page_node.image.copy()
+
+        progress((page_idx + 1) / total_steps, desc=f"Running Visualizer")
+
+        lines = list(page_node.traverse(lambda node: node.is_line()))
+
+        recog_conf_values = {
+            i: list(zip(tr.texts, tr.scores)) if (tr := ln.text_result) else []
+            for i, ln in enumerate(lines)
+        }
+
+        recog_df = pd.DataFrame(
+            [
+                {"Transcription": text, "Confidence Score": f"{score:.4f}"}
+                for values in recog_conf_values.values()
+                for text, score in values
+            ]
+        )
+
+        line_polygons = []
+        line_crops = []
+        for ln in lines:
+            seg: Segment = ln.data.get("segment")
+            if not seg:
+                continue
+
+            cropped_line_img = imgproc.crop(page_image, seg.bbox)
+            cropped_line_img = np.clip(cropped_line_img, 0, 255).astype(np.uint8)
+            line_crops.append(cropped_line_img)
+
+            if seg.polygon is not None:
+                line_polygons.append(seg.polygon)
+
+        annotated_image = draw_polygons(page_image, line_polygons)
+        annotated_page_node = np.clip(annotated_image, 0, 255).astype(np.uint8)
+        results.append(
+            {
+                "page_image": page_node,
+                "annotated_page_node": annotated_page_node,
+                "line_crops": line_crops,
+                "recog_conf_values": recog_df,
+            }
+        )
+
+    return results
+
+
 with gr.Blocks() as visualizer:
     with gr.Column(variant="panel"):
         with gr.Row():
@@ -28,73 +83,92 @@ with gr.Blocks() as visualizer:
                     "Visualize", scale=0, min_width=200, variant="secondary"
                 )
 
-            with gr.Column():
-                # image_visualizer_annotation = gr.Image(
-                #     interactive=False,
-                # )
+                progress_bar = gr.Textbox(visible=False, show_label=False)
 
-                line2 = gr.Gallery(
+            with gr.Column():
+                cropped_image_gallery = gr.Gallery(
+                    interactive=False,
+                    preview=True,
+                    label="Cropped Polygons",
+                    height=200,
+                )
+                df_for_cropped_images = gr.Dataframe(
+                    label="Cropped Transcriptions",
+                    headers=["Transcription", "Confidence Score"],
                     interactive=False,
                 )
-                textlines = gr.Dataframe()
 
-    # @viz_image_gallery.select(outputs=image_visualizer_annotation)
-    # def return_image_from_gallery(evt: gr.SelectData):
-    #     return evt.value["image"]["path"]
+    def on_visualize_button_clicked(collection_viz, progress=gr.Progress()):
+        """
+        This function:
+        - Receives the collection (collection_viz).
+        - Processes it into 'results' (list of dicts with annotated_page_node, line_crops, dataframe).
+        - Returns:
+            1) 'results' as state
+            2) List of annotated_page_node images (one per page) to populate viz_image_gallery
+        """
+        if not collection_viz:
+            return None, []
+
+        results = load_visualize_state_from_submit(collection_viz, progress)
+        annotated_images = [r["annotated_page_node"] for r in results]
+        return results, annotated_images, gr.skip()
+
+    visualize_button.click(lambda: gr.update(visible=True), outputs=progress_bar).then(
+        fn=on_visualize_button_clicked,
+        inputs=collection_viz_state,
+        outputs=[result_collection_viz_state, viz_image_gallery, progress_bar],
+    ).then(lambda: gr.update(visible=False), outputs=progress_bar)
+
+    @viz_image_gallery.change(
+        inputs=result_collection_viz_state,
+        outputs=[cropped_image_gallery, df_for_cropped_images],
+    )
+    def update_c_gallery_and_dataframe(results):
+        selected = results[0]
+        return selected["line_crops"], selected["recog_conf_values"]
 
-    @visualize_button.click(
-        outputs=[result_collection_viz_state, viz_image_gallery, line2, textlines]
+    @viz_image_gallery.select(
+        inputs=result_collection_viz_state,
+        outputs=[cropped_image_gallery, df_for_cropped_images],
     )
-    def testie_load_pickle():
-        col = Collection.from_pickle(".cache/HTRflow_demo_output.pickle")
+    def on_dataframe_select(evt: gr.SelectData, results):
+        """
+        evt.index => the index of the selected image in the gallery
+        results => the state object from result_collection_viz_state
 
-        results = []
-        for page_idx, page_node in enumerate(col):
-            page_image = page_node.image.copy()
+        Return the line crops and the recognized text for that index.
+        """
+        if results is None or evt.index is None:
+            return [], pd.DataFrame(columns=["Transcription", "Confidence Score"])
 
-            lines = list(page_node.traverse(lambda node: node.is_line()))
+        idx = evt.index
+        selected = results[idx]
 
-            recog_conf_values = {
-                i: list(zip(tr.texts, tr.scores)) if (tr := ln.text_result) else []
-                for i, ln in enumerate(lines)
-            }
+        return selected["line_crops"], selected["recog_conf_values"]
 
-            recog_df = pd.DataFrame(
-                [
-                    {"Transcription": text, "Confidence Score": f"{score:.4f}"}
-                    for values in recog_conf_values.values()
-                    for text, score in values
-                ]
-            )
-
-            line_polygons = []
-            line_crops = []
-            for ln in lines:
-                seg: Segment = ln.data.get("segment")
-                if not seg:
-                    continue
-
-                cropped_line_img = imgproc.crop(page_image, seg.bbox)
-                cropped_line_img = np.clip(cropped_line_img, 0, 255).astype(np.uint8)
-                line_crops.append(cropped_line_img)
-
-                if seg.polygon is not None:
-                    line_polygons.append(seg.polygon)
-
-            annotated_image = draw_polygons(page_image, line_polygons)
-            annotated_page_node = np.clip(annotated_image, 0, 255).astype(np.uint8)
-            results.append(
-                {
-                    "page_image": page_node,
-                    "annotated_page_node": annotated_page_node,
-                    "line_crops": line_crops,
-                    "recog_conf_values": recog_df,
-                }
-            )
-
-        return (
-            results,
-            [results[0]["annotated_page_node"]],
-            results[0]["line_crops"],
-            results[0]["recog_conf_values"],
-        )
+    @df_for_cropped_images.select(
+        outputs=[cropped_image_gallery],
+    )
+    def on_dataframe_select(evt: gr.SelectData):
+        return gr.update(selected_index=evt.index[0])
+
+    @cropped_image_gallery.select(
+        inputs=df_for_cropped_images, outputs=df_for_cropped_images
+    )
+    def return_image_from_gallery(df, evt: gr.SelectData):
+        selected_index = evt.index
+
+        def highlight_row(row):
+            return [
+                (
+                    "border: 1px solid blue; font-weight: bold"
+                    if row.name == selected_index
+                    else ""
+                )
+                for _ in row
+            ]
+
+        styler = df.style.apply(highlight_row, axis=1)
+
+        return styler