def __call__(self, img, boxes, segments, labels):
'''
apply transform when calling
:param img: numpy.ndarray,(h,w,c)
:param boxes: numpy.ndarray,(n,4), x1,y1,x2,y2
:param segments: list(list),(n)(x),x is variant
:param labels: numpy.ndarray,(n)
:return:
'''
h, w, c = img.shape
while True:
mode = random.choice(self.sample_mode)
if mode == 1:
return img, boxes, segments, labels
min_iou = mode
for i in range(50):
new_w = random.uniform(self.min_crop_size * w, w)
new_h = random.uniform(self.min_crop_size * h, h)
# h / w in [0.5, 2]
if new_h / new_w < 0.5 or new_h / new_w > 2:
continue
left = random.uniform(w - new_w)
top = random.uniform(h - new_h)
patch = np.array(
(int(left), int(top), int(left + new_w), int(top + new_h)))
overlaps = bbox_overlaps(patch.reshape(-1, 4),
boxes.reshape(-1, 4)).reshape(-1)
if overlaps.min() < min_iou:
continue
# center of boxes should inside the crop img
center = (boxes[:, :2] + boxes[:, 2:]) / 2
mask = (center[:, 0] > patch[0]) * (
center[:, 1] > patch[1]) * (center[:, 0] < patch[2]) * (
center[:, 1] < patch[3])
if not mask.any():
continue
# adjust boxes
img = img[patch[1]:patch[3], patch[0]:patch[2]]
# box here is not valid enough
# we shound use mask to generate the new box
# boxes[:, 2:] = boxes[:, 2:].clip(max=patch[2:])
# boxes[:, :2] = boxes[:, :2].clip(min=patch[:2])
# boxes -= np.tile(patch[:2], 2)
boxes = []
out_segments = []
out_labels = []
for j, segment in enumerate(segments):
if not mask[j]:
continue
segment = np.array(segment).reshape(-1, 2)
segment = Polygon(segment)
bound = patch.copy()
bound = bound + np.array([1, 1, -1, -1])
bound = np.vstack((bound[[0, 2, 2,
0]], bound[[1, 1, 3,
3]])).transpose()
bound = Polygon(bound)
segment = bound.intersection(segment)
try:
segment = np.array(segment.exterior.coords)
except Exception as e:
print(e)
continue
segment -= patch[:2]
x1, y1 = np.min(segment, 0)
x2, y2 = np.max(segment, 0)
boxes.append([x1, y1, x2, y2])
out_labels.append(labels[j])
out_segments.append(list(segment.reshape(-1)))
boxes = np.array(boxes).astype(np.int32)
return img, boxes, out_segments, np.array(out_labels)
