优化图片分割和拓展方法，兼容横向过长的图片

photo_mask/photo_mask.py

@@ -57,7 +57,7 @@ def find_boxes(content, layout, offset=0, length=None, improve=False, image_path
         capture_box = util.zoom_rectangle(box, 0.2)
         captured_image = image_util.capture(image, capture_box)
         with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
-            captured_image, offset_x, offset_y = image_util.expand_to_a4_size(captured_image, True)
+            captured_image, offset_x, offset_y = image_util.expand_to_a4_size(captured_image)
             cv2.imwrite(temp_file.name, captured_image)
         try:
             layouts = util.get_ocr_layout(OCR, temp_file.name)
@@ -166,7 +166,7 @@ def get_mask_layout(image, name, id_card_num):
 
 
 def handle_image_for_mask(split_result):
-    expand_img, offset_x, offset_y = image_util.expand_to_a4_size(split_result["img"], True)
+    expand_img, offset_x, offset_y = image_util.expand_to_a4_size(split_result["img"])
     split_result["x_offset"] -= offset_x
     split_result["y_offset"] -= offset_y
     return expand_img, split_result["x_offset"], split_result["y_offset"]

util/image_util.py

@@ -41,31 +41,36 @@ def capture(image, rectangle):
     return image[int(y1):int(y2), int(x1):int(x2)]
 
 
-def split(image, ratio=1.414, overlap=0.05):
+def split(image, ratio=1.414, overlap=0.05, x_compensation=3):
     """
-    分割图片，只分割过长的图片，暂不处理过宽的图片
+    分割图片
     :param image:图片，可以是NumPy数组或文件路径
     :param ratio: 分割后的比例
     :param overlap: 图片之间的覆盖比例
+    :param x_compensation: 横向补偿倍率
     :return: 分割后的图片组(NumPy数组形式)
     """
     split_result = []
     if isinstance(image, str):
         image = read(image)
-    # 获取图片的宽度和高度
     height, width = image.shape[:2]
-    # 计算宽高比
-    img_ratio = height / width
-    # 检查是否需要裁剪
-    if img_ratio > ratio:
-        split_ratio = ratio - overlap
-        # 分割后的高度
+    hw_ratio = height / width
+    wh_ratio = width / height
+
+    if hw_ratio > ratio:  # 纵向过长
         new_img_height = width * ratio
-        for i in range(math.ceil(height / (width * split_ratio))):
-            offset = round(width * split_ratio * i)
-            # 参数形式为[y1:y2, x1:x2]
+        step = width * (ratio - overlap)  # 偏移步长
+        for i in range(math.ceil(height / step)):
+            offset = round(step * i)
             cropped_img = capture(image, [0, offset, width, offset + new_img_height])
             split_result.append({"img": cropped_img, "x_offset": 0, "y_offset": offset})
+    elif wh_ratio > ratio:  # 横向过长
+        new_img_width = height * ratio
+        step = height * (ratio - overlap * x_compensation)  # 一般文字是横向的，所以横向截取时增大重叠部分
+        for i in range(math.ceil(width / step)):
+            offset = round(step * i)
+            cropped_img = capture(image, [offset, 0, offset + new_img_width, width])
+            split_result.append({"img": cropped_img, "x_offset": offset, "y_offset": 0})
     else:
         split_result.append({"img": image, "x_offset": 0, "y_offset": 0})
     return split_result
@@ -171,38 +176,40 @@ def invert_rotate_rectangle(rectangle, center, angle):
     return [new_top_left[0], new_top_left[1], new_bot_right[0], new_bot_right[1]]
 
 
-def expand_to_a4_size(image, center=False):
+def expand_to_a4_size(image):
     """
-    将图片扩充到a4大小
+    以尽量少的方式将图片扩充到a4大小
     :param image: 图片NumPy数组
-    :param center: 是否将原图置于中间
     :return: 扩充后的图片NumPy数组和偏移量
     """
-    h, w = image.shape[:2]
-    offset_x, offset_y = 0, 0
-    if h * 1.0 / w >= 1.42:
-        exp_w = int(h / 1.414 - w)
-        if center:
-            offset_x = int(exp_w / 2)
-            exp_img = numpy.zeros((h, offset_x, 3), dtype="uint8")
+    height, width = image.shape[:2]
+    x_offset, y_offset = 0, 0
+    hw_ratio = height / width
+    if hw_ratio >= 1.42:
+        exp_w = int(height / 1.414 - width)
+        x_offset = int(exp_w / 2)
+        exp_img = numpy.zeros((height, x_offset, 3), dtype="uint8")
         exp_img.fill(255)
         image = numpy.hstack([exp_img, image, exp_img])
-        else:
-            exp_img = numpy.zeros((h, exp_w, 3), dtype="uint8")
-            exp_img.fill(255)
-            image = numpy.hstack([image, exp_img])
-    elif h * 1.0 / w <= 1.40:
-        exp_h = int(w * 1.414 - h)
-        if center:
-            offset_y = int(exp_h / 2)
-            exp_img = numpy.zeros((offset_y, w, 3), dtype="uint8")
+    elif 1 <= hw_ratio <= 1.40:
+        exp_h = int(width * 1.414 - height)
+        y_offset = int(exp_h / 2)
+        exp_img = numpy.zeros((y_offset, width, 3), dtype="uint8")
         exp_img.fill(255)
         image = numpy.vstack([exp_img, image, exp_img])
-        else:
-            exp_img = numpy.zeros((exp_h, w, 3), dtype="uint8")
+    elif 0.72 <= hw_ratio < 1:
+        exp_w = int(height * 1.414 - width)
+        x_offset = int(exp_w / 2)
+        exp_img = numpy.zeros((height, x_offset, 3), dtype="uint8")
         exp_img.fill(255)
-            image = numpy.vstack([image, exp_img])
-    return image, offset_x, offset_y
+        image = numpy.hstack([exp_img, image, exp_img])
+    elif hw_ratio <= 0.7:
+        exp_h = int(width / 1.414 - height)
+        y_offset = int(exp_h / 2)
+        exp_img = numpy.zeros((y_offset, width, 3), dtype="uint8")
+        exp_img.fill(255)
+        image = numpy.vstack([exp_img, image, exp_img])
+    return image, x_offset, y_offset
 
 
 def combined(img1, img2):

visual_model_test/visual_model_test.py

@@ -56,7 +56,7 @@ def visual_model_test(model_type, test_img, task_path, schema):
                 # angle = image_util.parse_rotation_angles(img["img"])[0]
                 angle = 0
                 rotated_img = image_util.rotate(img["img"], angle)
-                rotated_img, offset_x, offset_y = image_util.expand_to_a4_size(rotated_img, True)
+                rotated_img, offset_x, offset_y = image_util.expand_to_a4_size(rotated_img)
                 cv2.imwrite(temp_file.name, rotated_img)
 
                 img["x_offset"] -= offset_x