def translate(image, annotations, prob=0.5, border_value=(128, 128, 128)): assert 'bboxes' in annotations, 'annotations should contain bboxes even if it is empty' random_prob = np.random.uniform() if random_prob < (1 - prob): return image, annotations h, w = image.shape[:2] bboxes = annotations['bboxes'] if bboxes.shape[0] != 0: min_x1, min_y1 = np.min(bboxes, axis=0)[:2].astype(np.int32) max_x2, max_y2 = np.max(bboxes, axis=0)[2:].astype(np.int32) translation_matrix = translation_xy(min=(min(-(min_x1 // 2), 0), min(-(min_y1 // 2), 0)), max=(max((w - 1 - max_x2) // 2, 1), max((h - 1 - max_y2) // 2, 1)), prob=1.) else: translation_matrix = translation_xy(min=(min(-w // 8, 0), min(-h // 8, 0)), max=(max(w // 8, 1), max(h // 8, 1))) translation_matrix = change_transform_origin(translation_matrix, (w / 2, h / 2)) image = cv2.warpAffine( image, translation_matrix[:2, :], dsize=(w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue=border_value, ) if bboxes.shape[0] != 0: new_bboxes = [] for bbox in bboxes: x1, y1, x2, y2 = bbox points = translation_matrix.dot([ [x1, x2, x1, x2], [y1, y2, y2, y1], [1, 1, 1, 1], ]) min_x, min_y = np.min(points, axis=1)[:2] max_x, max_y = np.max(points, axis=1)[:2] new_bboxes.append([min_x, min_y, max_x, max_y]) annotations['bboxes'] = np.array(new_bboxes).astype(np.float32) if 'quadrangles' in annotations and annotations['quadrangles'].shape[0] != 0: quadrangles = annotations['quadrangles'] translated_quadrangles = [] for quadrangle in quadrangles: quadrangle = np.concatenate([quadrangle, np.ones((4, 1))], axis=-1) translated_quadrangle = translation_matrix.dot(quadrangle.T).T[:, :2] quadrangle = reorder_vertexes(translated_quadrangle) translated_quadrangles.append(quadrangle) quadrangles = np.stack(translated_quadrangles) annotations['quadrangles'] = quadrangles xmin = np.min(quadrangles, axis=1)[:, 0] ymin = np.min(quadrangles, axis=1)[:, 1] xmax = np.max(quadrangles, axis=1)[:, 0] ymax = np.max(quadrangles, axis=1)[:, 1] bboxes = np.stack([xmin, ymin, xmax, ymax], axis=1) annotations['bboxes'] = bboxes return image, annotations
def rotate(image, annotations, prob=0.5, border_value=(128, 128, 128)): assert 'bboxes' in annotations, 'annotations should contain bboxes even if it is empty' random_prob = np.random.uniform() if random_prob < (1 - prob): return image, annotations rotate_degree = np.random.uniform(low=-45, high=45) h, w = image.shape[:2] # Compute the rotation matrix. M = cv2.getRotationMatrix2D(center=(w / 2, h / 2), angle=rotate_degree, scale=1) # Get the sine and cosine from the rotation matrix. abs_cos_angle = np.abs(M[0, 0]) abs_sin_angle = np.abs(M[0, 1]) # Compute the new bounding dimensions of the image. new_w = int(h * abs_sin_angle + w * abs_cos_angle) new_h = int(h * abs_cos_angle + w * abs_sin_angle) # Adjust the rotation matrix to take into account the translation. M[0, 2] += new_w // 2 - w // 2 M[1, 2] += new_h // 2 - h // 2 # Rotate the image. image = cv2.warpAffine(image, M=M, dsize=(new_w, new_h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue=border_value) bboxes = annotations['bboxes'] if bboxes.shape[0] != 0: new_bboxes = [] for bbox in bboxes: x1, y1, x2, y2 = bbox points = M.dot([ [x1, x2, x1, x2], [y1, y2, y2, y1], [1, 1, 1, 1], ]) # Extract the min and max corners again. min_xy = np.sort(points, axis=1)[:, :2] min_x = np.mean(min_xy[0]) min_y = np.mean(min_xy[1]) max_xy = np.sort(points, axis=1)[:, 2:] max_x = np.mean(max_xy[0]) max_y = np.mean(max_xy[1]) new_bboxes.append([min_x, min_y, max_x, max_y]) annotations['bboxes'] = np.array(new_bboxes, dtype=np.float32) if 'quadrangles' in annotations and annotations['quadrangles'].shape[ 0] != 0: quadrangles = annotations['quadrangles'] rotated_quadrangles = [] for quadrangle in quadrangles: quadrangle = np.concatenate( [quadrangle, np.ones((4, 1))], axis=-1) rotated_quadrangle = M.dot(quadrangle.T).T[:, :2] quadrangle = reorder_vertexes(rotated_quadrangle) rotated_quadrangles.append(quadrangle) quadrangles = np.stack(rotated_quadrangles) annotations['quadrangles'] = quadrangles xmin = np.min(quadrangles, axis=1)[:, 0] ymin = np.min(quadrangles, axis=1)[:, 1] xmax = np.max(quadrangles, axis=1)[:, 0] ymax = np.max(quadrangles, axis=1)[:, 1] bboxes = np.stack([xmin, ymin, xmax, ymax], axis=1) annotations['bboxes'] = bboxes return image, annotations