box_utils.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430
  1. # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import numpy as np
  15. import random
  16. import math
  17. import cv2
  18. import scipy
  19. def meet_emit_constraint(src_bbox, sample_bbox):
  20. center_x = (src_bbox[2] + src_bbox[0]) / 2
  21. center_y = (src_bbox[3] + src_bbox[1]) / 2
  22. if center_x >= sample_bbox[0] and \
  23. center_x <= sample_bbox[2] and \
  24. center_y >= sample_bbox[1] and \
  25. center_y <= sample_bbox[3]:
  26. return True
  27. return False
  28. def clip_bbox(src_bbox):
  29. src_bbox[0] = max(min(src_bbox[0], 1.0), 0.0)
  30. src_bbox[1] = max(min(src_bbox[1], 1.0), 0.0)
  31. src_bbox[2] = max(min(src_bbox[2], 1.0), 0.0)
  32. src_bbox[3] = max(min(src_bbox[3], 1.0), 0.0)
  33. return src_bbox
  34. def bbox_area(src_bbox):
  35. if src_bbox[2] < src_bbox[0] or src_bbox[3] < src_bbox[1]:
  36. return 0.
  37. else:
  38. width = src_bbox[2] - src_bbox[0]
  39. height = src_bbox[3] - src_bbox[1]
  40. return width * height
  41. def is_overlap(object_bbox, sample_bbox):
  42. if object_bbox[0] >= sample_bbox[2] or \
  43. object_bbox[2] <= sample_bbox[0] or \
  44. object_bbox[1] >= sample_bbox[3] or \
  45. object_bbox[3] <= sample_bbox[1]:
  46. return False
  47. else:
  48. return True
  49. def filter_and_process(sample_bbox, bboxes, labels, scores=None):
  50. new_bboxes = []
  51. new_labels = []
  52. new_scores = []
  53. for i in range(len(bboxes)):
  54. new_bbox = [0, 0, 0, 0]
  55. obj_bbox = [bboxes[i][0], bboxes[i][1], bboxes[i][2], bboxes[i][3]]
  56. if not meet_emit_constraint(obj_bbox, sample_bbox):
  57. continue
  58. if not is_overlap(obj_bbox, sample_bbox):
  59. continue
  60. sample_width = sample_bbox[2] - sample_bbox[0]
  61. sample_height = sample_bbox[3] - sample_bbox[1]
  62. new_bbox[0] = (obj_bbox[0] - sample_bbox[0]) / sample_width
  63. new_bbox[1] = (obj_bbox[1] - sample_bbox[1]) / sample_height
  64. new_bbox[2] = (obj_bbox[2] - sample_bbox[0]) / sample_width
  65. new_bbox[3] = (obj_bbox[3] - sample_bbox[1]) / sample_height
  66. new_bbox = clip_bbox(new_bbox)
  67. if bbox_area(new_bbox) > 0:
  68. new_bboxes.append(new_bbox)
  69. new_labels.append([labels[i][0]])
  70. if scores is not None:
  71. new_scores.append([scores[i][0]])
  72. bboxes = np.array(new_bboxes)
  73. labels = np.array(new_labels)
  74. scores = np.array(new_scores)
  75. return bboxes, labels, scores
  76. def bbox_area_sampling(bboxes, labels, scores, target_size, min_size):
  77. new_bboxes = []
  78. new_labels = []
  79. new_scores = []
  80. for i, bbox in enumerate(bboxes):
  81. w = float((bbox[2] - bbox[0]) * target_size)
  82. h = float((bbox[3] - bbox[1]) * target_size)
  83. if w * h < float(min_size * min_size):
  84. continue
  85. else:
  86. new_bboxes.append(bbox)
  87. new_labels.append(labels[i])
  88. if scores is not None and scores.size != 0:
  89. new_scores.append(scores[i])
  90. bboxes = np.array(new_bboxes)
  91. labels = np.array(new_labels)
  92. scores = np.array(new_scores)
  93. return bboxes, labels, scores
  94. def generate_sample_bbox(sampler):
  95. scale = np.random.uniform(sampler[2], sampler[3])
  96. aspect_ratio = np.random.uniform(sampler[4], sampler[5])
  97. aspect_ratio = max(aspect_ratio, (scale**2.0))
  98. aspect_ratio = min(aspect_ratio, 1 / (scale**2.0))
  99. bbox_width = scale * (aspect_ratio**0.5)
  100. bbox_height = scale / (aspect_ratio**0.5)
  101. xmin_bound = 1 - bbox_width
  102. ymin_bound = 1 - bbox_height
  103. xmin = np.random.uniform(0, xmin_bound)
  104. ymin = np.random.uniform(0, ymin_bound)
  105. xmax = xmin + bbox_width
  106. ymax = ymin + bbox_height
  107. sampled_bbox = [xmin, ymin, xmax, ymax]
  108. return sampled_bbox
  109. def generate_sample_bbox_square(sampler, image_width, image_height):
  110. scale = np.random.uniform(sampler[2], sampler[3])
  111. aspect_ratio = np.random.uniform(sampler[4], sampler[5])
  112. aspect_ratio = max(aspect_ratio, (scale**2.0))
  113. aspect_ratio = min(aspect_ratio, 1 / (scale**2.0))
  114. bbox_width = scale * (aspect_ratio**0.5)
  115. bbox_height = scale / (aspect_ratio**0.5)
  116. if image_height < image_width:
  117. bbox_width = bbox_height * image_height / image_width
  118. else:
  119. bbox_height = bbox_width * image_width / image_height
  120. xmin_bound = 1 - bbox_width
  121. ymin_bound = 1 - bbox_height
  122. xmin = np.random.uniform(0, xmin_bound)
  123. ymin = np.random.uniform(0, ymin_bound)
  124. xmax = xmin + bbox_width
  125. ymax = ymin + bbox_height
  126. sampled_bbox = [xmin, ymin, xmax, ymax]
  127. return sampled_bbox
  128. def data_anchor_sampling(bbox_labels, image_width, image_height, scale_array,
  129. resize_width):
  130. num_gt = len(bbox_labels)
  131. # np.random.randint range: [low, high)
  132. rand_idx = np.random.randint(0, num_gt) if num_gt != 0 else 0
  133. if num_gt != 0:
  134. norm_xmin = bbox_labels[rand_idx][0]
  135. norm_ymin = bbox_labels[rand_idx][1]
  136. norm_xmax = bbox_labels[rand_idx][2]
  137. norm_ymax = bbox_labels[rand_idx][3]
  138. xmin = norm_xmin * image_width
  139. ymin = norm_ymin * image_height
  140. wid = image_width * (norm_xmax - norm_xmin)
  141. hei = image_height * (norm_ymax - norm_ymin)
  142. range_size = 0
  143. area = wid * hei
  144. for scale_ind in range(0, len(scale_array) - 1):
  145. if area > scale_array[scale_ind] ** 2 and area < \
  146. scale_array[scale_ind + 1] ** 2:
  147. range_size = scale_ind + 1
  148. break
  149. if area > scale_array[len(scale_array) - 2]**2:
  150. range_size = len(scale_array) - 2
  151. scale_choose = 0.0
  152. if range_size == 0:
  153. rand_idx_size = 0
  154. else:
  155. # np.random.randint range: [low, high)
  156. rng_rand_size = np.random.randint(0, range_size + 1)
  157. rand_idx_size = rng_rand_size % (range_size + 1)
  158. if rand_idx_size == range_size:
  159. min_resize_val = scale_array[rand_idx_size] / 2.0
  160. max_resize_val = min(2.0 * scale_array[rand_idx_size],
  161. 2 * math.sqrt(wid * hei))
  162. scale_choose = random.uniform(min_resize_val, max_resize_val)
  163. else:
  164. min_resize_val = scale_array[rand_idx_size] / 2.0
  165. max_resize_val = 2.0 * scale_array[rand_idx_size]
  166. scale_choose = random.uniform(min_resize_val, max_resize_val)
  167. sample_bbox_size = wid * resize_width / scale_choose
  168. w_off_orig = 0.0
  169. h_off_orig = 0.0
  170. if sample_bbox_size < max(image_height, image_width):
  171. if wid <= sample_bbox_size:
  172. w_off_orig = np.random.uniform(xmin + wid - sample_bbox_size,
  173. xmin)
  174. else:
  175. w_off_orig = np.random.uniform(xmin,
  176. xmin + wid - sample_bbox_size)
  177. if hei <= sample_bbox_size:
  178. h_off_orig = np.random.uniform(ymin + hei - sample_bbox_size,
  179. ymin)
  180. else:
  181. h_off_orig = np.random.uniform(ymin,
  182. ymin + hei - sample_bbox_size)
  183. else:
  184. w_off_orig = np.random.uniform(image_width - sample_bbox_size, 0.0)
  185. h_off_orig = np.random.uniform(image_height - sample_bbox_size,
  186. 0.0)
  187. w_off_orig = math.floor(w_off_orig)
  188. h_off_orig = math.floor(h_off_orig)
  189. # Figure out top left coordinates.
  190. w_off = float(w_off_orig / image_width)
  191. h_off = float(h_off_orig / image_height)
  192. sampled_bbox = [
  193. w_off, h_off, w_off + float(sample_bbox_size / image_width),
  194. h_off + float(sample_bbox_size / image_height)
  195. ]
  196. return sampled_bbox
  197. else:
  198. return 0
  199. def jaccard_overlap(sample_bbox, object_bbox):
  200. if sample_bbox[0] >= object_bbox[2] or \
  201. sample_bbox[2] <= object_bbox[0] or \
  202. sample_bbox[1] >= object_bbox[3] or \
  203. sample_bbox[3] <= object_bbox[1]:
  204. return 0
  205. intersect_xmin = max(sample_bbox[0], object_bbox[0])
  206. intersect_ymin = max(sample_bbox[1], object_bbox[1])
  207. intersect_xmax = min(sample_bbox[2], object_bbox[2])
  208. intersect_ymax = min(sample_bbox[3], object_bbox[3])
  209. intersect_size = (intersect_xmax - intersect_xmin) * (
  210. intersect_ymax - intersect_ymin)
  211. sample_bbox_size = bbox_area(sample_bbox)
  212. object_bbox_size = bbox_area(object_bbox)
  213. overlap = intersect_size / (
  214. sample_bbox_size + object_bbox_size - intersect_size)
  215. return overlap
  216. def intersect_bbox(bbox1, bbox2):
  217. if bbox2[0] > bbox1[2] or bbox2[2] < bbox1[0] or \
  218. bbox2[1] > bbox1[3] or bbox2[3] < bbox1[1]:
  219. intersection_box = [0.0, 0.0, 0.0, 0.0]
  220. else:
  221. intersection_box = [
  222. max(bbox1[0], bbox2[0]),
  223. max(bbox1[1], bbox2[1]),
  224. min(bbox1[2], bbox2[2]),
  225. min(bbox1[3], bbox2[3])
  226. ]
  227. return intersection_box
  228. def bbox_coverage(bbox1, bbox2):
  229. inter_box = intersect_bbox(bbox1, bbox2)
  230. intersect_size = bbox_area(inter_box)
  231. if intersect_size > 0:
  232. bbox1_size = bbox_area(bbox1)
  233. return intersect_size / bbox1_size
  234. else:
  235. return 0.
  236. def satisfy_sample_constraint(sampler,
  237. sample_bbox,
  238. gt_bboxes,
  239. satisfy_all=False):
  240. if sampler[6] == 0 and sampler[7] == 0:
  241. return True
  242. satisfied = []
  243. for i in range(len(gt_bboxes)):
  244. object_bbox = [
  245. gt_bboxes[i][0], gt_bboxes[i][1], gt_bboxes[i][2], gt_bboxes[i][3]
  246. ]
  247. overlap = jaccard_overlap(sample_bbox, object_bbox)
  248. if sampler[6] != 0 and \
  249. overlap < sampler[6]:
  250. satisfied.append(False)
  251. continue
  252. if sampler[7] != 0 and \
  253. overlap > sampler[7]:
  254. satisfied.append(False)
  255. continue
  256. satisfied.append(True)
  257. if not satisfy_all:
  258. return True
  259. if satisfy_all:
  260. return np.all(satisfied)
  261. else:
  262. return False
  263. def satisfy_sample_constraint_coverage(sampler, sample_bbox, gt_bboxes):
  264. if sampler[6] == 0 and sampler[7] == 0:
  265. has_jaccard_overlap = False
  266. else:
  267. has_jaccard_overlap = True
  268. if sampler[8] == 0 and sampler[9] == 0:
  269. has_object_coverage = False
  270. else:
  271. has_object_coverage = True
  272. if not has_jaccard_overlap and not has_object_coverage:
  273. return True
  274. found = False
  275. for i in range(len(gt_bboxes)):
  276. object_bbox = [
  277. gt_bboxes[i][0], gt_bboxes[i][1], gt_bboxes[i][2], gt_bboxes[i][3]
  278. ]
  279. if has_jaccard_overlap:
  280. overlap = jaccard_overlap(sample_bbox, object_bbox)
  281. if sampler[6] != 0 and \
  282. overlap < sampler[6]:
  283. continue
  284. if sampler[7] != 0 and \
  285. overlap > sampler[7]:
  286. continue
  287. found = True
  288. if has_object_coverage:
  289. object_coverage = bbox_coverage(object_bbox, sample_bbox)
  290. if sampler[8] != 0 and \
  291. object_coverage < sampler[8]:
  292. continue
  293. if sampler[9] != 0 and \
  294. object_coverage > sampler[9]:
  295. continue
  296. found = True
  297. if found:
  298. return True
  299. return found
  300. def crop_image_sampling(img, sample_bbox, image_width, image_height,
  301. target_size):
  302. # no clipping here
  303. xmin = int(sample_bbox[0] * image_width)
  304. xmax = int(sample_bbox[2] * image_width)
  305. ymin = int(sample_bbox[1] * image_height)
  306. ymax = int(sample_bbox[3] * image_height)
  307. w_off = xmin
  308. h_off = ymin
  309. width = xmax - xmin
  310. height = ymax - ymin
  311. cross_xmin = max(0.0, float(w_off))
  312. cross_ymin = max(0.0, float(h_off))
  313. cross_xmax = min(float(w_off + width - 1.0), float(image_width))
  314. cross_ymax = min(float(h_off + height - 1.0), float(image_height))
  315. cross_width = cross_xmax - cross_xmin
  316. cross_height = cross_ymax - cross_ymin
  317. roi_xmin = 0 if w_off >= 0 else abs(w_off)
  318. roi_ymin = 0 if h_off >= 0 else abs(h_off)
  319. roi_width = cross_width
  320. roi_height = cross_height
  321. roi_y1 = int(roi_ymin)
  322. roi_y2 = int(roi_ymin + roi_height)
  323. roi_x1 = int(roi_xmin)
  324. roi_x2 = int(roi_xmin + roi_width)
  325. cross_y1 = int(cross_ymin)
  326. cross_y2 = int(cross_ymin + cross_height)
  327. cross_x1 = int(cross_xmin)
  328. cross_x2 = int(cross_xmin + cross_width)
  329. sample_img = np.zeros((height, width, 3))
  330. sample_img[roi_y1: roi_y2, roi_x1: roi_x2] = \
  331. img[cross_y1: cross_y2, cross_x1: cross_x2]
  332. sample_img = cv2.resize(
  333. sample_img, (target_size, target_size), interpolation=cv2.INTER_AREA)
  334. return sample_img
  335. def box_horizontal_flip(bboxes, width):
  336. oldx1 = bboxes[:, 0].copy()
  337. oldx2 = bboxes[:, 2].copy()
  338. bboxes[:, 0] = width - oldx2 - 1
  339. bboxes[:, 2] = width - oldx1 - 1
  340. if bboxes.shape[0] != 0 and (bboxes[:, 2] < bboxes[:, 0]).all():
  341. raise ValueError(
  342. "RandomHorizontalFlip: invalid box, x2 should be greater than x1")
  343. return bboxes
  344. def segms_horizontal_flip(segms, height, width):
  345. def _flip_poly(poly, width):
  346. flipped_poly = np.array(poly)
  347. flipped_poly[0::2] = width - np.array(poly[0::2]) - 1
  348. return flipped_poly.tolist()
  349. def _flip_rle(rle, height, width):
  350. if 'counts' in rle and type(rle['counts']) == list:
  351. rle = mask_util.frPyObjects([rle], height, width)
  352. mask = mask_util.decode(rle)
  353. mask = mask[:, ::-1, :]
  354. rle = mask_util.encode(np.array(mask, order='F', dtype=np.uint8))
  355. return rle
  356. def is_poly(segm):
  357. if not isinstance(segm, (list, dict)):
  358. raise Exception("Invalid segm type: {}".format(type(segm)))
  359. return isinstance(segm, list)
  360. flipped_segms = []
  361. for segm in segms:
  362. if is_poly(segm):
  363. # Polygon format
  364. flipped_segms.append([_flip_poly(poly, width) for poly in segm])
  365. else:
  366. # RLE format
  367. import pycocotools.mask as mask_util
  368. flipped_segms.append(_flip_rle(segm, height, width))
  369. return flipped_segms