First model version

app.py

@@ -10,20 +10,46 @@ import glob
 import gradio as gr
 from demo import get_model, preprocess, postprocess, load
 from utils import Config, Logger, CharsetMapper
-from accelerate import Accelerator
-
-accelerator = Accelerator()
-device = accelerator.device
 
 def process_image(image):
+  # 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').to(device)
+  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
+  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)
+  
 
   img = image.convert('RGB')
-  img = preprocess(img, config.dataset_image_width, config.dataset_image_height).to(device)
+  img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
   res = model(img)
   return postprocess(res, charset, 'alignment')[0][0]
 
@@ -32,8 +58,8 @@ description = "西北工业大学航海学院张博强毕设，目前识别部
 #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(type="pil")], 
-                     outputs=[gr.outputs.Textbox()],
+                     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'))