Spaces:
Runtime error
Runtime error
File size: 2,163 Bytes
44d692b 206c3a7 d61c233 44d692b 20ae047 531a42f d6c5ffb 44d692b 20ae047 44d692b 1b6d59b 44d692b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import re
import json
import torch
import numpy as np
import gradio as gr
from PIL import Image
from transformers import DonutProcessor, VisionEncoderDecoderModel
auth_tok="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"
def demo_process(input_img):
global processor, pretrained_model, task_prompt, task_name
input_img = Image.fromarray(input_img)
# prepare encoder inputs
pixel_values = processor(input_img.convert("RGB"), return_tensors="pt").pixel_values
pixel_values = pixel_values.to(device)
# prepare decoder inputs
task_prompt = "<s_lotoquine>"
decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
decoder_input_ids = decoder_input_ids.to(device)
# autoregressively generate sequence
outputs = pretrained_model.generate(
pixel_values,
decoder_input_ids=decoder_input_ids,
max_length=pretrained_model.decoder.config.max_position_embeddings,
early_stopping=True,
pad_token_id=processor.tokenizer.pad_token_id,
eos_token_id=processor.tokenizer.eos_token_id,
use_cache=True,
num_beams=1,
bad_words_ids=[[processor.tokenizer.unk_token_id]],
return_dict_in_generate=True,
)
# turn into JSON
seq = processor.batch_decode(outputs.sequences)[0]
seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
seq = re.sub(r"<.*?>", "", seq, count=1).strip() # remove first task start token
seq = processor.token2json(seq)
return seq
processor = DonutProcessor.from_pretrained("Aigle974/donut-lotoquine",use_auth_token=auth_tok)
pretrained_model = VisionEncoderDecoderModel.from_pretrained("Aigle974/donut-lotoquine",use_auth_token=auth_tok)
processor.feature_extractor.do_align_long_axis = True
device ="cuda" if torch.cuda.is_available() else "cpu"
pretrained_model.to(device)
pretrained_model.eval()
demo = gr.Interface(
fn=demo_process,
inputs="image",
outputs="json",
title=f"Lotoquine Automatic Extraction by Fab",
)
demo.launch() |