Initial version

app.py

@@ -0,0 +1,53 @@
+import torch
+from sconf import Config
+from PIL import Image, ImageOps
+from donut import DonutConfig, DonutModel
+import warnings
+
+warnings.filterwarnings("ignore")
+
+from transformers import logging
+
+logging.set_verbosity_warning()
+
+config = Config(default="./config.yaml")
+
+model = DonutModel.from_pretrained(
+                config.pretrained_model_name_or_path,
+                input_size=config.input_size,
+                max_length=config.max_position_embeddings, #self.config.max_length,
+                align_long_axis=config.align_long_axis,
+                ignore_mismatched_sizes=True,
+            )
+
+task_name = "matricula"
+task_prompt = f"<s_{task_name}>"
+    
+def predict_matricula(model, task_name, image):
+    image = ImageOps.exif_transpose(image)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    model.eval()
+    model.to(device)
+
+    result = model.inference(image=image, prompt=f"<s_{task_name}>")["predictions"][0]
+    return result
+
+
+import gradio as gr
+
+with gr.Blocks() as demo:
+    fn_predict=lambda x:predict_matricula(model, task_name="matricula", image=x)
+    gr.Markdown(
+    """
+    # Demo: Donut 🍩 for DR Matriculas
+    Dominican Vehicle **Matriculas OCR** Infering
+    """)
+    with gr.Row():
+        input_image=gr.Image(label="Matricula", sources="upload", type="pil", show_label=True)
+    with gr.Row():
+        output_json=gr.JSON(label="Matricula JSON", show_label=True, value={})
+    submit_btn = gr.Button("Submit")
+    submit_btn.click(fn=fn_predict, inputs=[input_image], outputs=[output_json])
+
+demo.launch(share=True)

config.yaml

@@ -0,0 +1,35 @@
+resume_from_checkpoint_path: None
+result_path: 'trainer/result'
+pretrained_model_name_or_path: 'marzanconsulting/donut-dr-matriculas-ocr'
+dataset_name_or_paths: 
+  - 'trainer/dataset'
+task_start_tokens: 
+  - '<s_matricula>'
+sort_json_key: False
+train_batch_sizes: 
+  - 5
+val_batch_sizes: 
+  - 1
+input_size: 
+  - 960
+  - 1280
+max_length: 868
+max_position_embeddings: 868
+align_long_axis: False
+num_nodes: 1
+seed: 2022
+lr: 3e-05
+warmup_steps: 300
+num_training_samples_per_epoch: 800
+max_epochs: 50
+max_steps: -1
+num_workers: 12
+val_check_interval: 1.0
+check_val_every_n_epoch: 3
+gradient_clip_val: 1.0
+verbose: True
+model_dir: ''
+tensorboard_dir: ''
+checkpoint_dir: 
+exp_name: 'train_cord_matriculas'
+exp_version: 'mmc_v1'

requirements.txt

@@ -0,0 +1,8 @@
+donut-python==1.0.9
+torchvision==0.14.0
+pytorch-lightning==1.6.4
+tokenizers>=0.11,<=0.12.1
+transformers==4.25.1
+sentencepiece
+timm==0.6.13
+datasets[vision]