server.py

@@ -54,10 +54,172 @@ def ping():
 
 @app.post("/ocr_system/table")
 @web_try()
+<<<<<<< HEAD
+def rotate_bound_white_bg(self, image, angle):
+    
+    (h, w) = image.shape[:2]
+    (cX, cY) = (w // 2, h // 2)
+
+    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
+    cos = np.abs(M[0, 0])
+    sin = np.abs(M[0, 1])
+
+    nW = int((h * sin) + (w * cos))
+    nH = int((h * cos) + (w * sin))
+
+    M[0, 2] += (nW / 2) - cX
+    M[1, 2] += (nH / 2) - cY
+    return cv2.warpAffine(image, M, (nW, nH), borderValue=(255, 255, 255))
+
+
+class GetImageRotation(object):
+    def __init__(self):
+        self.ocr = PaddleOCR(use_angle_cls=True)
+        self.ocr_angle = PaddleOCR(use_angle_cls=True)
+
+    def get_real_rotation_when_null_rect(self, rect_list):
+        w_div_h_sum = 0
+        count = 0
+        for rect in rect_list:
+            p0 = rect[0]
+            p1 = rect[1]
+            p2 = rect[2]
+            p3 = rect[3]
+            width = abs(p1[0] - p0[0])
+            height = abs(p3[1] - p0[1])
+            w_div_h = width / height
+            if abs(w_div_h - 1.0) < 0.5:
+                count += 1
+                continue
+            w_div_h_sum += w_div_h
+        length = len(rect_list) - count
+        if length == 0:
+            length = 1
+        if w_div_h_sum / length >= 1.5:
+            return 1
+        else:
+            return 0
+            
+    def get_real_rotation_flag(self, rect_list):
+        ret_rect = []
+        w_div_h_mean = 0
+        real_rect_count = 0
+        rect_big_list = []
+        rect_small_list = []
+        w_div_h_sum_big = []
+        w_div_h_sum_small = []
+        for rect in rect_list:
+            p0 = rect[0]
+            p1 = rect[1]
+            p2 = rect[2]
+            p3 = rect[3]
+            width = abs(p1[0] - p0[0])
+            height = abs(p3[1] - p0[1])
+            w_div_h = width / height
+            if 5 <= w_div_h <= 25:
+                real_rect_count +=1
+                rect_big_list.append(rect)
+                w_div_h_sum_big.append(w_div_h)
+        
+            if 0.04 <= w_div_h <= 0.2:
+                real_rect_count -=1
+                rect_small_list.append(rect)
+                w_div_h_sum_small.append(w_div_h)
+        if real_rect_count > 0:
+            ret_rect = rect_big_list
+            w_div_h_mean = np.mean(w_div_h_sum_big)
+        else:
+            ret_rect = rect_small_list
+            w_div_h_mean = np.mean(w_div_h_sum_small)
+            
+        if w_div_h_mean >= 1.5:
+            return 1, ret_rect
+        else:
+            return 0, ret_rect
+
+    def crop_image(self, rect, image):
+        p0 = rect[0]
+        p1 = rect[1]
+        p2 = rect[2]
+        p3 = rect[3]
+        crop = image[int(p0[1]):int(p2[1]), int(p0[0]):int(p2[0])]
+        # crop_image = Image.fromarray(crop)
+        return crop
+
+    def get_img_real_angle(self, img):
+        ret_angle = 0
+        image = img
+        # ocr = PaddleOCR(use_angle_cls=True)
+        # angle_cls = ocr.ocr(img_path, det=False, rec=False, cls=True)
+
+        rect_list = self.ocr.ocr(image, rec=False)
+        if rect_list != [[]]:
+            except_flag = False
+            try:
+                real_angle_flag, rect_good = self.get_real_rotation_flag(
+                    rect_list)
+                rect_crop = choice(rect_good)
+                # rect_crop = rect_good[0]
+                image_crop = self.crop_image(rect_crop, image)
+                # ocr_angle = PaddleOCR(use_angle_cls=True)
+                angle_cls = self.ocr_angle.ocr(
+                    image_crop, det=False, rec=False, cls=True)
+            except:
+                except_flag = True
+                real_angle_flag = self.get_real_rotation_when_null_rect(
+                    rect_list)
+                # ocr_angle = PaddleOCR(use_angle_cls=True)
+                angle_cls = self.ocr_angle.ocr(
+                    image, det=False, rec=False, cls=True)
+        else:
+            return 0
+        if angle_cls[0][0] == '0':
+            if real_angle_flag:
+                ret_angle = 0
+            else:
+                ret_angle = 270
+                if not except_flag:
+                    anticlockwise_90 = rotate_bound_white_bg(image_crop, 90)
+                    angle_cls = self.ocr_angle.ocr(anticlockwise_90, det=False, rec=False, cls=True)
+                    if angle_cls[0][0] == '0':
+                        ret_angle = 270
+                    if angle_cls[0][0] == '180':
+                        ret_angle = 90
+        if angle_cls[0][0] == '180':
+            if real_angle_flag:
+                ret_angle = 180
+            else:
+                ret_angle = 90
+        return ret_angle
+def paddle(request: Request,info: ImageListInfo):
+    logger.info(f"->图片数量:{len(info.images)}")
+    res_list = []
+    for b_img in info.images:
+        img = base64_to_np(b_img)
+        route=GetImageRotation()
+        route2=route.get_img_real_angle(img)
+        if route2==90 or route2== 270:
+            img=im.transpose(img.ROTATE_90)
+        result=ocr.ocr(img,cls=True)
+        r_list = []
+        for text_list in result:
+            if len(text_list) >= 1:
+                data = {}
+                data["confidence"]= text_list[1][1]
+                data["text"] = text_list[1][0]
+                data["type"] = info.img_type
+                data["text_region"]= text_list[0]
+                r_list.append(data)
+        res_list.append(r_list)
+    return res_list
+
+
+=======
 def table(image: TableInfo):
     img = base64_to_np(image.image)
     res = table_engine(img)
     return res[0]['res']
+>>>>>>> 658d5ab26bfc557357332fa8aea4b547f20c5ff2
 
 
 if __name__ == '__main__':