当前位置：首页 > news >正文

【深度学习】OCR中的Shrink操作详解

news 来源：原创 2024/9/20 13:27:16

OCR中的Shrink操作详解

在光学字符识别（OCR）中，shrink操作用于对文本框多边形进行缩放，以生成用于训练和检测的特征图。本文将介绍shrink操作的背景、实现方法及其应用。以下是用户提供的代码，详细展示了如何实现这一过程。

背景介绍

在OCR任务中，文本通常以多边形的形式标注于图像中。为了更好地训练检测模型，通常需要将这些多边形进行一定比例的缩放（shrink），以生成不同大小的特征图，从而提高模型的泛化能力和精度。shrink操作的目标是将文本框缩小，以减少噪声对检测结果的影响。

代码实现

以下是实现shrink操作的详细代码：

import numpy as np
import cv2
import pyclipper
from shapely.geometry import Polygondef shrink_polygon_py(polygon, shrink_ratio):"""对框进行缩放，返回去的比例为1/shrink_ratio 即可"""cx = polygon[:, 0].mean()cy = polygon[:, 1].mean()polygon[:, 0] = cx + (polygon[:, 0] - cx) * shrink_ratiopolygon[:, 1] = cy + (polygon[:, 1] - cy) * shrink_ratioreturn polygondef shrink_polygon_pyclipper(polygon, shrink_ratio):polygon_shape = Polygon(polygon)distance = (polygon_shape.area * (1 - np.power(shrink_ratio, 2)) / polygon_shape.length)subject = [tuple(l) for l in polygon]padding = pyclipper.PyclipperOffset()padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)shrinked = padding.Execute(-distance)if shrinked == []:shrinked = np.array(shrinked)else:shrinked = np.array(shrinked[0]).reshape(-1, 2)return shrinkedclass MakeShrinkMap:def __init__(self, min_text_size=8, shrink_ratio=0.4, shrink_type="pyclipper"):shrink_func_dict = {"py": shrink_polygon_py,"pyclipper": shrink_polygon_pyclipper,}self.shrink_func = shrink_func_dict[shrink_type]self.min_text_size = min_text_sizeself.shrink_ratio = shrink_ratiodef __call__(self, data: dict) -> dict:image = data["img"]text_polys = data["text_polys"]ignore_tags = data["ignore_tags"]h, w = image.shape[:2]text_polys, ignore_tags = self.validate_polygons(text_polys, ignore_tags, h, w)gt = np.zeros((h, w), dtype=np.float32)mask = np.ones((h, w), dtype=np.float32)shrinked_polygons = []for i in range(len(text_polys)):polygon = text_polys[i]height = max(polygon[:, 1]) - min(polygon[:, 1])width = max(polygon[:, 0]) - min(polygon[:, 0])if ignore_tags[i] or min(height, width) < self.min_text_size:cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)ignore_tags[i] = Trueelse:shrinked = self.shrink_func(polygon, self.shrink_ratio)shrinked_polygons.append(shrinked)if shrinked.size == 0:cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)ignore_tags[i] = Truecontinuecv2.fillPoly(gt, [shrinked.astype(np.int32)], 1)data["shrink_map"] = gtdata["shrink_mask"] = maskdata["shrinked_polygons"] = shrinked_polygonsreturn datadef validate_polygons(self, polygons, ignore_tags, h, w):if len(polygons) == 0:return polygons, ignore_tagsassert len(polygons) == len(ignore_tags)for polygon in polygons:polygon[:, 0] = np.clip(polygon[:, 0], 0, w - 1)polygon[:, 1] = np.clip(polygon[:, 1], 0, h - 1)for i in range(len(polygons)):area = self.polygon_area(polygons[i])if abs(area) < 1:ignore_tags[i] = Trueif area > 0:polygons[i] = polygons[i][::-1, :]return polygons, ignore_tagsdef polygon_area(self, polygon):return cv2.contourArea(polygon)if __name__ == "__main__":# 示例图像image = np.ones((200, 200, 3), dtype=np.uint8) * 255# 示例文本框多边形text_polys = [np.array([[50, 50], [150, 50], [150, 100], [50, 100]]),np.array([[60, 120], [140, 120], [140, 160], [60, 160]])]# 示例忽略标志ignore_tags = [False, False]# 构建输入数据字典data = {"img": image,"text_polys": text_polys,"ignore_tags": ignore_tags}# 初始化 MakeShrinkMap 类make_shrink_map = MakeShrinkMap(min_text_size=8, shrink_ratio=0.4, shrink_type="pyclipper")# 调用类处理数据result = make_shrink_map(data)# 获取生成的shrink_map和shrink_maskshrink_map = result["shrink_map"]shrink_mask = result["shrink_mask"]shrinked_polygons = result["shrinked_polygons"]# 在原图上绘制shrink前的多边形original_image = image.copy()for polygon in text_polys:cv2.polylines(original_image, [polygon.astype(np.int32)], True, (0, 0, 255), 2)# 在原图上绘制shrink后的多边形shrinked_image = image.copy()for polygon in shrinked_polygons:cv2.polylines(shrinked_image, [polygon.astype(np.int32)], True, (0, 255, 0), 2)# 保存结果图像cv2.imwrite("original_image.png", original_image)cv2.imwrite("shrinked_image.png", shrinked_image)cv2.imwrite("shrink_map.png", shrink_map * 255)  # 将shrink_map转换为图像cv2.imwrite("shrink_mask.png", shrink_mask * 255)  # 将shrink_mask转换为图像# 显示结果# cv2.imshow("Original Image", original_image)# cv2.imshow("Shrinked Image", shrinked_image)# cv2.imshow("Shrink Map", shrink_map)# cv2.imshow("Shrink Mask", shrink_mask)# cv2.waitKey(0)# cv2.destroyAllWindows()

代码详解

Shrink算法实现

代码中实现了两种不同的shrink算法：shrink_polygon_py和shrink_polygon_pyclipper。
- shrink_polygon_py：通过计算多边形的中心点，然后将多边形的每个点按照缩放比例向中心点收缩。
- shrink_polygon_pyclipper：使用pyclipper库进行多边形缩放，计算更为精确，适用于复杂多边形。
MakeShrinkMap类

MakeShrinkMap类用于将图像中的文本多边形进行shrink操作。类的构造函数接受最小文本尺寸、缩放比例和缩放类型作为参数。__call__方法处理输入数据字典，并生成缩放后的特征图和掩码。
代码示例

在__main__部分，创建了一个示例图像和文本多边形，并使用MakeShrinkMap类进行shrink操作。结果图像包括原始多边形和缩放后的多边形，并将生成的特征图和掩码保存为图像文件。