coordinate_utils.py 9.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308
  1. """
  2. 通用坐标转换工具模块
  3. 提供通用的坐标计算和转换功能:
  4. - 底层坐标计算(IoU、重叠比例)
  5. - 多边形/bbox 格式转换
  6. - 相对坐标 → 绝对坐标转换
  7. - 图像裁剪
  8. - 格式检测
  9. 此模块从 universal_doc_parser 中提取,供多个模块共享使用。
  10. """
  11. from typing import List, Tuple, Union, Any
  12. import numpy as np
  13. # 导入 MinerU 组件(用于 IoU 计算优化)
  14. try:
  15. from mineru.utils.boxbase import calculate_iou as mineru_calculate_iou
  16. from mineru.utils.boxbase import calculate_overlap_area_2_minbox_area_ratio
  17. MINERU_BOXBASE_AVAILABLE = True
  18. except ImportError:
  19. MINERU_BOXBASE_AVAILABLE = False
  20. mineru_calculate_iou = None
  21. calculate_overlap_area_2_minbox_area_ratio = None
  22. class CoordinateUtils:
  23. """通用坐标转换工具类"""
  24. # ==================== 底层坐标计算方法 ====================
  25. @staticmethod
  26. def calculate_iou(bbox1: List[float], bbox2: List[float]) -> float:
  27. """
  28. 计算两个 bbox 的 IoU(交并比)
  29. Args:
  30. bbox1: 第一个 bbox [x1, y1, x2, y2]
  31. bbox2: 第二个 bbox [x1, y1, x2, y2]
  32. Returns:
  33. IoU 值
  34. """
  35. if MINERU_BOXBASE_AVAILABLE and mineru_calculate_iou is not None:
  36. return mineru_calculate_iou(bbox1, bbox2)
  37. # 备用实现
  38. x_left = max(bbox1[0], bbox2[0])
  39. y_top = max(bbox1[1], bbox2[1])
  40. x_right = min(bbox1[2], bbox2[2])
  41. y_bottom = min(bbox1[3], bbox2[3])
  42. if x_right < x_left or y_bottom < y_top:
  43. return 0.0
  44. intersection_area = (x_right - x_left) * (y_bottom - y_top)
  45. bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
  46. bbox2_area = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
  47. if bbox1_area == 0 or bbox2_area == 0:
  48. return 0.0
  49. return intersection_area / float(bbox1_area + bbox2_area - intersection_area)
  50. @staticmethod
  51. def calculate_overlap_ratio(bbox1: List[float], bbox2: List[float]) -> float:
  52. """
  53. 计算重叠面积占小框面积的比例
  54. Args:
  55. bbox1: 第一个 bbox [x1, y1, x2, y2]
  56. bbox2: 第二个 bbox [x1, y1, x2, y2]
  57. Returns:
  58. 重叠比例
  59. """
  60. if MINERU_BOXBASE_AVAILABLE and calculate_overlap_area_2_minbox_area_ratio is not None:
  61. return calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2)
  62. # 备用实现
  63. x_left = max(bbox1[0], bbox2[0])
  64. y_top = max(bbox1[1], bbox2[1])
  65. x_right = min(bbox1[2], bbox2[2])
  66. y_bottom = min(bbox1[3], bbox2[3])
  67. if x_right < x_left or y_bottom < y_top:
  68. return 0.0
  69. intersection_area = (x_right - x_left) * (y_bottom - y_top)
  70. area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
  71. area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
  72. min_area = min(area1, area2)
  73. if min_area == 0:
  74. return 0.0
  75. return intersection_area / min_area
  76. @staticmethod
  77. def calculate_overlap_in_bbox1_ratio(
  78. bbox1: List[float],
  79. bbox2: List[float]
  80. ) -> float:
  81. """
  82. 计算 bbox1 被 bbox2 覆盖的面积比例
  83. Args:
  84. bbox1: 第一个 bbox [x1, y1, x2, y2]
  85. bbox2: 第二个 bbox [x1, y1, x2, y2]
  86. Returns:
  87. bbox1 被覆盖的比例
  88. """
  89. x_left = max(bbox1[0], bbox2[0])
  90. y_top = max(bbox1[1], bbox2[1])
  91. x_right = min(bbox1[2], bbox2[2])
  92. y_bottom = min(bbox1[3], bbox2[3])
  93. if x_right < x_left or y_bottom < y_top:
  94. return 0.0
  95. intersection_area = (x_right - x_left) * (y_bottom - y_top)
  96. bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
  97. if bbox1_area == 0:
  98. return 0.0
  99. return intersection_area / bbox1_area
  100. @staticmethod
  101. def poly_to_bbox(poly: Union[List, None]) -> List[float]:
  102. """
  103. 将多边形坐标转换为 bbox 格式
  104. Args:
  105. poly: 多边形坐标,支持以下格式:
  106. - [[x1,y1], [x2,y1], [x2,y2], [x1,y2]] (4个点)
  107. - [x1, y1, x2, y1, x2, y2, x1, y2] (8个值)
  108. - [x1, y1, x2, y2] (4个值,已是bbox)
  109. Returns:
  110. bbox [x1, y1, x2, y2]
  111. """
  112. if not poly:
  113. return [0, 0, 0, 0]
  114. # 处理嵌套列表格式 [[x1,y1], [x2,y1], ...]
  115. if isinstance(poly[0], (list, tuple)):
  116. xs = [p[0] for p in poly]
  117. ys = [p[1] for p in poly]
  118. return [min(xs), min(ys), max(xs), max(ys)]
  119. # 处理平面列表格式
  120. if len(poly) == 4:
  121. # 已经是 bbox 格式
  122. return list(poly)
  123. elif len(poly) >= 8:
  124. # 8点格式:[x1, y1, x2, y1, x2, y2, x1, y2]
  125. xs = [poly[i] for i in range(0, len(poly), 2)]
  126. ys = [poly[i] for i in range(1, len(poly), 2)]
  127. return [min(xs), min(ys), max(xs), max(ys)]
  128. return [0, 0, 0, 0]
  129. @staticmethod
  130. def bbox_to_poly(bbox: List[float]) -> List[List[float]]:
  131. """
  132. 将 bbox 转换为多边形坐标
  133. Args:
  134. bbox: [x1, y1, x2, y2]
  135. Returns:
  136. [[x1,y1], [x2,y1], [x2,y2], [x1,y2]]
  137. """
  138. if not bbox or len(bbox) < 4:
  139. return [[0, 0], [0, 0], [0, 0], [0, 0]]
  140. x1, y1, x2, y2 = bbox[:4]
  141. return [
  142. [float(x1), float(y1)],
  143. [float(x2), float(y1)],
  144. [float(x2), float(y2)],
  145. [float(x1), float(y2)]
  146. ]
  147. # ==================== 图像裁剪 ====================
  148. @staticmethod
  149. def crop_region(image: np.ndarray, bbox: List[float], padding: int = 0) -> np.ndarray:
  150. """
  151. 裁剪图像区域
  152. Args:
  153. image: 原始图像
  154. bbox: 裁剪区域 [x1, y1, x2, y2]
  155. padding: 边缘padding(像素),可以为正数(扩展裁剪区域)或负数(收缩裁剪区域)
  156. Returns:
  157. 裁剪后的图像
  158. """
  159. if len(bbox) < 4:
  160. return image
  161. h, w = image.shape[:2]
  162. # 解析padding(支持单个值或四个值)
  163. if isinstance(padding, (int, float)):
  164. pad_left = pad_right = pad_top = pad_bottom = int(padding)
  165. else:
  166. # 假设是长度为4的元组/列表 [left, top, right, bottom]
  167. if len(padding) >= 4:
  168. pad_left, pad_top, pad_right, pad_bottom = [int(p) for p in padding[:4]]
  169. else:
  170. pad_left = pad_top = pad_right = pad_bottom = 0
  171. x1 = max(0 - pad_left, int(bbox[0]) - pad_left)
  172. y1 = max(0 - pad_top, int(bbox[1]) - pad_top)
  173. x2 = min(w + pad_right, int(bbox[2]) + pad_right)
  174. y2 = min(h + pad_bottom, int(bbox[3]) + pad_bottom)
  175. # 确保坐标有效
  176. x1 = max(0, x1)
  177. y1 = max(0, y1)
  178. x2 = min(w, x2)
  179. y2 = min(h, y2)
  180. # 检查是否有效区域
  181. if x2 <= x1 or y2 <= y1:
  182. return image
  183. return image[y1:y2, x1:x2]
  184. @staticmethod
  185. def bbox_overlap(bbox1: List[float], bbox2: List[float]) -> bool:
  186. """
  187. 检查两个 bbox 是否重叠
  188. Args:
  189. bbox1: 第一个 bbox [x1, y1, x2, y2]
  190. bbox2: 第二个 bbox [x1, y1, x2, y2]
  191. Returns:
  192. 是否重叠
  193. """
  194. if len(bbox1) < 4 or len(bbox2) < 4:
  195. return False
  196. x1_1, y1_1, x2_1, y2_1 = bbox1[:4]
  197. x1_2, y1_2, x2_2, y2_2 = bbox2[:4]
  198. if x2_1 < x1_2 or x2_2 < x1_1:
  199. return False
  200. if y2_1 < y1_2 or y2_2 < y1_1:
  201. return False
  202. return True
  203. @staticmethod
  204. def convert_to_absolute_coords(
  205. relative_bbox: List,
  206. region_bbox: List[float]
  207. ) -> List:
  208. """
  209. 将相对坐标转换为绝对坐标
  210. Args:
  211. relative_bbox: 相对坐标
  212. region_bbox: 区域的绝对坐标 [x1, y1, x2, y2]
  213. Returns:
  214. 绝对坐标
  215. """
  216. if not relative_bbox or len(region_bbox) < 4:
  217. return relative_bbox
  218. bx1, by1 = region_bbox[0], region_bbox[1]
  219. # 处理4点坐标格式 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
  220. if isinstance(relative_bbox[0], (list, tuple)):
  221. return [
  222. [p[0] + bx1, p[1] + by1] for p in relative_bbox
  223. ]
  224. # 处理4值坐标格式 [x1, y1, x2, y2]
  225. if len(relative_bbox) >= 4:
  226. return [
  227. relative_bbox[0] + bx1,
  228. relative_bbox[1] + by1,
  229. relative_bbox[2] + bx1,
  230. relative_bbox[3] + by1
  231. ]
  232. return relative_bbox
  233. @staticmethod
  234. def is_poly_format(bbox: Any) -> bool:
  235. """
  236. 检测 bbox 是否为四点多边形格式
  237. 四点格式: [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
  238. 矩形格式: [x_min, y_min, x_max, y_max]
  239. """
  240. if not bbox or not isinstance(bbox, list):
  241. return False
  242. if len(bbox) != 4:
  243. return False
  244. return isinstance(bbox[0], (list, tuple))