First model version

app.py

@@ -43,10 +43,9 @@ def infer(filepath):
     patchs = [preprocess(patch, config.dataset_image_width, config.dataset_image_height) for patch in patchs]
     patchs = torch.cat(patchs, dim=0)
     res = model(patchs)
-    rec_result = postprocess(res, charset, 'alignment')[0]
-    print(rec_result)
+    result_words = postprocess(res, charset, 'alignment')[0]
     
-    visual_image = det_demo.visualization(image.copy(), result_polygons, result_masks, result_boxes)
+    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)

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, 1, (0,0,255), 2)
         '''
         for box in boxes:
             cv2.rectangle(image,(box[0], box[1]), (box[2], box[3]), (0,0,255), 2)