import os os.system('pip install --upgrade --no-cache-dir gdown') os.system('gdown -O ./model_ctw.pth 16qgtD4UOhp0q5e2RYXE1dvuTz_ylZMyb') #os.system('unzip model_ctw.zip') os.system('gdown -O ./workdir.zip 10HxLehcJMY9rLd_OyH40HmrySZItuNDt') os.system('unzip workdir.zip') os.system('pip install "git+https://github.com/philferriere/cocoapi.git#egg=pycocotools&subdirectory=PythonAPI"') os.system('python setup.py build develop --user') import cv2 import pandas as pd import gradio as gr from det_demo import DetDemo from maskrcnn_benchmark.config import cfg from demo import get_model, preprocess, postprocess, load from utils import Config, Logger, CharsetMapper import torch def infer(img): filepath = './input.png' img.save(filepath) config = Config('configs/rec/train_abinet.yaml') config.model_vision_checkpoint = None model = get_model(config) model = load(model, 'workdir/train-abinet/best-train-abinet.pth') charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1) cfg.merge_from_file('./configs/det/r50_baseline.yaml') # manual override some options cfg.merge_from_list(["MODEL.DEVICE", "cpu"]) det_demo = DetDemo( cfg, min_image_size=800, confidence_threshold=0.7, output_polygon=True ) image = cv2.imread(filepath) print(image.shape) result_polygons, result_masks, result_boxes = det_demo.run_on_opencv_image(image) patchs = [image[box[1]:box[3], box[0]:box[2], :] for box in result_boxes] patchs = [preprocess(patch, config.dataset_image_width, config.dataset_image_height) for patch in patchs] patchs = torch.cat(patchs, dim=0) res = model(patchs) result_words = postprocess(res, charset, 'alignment')[0] visual_image = det_demo.visualization(image.copy(), result_polygons, result_masks, result_boxes, result_words) print(visual_image.shape) cv2.imwrite('result.jpg', visual_image) return ['result.jpg', pd.DataFrame(result_words)] blocks = gr.Blocks() input_image = gr.Image(label="image", type="pil") output_image = gr.Image(label="out_img", type="filepath") output_word = gr.Dataframe(label="out_word", headers=['word']) with blocks: gr.Markdown('''