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
Beispiel #2
0
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