File size: 4,184 Bytes
36c95ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import argparse
import os
import sys
import cv2
import torch
import kornia as dgm
def load_depth(file_name):
"""Load the depth using the syntel SDK and converts to torch.Tensor."""
if not os.path.isfile(file_name):
raise AssertionError(f"Invalid file {file_name}")
import sintel_io
depth = sintel_io.depth_read(file_name)
return torch.from_numpy(depth).view(1, 1, *depth.shape).float()
def load_camera_data(file_name):
"""Load the camera data using the syntel SDK and converts to torch.Tensor."""
if not os.path.isfile(file_name):
raise AssertionError(f"Invalid file {file_name}")
import sintel_io
intrinsic, extrinsic = sintel_io.cam_read(file_name)
return intrinsic, extrinsic
def load_image(file_name):
"""Load the image with OpenCV and converts to torch.Tensor."""
if not os.path.isfile(file_name):
raise AssertionError(f"Invalid file {file_name}")
# load image with OpenCV
img = cv2.imread(file_name, cv2.IMREAD_COLOR)
# convert image to torch tensor
tensor = dgm.utils.image_to_tensor(img).float() / 255.0
return tensor.view(1, *tensor.shape) # 1xCxHxW
def load_data(root_path, sequence_name, frame_id):
# index paths
file_name = 'frame_%04d' % (frame_id)
image_file = os.path.join(root_path, 'clean', sequence_name, file_name + '.png')
depth_file = os.path.join(root_path, 'depth', sequence_name, file_name + '.dpt')
camera_file = os.path.join(root_path, 'camdata_left', sequence_name, file_name + '.cam')
# load the actual data
image = load_image(image_file)
depth = load_depth(depth_file)
# load camera data and create pinhole
height, width = image.shape[-2:]
intrinsics, extrinsics = load_camera_data(camera_file)
camera = dgm.utils.create_pinhole(intrinsics, extrinsics, height, width)
return image, depth, camera
def DepthWarperApp():
parser = argparse.ArgumentParser(description='Warp images by depth application.')
# data parameters
parser.add_argument('--input-dir', type=str, required=True, help='the path to the directory with the input data.')
parser.add_argument('--output-dir', type=str, required=True, help='the path to output the results.')
parser.add_argument('--sequence-name', type=str, default='alley_1', help='the name of the sequence.')
parser.add_argument('--frame-ref-id', type=int, default=1, help='the id for the reference image in the sequence.')
parser.add_argument('--frame-i-id', type=int, default=2, help='the id for the image i in the sequence.')
# device parameters
parser.add_argument('--cuda', action='store_true', default=False, help='enables CUDA training')
parser.add_argument('--seed', type=int, default=666, metavar='S', help='random seed (default: 666)')
args = parser.parse_args()
torch.manual_seed(args.seed)
# configure syntel SDK path
root_path = os.path.abspath(args.input_dir)
sys.path.append(os.path.join(root_path, 'sdk/python'))
# load the data
root_dir = os.path.join(root_path, 'training')
img_ref, depth_ref, cam_ref = load_data(root_dir, args.sequence_name, args.frame_ref_id)
img_i, _, cam_i = load_data(root_dir, args.sequence_name, args.frame_i_id)
# instantiate the homography warper from `kornia`
warper = dgm.DepthWarper(cam_i)
warper.compute_homographies(cam_ref)
# compute the inverse depth and warp the source image
inv_depth_ref = 1.0 / depth_ref
img_i_to_ref = warper(inv_depth_ref, img_i)
# generate occlusion mask
mask = ((img_ref - img_i_to_ref).mean(1) < 1e-1).float()
img_vis_warped = 0.5 * img_i_to_ref + img_ref
img_vis_warped_masked = mask * (0.5 * img_i_to_ref + img_ref)
# save warped image to disk
file_name = os.path.join(args.output_dir, f'warped_{args.frame_i_id}_to_{args.frame_ref_id}.png')
cv2.imwrite(file_name, dgm.utils.tensor_to_image(255.0 * img_vis_warped))
cv2.imwrite(file_name + 'mask.png', dgm.utils.tensor_to_image(255.0 * mask))
cv2.imwrite(file_name + 'warpedmask.png', dgm.utils.tensor_to_image(255.0 * img_vis_warped_masked))
if __name__ == "__main__":
DepthWarperApp()
|