Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| """Demo.ipynb | |
| Automatically generated by Colaboratory. | |
| Original file is located at | |
| https://colab.research.google.com/drive/1Icb8zeoaudyTDOKM1QySNay1cXzltRAp | |
| """ | |
| import gradio as gr | |
| from PIL import Image | |
| import re | |
| import torch | |
| import torch.nn as nn | |
| from warnings import simplefilter | |
| simplefilter('ignore') | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| # Seting up the model | |
| from transformers import DonutProcessor, VisionEncoderDecoderModel | |
| print('Loading the base model ....') | |
| base_model = VisionEncoderDecoderModel.from_pretrained('Edgar404/donut-shivi-recognition') | |
| base_processor = DonutProcessor.from_pretrained('Edgar404/donut-shivi-recognition') | |
| print('Loading complete') | |
| print('Loading the optimized model ....') | |
| optimized_model = VisionEncoderDecoderModel.from_pretrained('Edgar404/donut-shivi-cheques_KD_320', torch_dtype = torch.bfloat16 ) | |
| optimized_processor = DonutProcessor.from_pretrained('Edgar404/donut-shivi-cheques_KD_320') | |
| print('Loading complete') | |
| # setting | |
| def process_image(image , mode = 'optimized' ): | |
| """ Function that takes an image and perform an OCR using the model DonUT via the task document | |
| parsing | |
| parameters | |
| __________ | |
| image : a machine readable image of class PIL or numpy""" | |
| model = optimized_model if mode == 'optimized' else base_model | |
| processor = optimized_processor if mode == 'optimized' else base_processor | |
| d_type = torch.bfloat16 if ((mode == 'optimized') & (device =='cuda')) else torch.float32 | |
| model.to(device) | |
| model.eval() | |
| task_prompt = "<s_cord-v2>" | |
| decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids | |
| pixel_values = processor(image, return_tensors="pt").pixel_values | |
| outputs = model.generate( | |
| pixel_values.to(device , dtype = d_type), | |
| decoder_input_ids=decoder_input_ids.to(device), | |
| max_length=model.decoder.config.max_position_embeddings, | |
| pad_token_id=processor.tokenizer.pad_token_id, | |
| eos_token_id=processor.tokenizer.eos_token_id, | |
| use_cache=True, | |
| bad_words_ids=[[processor.tokenizer.unk_token_id]], | |
| return_dict_in_generate=True, | |
| ) | |
| sequence = processor.batch_decode(outputs.sequences)[0] | |
| sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "") | |
| sequence = re.sub(r"<.*?>", "", sequence, count=1).strip() | |
| output = processor.token2json(sequence) | |
| return output | |
| def image_classifier(image , mode): | |
| return process_image(image , mode) | |
| examples_list = [['./test_images/test_0.jpg' ,"base"] , | |
| ['./test_images/test_1.jpg','base'], | |
| ['./test_images/test_2.jpg' ,"base"], | |
| ['./test_images/test_3.jpg','base'], | |
| ['./test_images/test_4.jpg','base'], | |
| ['./test_images/test_5.jpg' ,"base"], | |
| ['./test_images/test_6.jpg' ,"base"], | |
| ['./test_images/test_7.jpg','base'], | |
| ['./test_images/test_8.jpg','base'], | |
| ['./test_images/test_9.jpg','base'] | |
| ] | |
| demo = gr.Interface(fn=image_classifier, inputs=["image", | |
| gr.Radio(["base" , "optimized"], label="mode")], | |
| outputs="text", | |
| examples = examples_list ) | |
| demo.launch(share = True , debug = True) |