Spaces:
Sleeping
Sleeping
| import os | |
| import gradio as gr | |
| from helper.examples.examples import DemoImages | |
| from helper.utils import TrafficDataHandler | |
| from src.htr_pipeline.gradio_backend import FastTrack, SingletonModelLoader | |
| model_loader = SingletonModelLoader() | |
| fast_track = FastTrack(model_loader) | |
| images_for_demo = DemoImages() | |
| terminate = False | |
| with gr.Blocks() as htr_tool_tab: | |
| with gr.Row(equal_height=True): | |
| with gr.Column(scale=2): | |
| with gr.Row(): | |
| fast_track_input_region_image = gr.Image( | |
| label="Image to run HTR on", type="numpy", tool="editor", elem_id="image_upload", height=395 | |
| ) | |
| with gr.Row(): | |
| with gr.Tab("HTRFLOW") as tab_output_and_setting_selector: | |
| with gr.Row(): | |
| stop_htr_button = gr.Button( | |
| value="Stop run", | |
| variant="stop", | |
| ) | |
| htr_pipeline_button = gr.Button( | |
| "Run ", | |
| variant="primary", | |
| visible=True, | |
| elem_id="run_pipeline_button", | |
| ) | |
| htr_pipeline_button_var = gr.State(value="htr_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("Visualize") as tab_image_viewer_selector: | |
| with gr.Row(): | |
| gr.Button( | |
| value="Image Viewer (demo)", | |
| 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.Tab("Compare") 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", "Page XML"], | |
| value=["Txt", "Page 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.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() | |
| fast_track_output_image = gr.Image( | |
| 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( | |
| fast_track.segment_to_xml, | |
| inputs=[fast_track_input_region_image, radio_file_input], | |
| outputs=[fast_file_downlod, fast_file_downlod], | |
| ) | |
| htr_pipeline_button_api.click( | |
| fast_track.segment_to_xml_api, | |
| inputs=[fast_track_input_region_image], | |
| outputs=[xml_rendered_placeholder_for_api], | |
| api_name="predict", | |
| ) | |
| def dummy_update_htr_tool_transcriber_model_dropdown(htr_tool_transcriber_model_dropdown): | |
| return gr.update(value="Riksarkivet/satrn_htr") | |
| htr_tool_transcriber_model_dropdown.change( | |
| fn=dummy_update_htr_tool_transcriber_model_dropdown, | |
| inputs=htr_tool_transcriber_model_dropdown, | |
| outputs=htr_tool_transcriber_model_dropdown, | |
| ) | |
| def update_selected_tab_output_and_setting(): | |
| 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), 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, model_compare_selector], | |
| ) | |
| tab_image_viewer_selector.select( | |
| 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(): | |
| 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 | |
| ) | |
| SECRET_KEY = os.environ.get("AM_I_IN_A_DOCKER_CONTAINER", False) | |
| if SECRET_KEY: | |
| htr_pipeline_button.click(fn=TrafficDataHandler.store_metric_data, inputs=htr_pipeline_button_var) | |