generate_real_data.py 文件源码

def mask_real_image(color, depth, depth_render):
    mask = (depth_render != 0).astype(np.uint8)[:, :, np.newaxis]
    masked_rgb = color * mask

    masked_hsv = cv2.cvtColor(masked_rgb, cv2.COLOR_BGR2HSV)
    saturation_mask = (masked_hsv[:, :, 2] <= SATURATION_THRESHOLD)[:, :, np.newaxis].astype(np.uint8)
    total_mask = np.bitwise_and(mask, saturation_mask)

    masked_color = color * total_mask
    masked_depth = depth[:total_mask.shape[0], :total_mask.shape[1]] * total_mask[:, :, 0]
    return masked_color, masked_depth