generate_real_data.py 文件源码

def random_z_rotation(rgb, depth, pose, camera):
    rotation = random.uniform(-180, 180)
    rotation_matrix = Transform()
    rotation_matrix.set_rotation(0, 0, math.radians(rotation))

    pixel = center_pixel(pose, camera)
    new_rgb = rotate_image(rgb, rotation, pixel[0])
    new_depth = rotate_image(depth, rotation, pixel[0])
    # treshold below 50 means we remove some interpolation noise, which cover small holes
    mask = (new_depth >= 50).astype(np.uint8)[:, :, np.newaxis]
    rgb_mask = np.all(new_rgb != 0, axis=2).astype(np.uint8)
    kernel = np.array([[0, 1, 0],
                       [1, 1, 1],
                       [0, 1, 0]], np.uint8)
    # erode rest of interpolation noise which will affect negatively future blendings
    eroded_mask = cv2.erode(mask, kernel, iterations=2)
    eroded_rgb_mask = cv2.erode(rgb_mask, kernel, iterations=2)
    new_depth = new_depth * eroded_mask
    new_rgb = new_rgb * eroded_rgb_mask[:, :, np.newaxis]
    new_pose = combine_view_transform(pose, rotation_matrix)
    return new_rgb, new_depth, new_pose