MikkoLipsanen's picture
Update app.py
7c01649 verified
raw
history blame
4.92 kB
from optimum.onnxruntime import ORTModelForVision2Seq
from transformers import TrOCRProcessor
from ultralytics import YOLO
import gradio as gr
import numpy as np
import onnxruntime
import time
from plotting_functions import PlotHTR
from segment_image import SegmentImage
from onnx_text_recognition import TextRecognition
LINE_MODEL_PATH = "Kansallisarkisto/multicentury-textline-detection"
REGION_MODEL_PATH = "Kansallisarkisto/court-records-region-detection"
TROCR_PROCESSOR_PATH = "Kansallisarkisto/multicentury-htr-model-onnx"
TROCR_MODEL_PATH = "Kansallisarkisto/multicentury-htr-model-onnx"
def get_segmenter():
"""Initialize segmentation class."""
try:
segmenter = SegmentImage(line_model_path=LINE_MODEL_PATH,
device='cpu',
line_iou=0.3,
region_iou=0.5,
line_overlap=0.5,
line_nms_iou=0.7,
region_nms_iou=0.3,
line_conf_threshold=0.25,
region_conf_threshold=0.5,
region_model_path=REGION_MODEL_PATH,
order_regions=True,
region_half_precision=False,
line_half_precision=False)
return segmenter
except Exception as e:
print('Failed to initialize SegmentImage class: %s' % e)
def get_recognizer():
"""Initialize text recognition class."""
try:
recognizer = TextRecognition(
processor_path = TROCR_PROCESSOR_PATH,
model_path = TROCR_MODEL_PATH,
device = 'cpu',
half_precision = True,
line_threshold = 100
)
return recognizer
except Exception as e:
print('Failed to initialize TextRecognition class: %s' % e)
segmenter = get_segmenter()
recognizer = get_recognizer()
plotter = PlotHTR()
color_codes = """**Text region type:** <br>
Paragraph ![#EE1289](https://placehold.co/15x15/EE1289/EE1289.png)
Marginalia ![#00C957](https://placehold.co/15x15/00C957/00C957.png)
Page number ![#0000FF](https://placehold.co/15x15/0000FF/0000FF.png)"""
def merge_lines(segment_predictions):
img_lines = []
for region in segment_predictions:
img_lines += region['lines']
return img_lines
def get_text_predictions(image, segment_predictions, recognizer):
"""Collects text prediction data into dicts based on detected text regions."""
img_lines = merge_lines(segment_predictions)
height, width = segment_predictions[0]['img_shape']
# Process all lines of an image
texts = recognizer.process_lines(img_lines, image, height, width)
return texts
# Run demo code
with gr.Blocks(theme=gr.themes.Monochrome(), title="Multicentury HTR Demo") as demo:
gr.Markdown("# Multicentury HTR Demo")
with gr.Tab("Text content"):
with gr.Row():
input_img = gr.Image(label="Input image", type="pil")
textbox = gr.Textbox(label="Predicted text content", lines=10)
button = gr.Button("Process image")
processing_time = gr.Markdown()
with gr.Tab("Text regions"):
region_img = gr.Image(label="Predicted text regions", type="numpy")
gr.Markdown(color_codes)
with gr.Tab("Text lines"):
line_img = gr.Image(label="Predicted text lines", type="numpy")
gr.Markdown(color_codes)
def run_pipeline(image):
# Predict region and line segments
start = time.time()
segment_predictions = segmenter.get_segmentation(image)
if segment_predictions:
region_plot = plotter.plot_regions(segment_predictions, image)
line_plot = plotter.plot_lines(segment_predictions, image)
text_predictions = get_text_predictions(np.array(image), segment_predictions, recognizer)
text = "\n".join(text_predictions)
end = time.time()
proc_time = end - start
proc_time_str = f"Processing time: {proc_time:.4f}s"
return {
region_img: region_plot,
line_img: line_plot,
textbox: text,
processing_time: proc_time_str
}
else:
end = time.time()
proc_time = end - start
proc_time_str = f"Processing time: {proc_time:.4f}s"
return {
region_img: None,
line_img: None,
textbox: None,
processing_time: proc_time_str
}
button.click(fn=run_pipeline,
inputs=input_img,
outputs=[region_img, line_img, textbox, processing_time])
if __name__ == "__main__":
demo.launch()