Spaces:
Runtime error
Runtime error
File size: 6,672 Bytes
ee21b96 08374eb 59fdcb0 08374eb 7883098 35a37ac 08374eb ee21b96 d23d897 8239234 ee21b96 803e48e ee21b96 fb33609 955051b ee21b96 cd31959 ee21b96 e235062 ee21b96 e235062 ee21b96 2915058 ee21b96 803e48e cd31959 0c80503 aa287a2 0c80503 ee21b96 8f7f77a 1cdafc5 ee21b96 ab591a3 ee21b96 e235062 ee21b96 e235062 ee21b96 0c80503 ab591a3 ee21b96 8f7f77a 2915058 ee21b96 2915058 8f7f77a ee21b96 |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 |
import os
os.system('cd fairseq;'
'pip install ./; cd ..')
os.system('cd ezocr;'
'pip install .; cd ..')
import torch
import numpy as np
from fairseq import utils, tasks
from fairseq import checkpoint_utils
from utils.eval_utils import eval_step
from data.mm_data.ocr_dataset import ocr_resize
from tasks.mm_tasks.ocr import OcrTask
from PIL import Image, ImageDraw
from torchvision import transforms
from typing import List, Tuple
import cv2
from easyocrlite import ReaderLite
import gradio as gr
# Register refcoco task
tasks.register_task('ocr', OcrTask)
os.system('wget https://shuangqing-multimodal.oss-cn-zhangjiakou.aliyuncs.com/ocr_general_clean.pt; '
'mkdir -p checkpoints; mv ocr_general_clean.pt checkpoints/ocr_general_clean.pt')
# turn on cuda if GPU is available
use_cuda = torch.cuda.is_available()
# use fp16 only when GPU is available
use_fp16 = True
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]
Rect = Tuple[int, int, int, int]
FourPoint = Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int], Tuple[int, int]]
def four_point_transform(image: np.ndarray, rect: FourPoint) -> np.ndarray:
(tl, tr, br, bl) = rect
widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
maxWidth = max(int(widthA), int(widthB))
# compute the height of the new image, which will be the
# maximum distance between the top-right and bottom-right
# y-coordinates or the top-left and bottom-left y-coordinates
heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
maxHeight = max(int(heightA), int(heightB))
dst = np.array(
[[0, 0], [maxWidth - 1, 0], [maxWidth - 1, maxHeight - 1], [0, maxHeight - 1]],
dtype="float32",
)
# compute the perspective transform matrix and then apply it
M = cv2.getPerspectiveTransform(rect, dst)
warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
return warped
def get_images(img: str, reader: ReaderLite, **kwargs):
results = reader.process(img, **kwargs)
return results
def draw_boxes(image, bounds, color='red', width=4):
draw = ImageDraw.Draw(image)
for i, bound in enumerate(bounds):
p0, p1, p2, p3 = bound
draw.text(p0, str(i+1), fill=color)
draw.line([*p0, *p1, *p2, *p3, *p0], fill=color, width=width)
return image
def encode_text(task, text, length=None, append_bos=False, append_eos=False):
bos_item = torch.LongTensor([task.src_dict.bos()])
eos_item = torch.LongTensor([task.src_dict.eos()])
pad_idx = task.src_dict.pad()
s = task.tgt_dict.encode_line(
line=task.bpe.encode(text),
add_if_not_exist=False,
append_eos=False
).long()
if length is not None:
s = s[:length]
if append_bos:
s = torch.cat([bos_item, s])
if append_eos:
s = torch.cat([s, eos_item])
return s
def patch_resize_transform(patch_image_size=480, is_document=False):
_patch_resize_transform = transforms.Compose(
[
lambda image: ocr_resize(
image, patch_image_size, is_document=is_document, split='test',
),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std),
]
)
return _patch_resize_transform
reader = ReaderLite(gpu=True)
overrides={"eval_cider": False, "beam": 5, "max_len_b": 64, "patch_image_size": 480,
"orig_patch_image_size": 224, "no_repeat_ngram_size": 0, "seed": 42}
models, cfg, task = checkpoint_utils.load_model_ensemble_and_task(
utils.split_paths('checkpoints/ocr_general_clean.pt'),
arg_overrides=overrides
)
# Move models to GPU
for model in models:
model.eval()
if use_fp16:
model.half()
if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
model.cuda()
model.prepare_for_inference_(cfg)
# Initialize generator
generator = task.build_generator(models, cfg.generation)
bos_item = torch.LongTensor([task.src_dict.bos()])
eos_item = torch.LongTensor([task.src_dict.eos()])
pad_idx = task.src_dict.pad()
# Construct input for caption task
def construct_sample(task, image: Image, patch_image_size=480):
patch_image = patch_resize_transform(patch_image_size)(image).unsqueeze(0)
patch_mask = torch.tensor([True])
src_text = encode_text(task, "图片上的文å—是什么?", append_bos=True, append_eos=True).unsqueeze(0)
src_length = torch.LongTensor([s.ne(pad_idx).long().sum() for s in src_text])
sample = {
"id":np.array(['42']),
"net_input": {
"src_tokens": src_text,
"src_lengths": src_length,
"patch_images": patch_image,
"patch_masks": patch_mask,
},
"target": None
}
return sample
# Function to turn FP32 to FP16
def apply_half(t):
if t.dtype is torch.float32:
return t.to(dtype=torch.half)
return t
def ocr(img):
out_img = Image.open(img)
results = get_images(img, reader, text_confidence=0.7, text_threshold=0.4, link_threshold=0.3)
box_list, image_list = zip(*results)
draw_boxes(out_img, box_list)
ocr_result = []
for i, (box, image) in enumerate(zip(box_list, image_list)):
image = Image.fromarray(image)
sample = construct_sample(task, image, cfg.task.patch_image_size)
sample = utils.move_to_cuda(sample) if use_cuda else sample
sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample
with torch.no_grad():
result, scores = eval_step(task, generator, models, sample)
ocr_result.append(str(i+1) + '\t' + result[0]['ocr'].replace(' ', ''))
result = '\n'.join(ocr_result)
return out_img, result
title = "OFA-OCR"
description = "Gradio Demo for OFA-OCR. Upload your own image or click any one of the examples, and click " \
"\"Submit\" and then wait for the generated OCR result. "
article = "<p style='text-align: center'><a href='https://github.com/OFA-Sys/OFA' target='_blank'>OFA Github " \
"Repo</a></p> "
examples = [['shupai.png'], ['chinese.jpg'], ['ms2.jpg']]
io = gr.Interface(fn=ocr, inputs=gr.inputs.Image(type='filepath', label='Image'),
outputs=[gr.outputs.Image(type='pil', label='Image'), gr.outputs.Textbox(label="OCR result")],
title=title, description=description, article=article, examples=examples,
allow_flagging='never', allow_screenshot=False)
io.launch(cache_examples=False)
|