First model version

README.md

@@ -1,8 +1,8 @@
 ---
-title: ContourNet ABI
-emoji: 🐢
-colorFrom: red
-colorTo: purple
+title: ContourNet
+emoji: 🚀
+colorFrom: blue
+colorTo: green
 sdk: gradio
 sdk_version: 3.0.10
 app_file: app.py

app.py

@@ -1,71 +1,62 @@
 import os
 os.system('pip install --upgrade --no-cache-dir gdown')
-os.system('gdown -O ./output/ctw/model_ctw.pth 1Ajslu_9WisuZ2nJGzE6qbD87aK6_ozzA')
-#os.system('gdown -O ./workdir.zip 1mYM_26qHUom_5NU7iutHneB_KHlLjL5y')
-#os.system('unzip workdir.zip')
+os.system('gdown -O ./output/ctw/model_ctw.pth 1ydYjuxJglLAJRVWjj6sKXM3aBL2pXRpG')
+os.system('gdown -O ./workdir.zip 1mYM_26qHUom_5NU7iutHneB_KHlLjL5y')
+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 glob
-import gradio as gr
-#from demo import get_model, preprocess, postprocess, load
-#from utils import Config, Logger, CharsetMapper
 import cv2
 import pandas as pd
+import gradio as gr
+
 from det_demo import DetDemo
 from maskrcnn_benchmark.config import cfg
 
-def process_image(filepath):
-  # rec model
-  '''
-  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)
-  '''
-  # det model
-  cfg.merge_from_file('./configs/det/r50_baseline.yaml')
-  cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
-  det_demo = DetDemo(
-      cfg,
-      min_image_size=800,
-      confidence_threshold=0.7,
-      output_polygon=True
-  )
-  
-  # detect
-  image = cv2.imread(filepath)
-  result_polygons, result_masks, result_boxes = det_demo.run_on_opencv_image(image)
-  
-  # cut patch
-  #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-  #patchs = [image[box[1]:box[3], box[0]:box[2], :] for box in result_boxes]
-  #patchs = [cv2.resize(patch, (128,32)) for patch in patchs]
-  #patchs = np.stack(patchs, axis=0).transpose(0,3,1,2)
-  '''
-  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.stack(patchs, dim=0)
-  print(patchs.shape)
-  res = model(patchs)
-  rec_result = postprocess(res, charset, 'alignment')[0]
-  print(rec_result)
-  '''
-  
-  # visual detect results
-  visual_image = det_demo.visualization(image.copy(), result_polygons, result_masks, result_boxes)
-  cv2.imwrite('result.jpg', visual_image)
-  return 'result.jpg'#, pd.DataFrame(result_words)
+from demo import get_model, preprocess, postprocess, load
+from utils import Config, Logger, CharsetMapper
+import torch
+
+
+def infer(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)
+    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)
+    
+    cv2.imwrite('result.jpg', visual_image)
+    return 'result.jpg', pd.DataFrame(result_words)
 
-title = "张博强毕设中期展示（文本识别部分）"
-description = "西北工业大学航海学院张博强毕设，目前识别部分进度为复现abinet，本网页为abinet复现的可视化web端展示"
-#article = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2103.06495.pdf'>Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition</a> | <a href='https://github.com/FangShancheng/ABINet'>Github Repo</a></p>"
 
-iface = gr.Interface(fn=process_image, 
-                     inputs=[gr.inputs.Image(label="image", type="filepath")], 
-                     outputs=[gr.outputs.Image()],#, gr.outputs.Textbox()
-                     title=title,
-                     description=description,
-                     examples=glob.glob('figs/test/*.png'))
-iface.launch(enable_queue=True)
+iface = gr.Interface(
+    fn=infer,
+    title="张博强毕设展示",
+    description="毕设题目：自然场景中任意形状文字的检测与识别\n目前进度：检测",
+    inputs=[gr.inputs.Image(label="image", type="filepath")],
+    outputs=[gr.outputs.Image(), gr.outputs.Dataframe(headers=['word'])], 
+    #examples=['example1.jpg', 'example2.jpg', 'example3.jpg'],
+    #article="<a href=\"https://github.com/MhLiao/MaskTextSpotterV3\">GitHub Repo</a>",
+).launch(enable_queue=True)

configs/det/r50_baseline.yaml

@@ -1,4 +1,4 @@
-OUTPUT_DIR: "./output/ctw/"
+OUTPUT_DIR: "./det_model/"
 MODEL:
   META_ARCHITECTURE: "GeneralizedRCNN"
   WEIGHT: "catalog://ImageNetPretrained/MSRA/R-50"

det_demo.py

@@ -573,13 +573,14 @@ class DetDemo(object):
             polygon = list(map(int, polygon))
         return polygon
 
-    def visualization(self, image, polygons, masks, boxes):
+    def visualization(self, image, polygons, masks, boxes, words):
         green = np.ones(image.shape).astype(np.uint8)
         green[...,0] = 0
         green[...,1] = 255
         green[...,2] = 0
-        for mask in masks:
+        for mask, word, box in zip(masks, words, boxes):
             image[mask] = image[mask] * 0.5 + green[mask] * 0.5
+            cv2.putText(image, word, (box[0], box[1]), cv2.FONT_HERSHEY_COMPLEX, 0.4, (0,0,255), 1)
         '''
         for box in boxes:
             cv2.rectangle(image,(box[0], box[1]), (box[2], box[3]), (0,0,255), 2)