import os import pandas as pd os.system('cd fairseq;' 'pip install ./; cd ..') os.system('cd ezocr;' 'pip install .; cd ..') os.system('pip install --upgrade tensorflow-gpu==1.15;' 'pip install "modelscope[cv]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html') 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 from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks from modelscope.outputs import OutputKeys from modelscope.preprocessors.image import load_image # 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[0]+5, p0[1]+5), str(i+1), fill=color, align='center') 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()]) 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) ocr_detection = pipeline(Tasks.ocr_detection, model='damo/cv_resnet18_ocr-detection-line-level_damo') 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): boxes = ocr_detection(img)[OutputKeys.POLYGONS] image = cv2.imread(img) out_img = Image.open(img) ocr_result = list() for i, box in boxes: # 因为检测结果是四边形，所以用透视变化转为长方形 post1 = box.reshape((4, 2)).astype(np.float32) width = box[4] - box[0] height = box[5] - box[1] post2 = np.float32([[0, 0], [width, 0], [width, height], [0, height]]) M = cv2.getPerspectiveTransform(post1, post2) new_img = cv2.warpPerspective(image, M, (width, height)) new_img_pil = Image.fromarray(cv2.cvtColor(new_img, cv2.COLOR_BGR2RGB)) # 开启文字识别 sample = construct_sample(task, new_img_pil, 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), result[0]['ocr'].replace(' ', '')]) result = pd.DataFrame(ocr_result, columns=['Box ID', 'Text']) # results = get_images(img, reader, text_confidence=0.7, text_threshold=0.4, # link_threshold=0.43, slope_ths=0., add_margin=0.02) # box_list, image_list = zip(*results) draw_boxes(out_img, boxes) # # 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), result[0]['ocr'].replace(' ', '')]) # # result = pd.DataFrame(ocr_result, columns=['Box ID', 'Text']) return out_img, result title = "Chinese OCR" description = "Gradio Demo for Chinese OCR based on OFA. "\ "Upload your own image or click any one of the examples, and click " \ "\"Submit\" and then wait for the generated OCR result. " \ "\n中文OCR体验区。欢迎上传图片，静待检测文字返回~" article = "

OFA Github " \ "Repo

" examples = [['shupai.png'], ['chinese.jpg'], ['gaidao.jpeg'], ['qiaodaima.png'], ['benpao.jpeg'], ['wanli.png'], ['xsd.jpg']] io = gr.Interface(fn=ocr, inputs=gr.inputs.Image(type='filepath', label='Image'), outputs=[gr.outputs.Image(type='pil', label='Image'), gr.outputs.Dataframe(headers=['Box ID', 'Text'], type='pandas', label='OCR Results')], title=title, description=description, article=article, examples=examples) io.launch()