From 47d0f7948ba99a96e9d2d10f264a6900b41ab6f7 Mon Sep 17 00:00:00 2001 From: Vlad St Date: Thu, 18 Feb 2021 16:38:17 +0200 Subject: [PATCH] added calibration script --- .gitignore | 2 + calibration.py | 128 ++++++++++++++++++++++++++ data/asymmetric_box.obj | 168 +++++++++++++++++------------------ data/structure.py | 27 ------ src/utils/__init__.py | 3 +- src/utils/geometric.py | 9 +- src/utils/save_pointcloud.py | 47 ++++++++++ 7 files changed, 270 insertions(+), 114 deletions(-) create mode 100644 calibration.py delete mode 100644 data/structure.py create mode 100644 src/utils/save_pointcloud.py diff --git a/.gitignore b/.gitignore index b6e4761..e701302 100644 --- a/.gitignore +++ b/.gitignore @@ -127,3 +127,5 @@ dmypy.json # Pyre type checker .pyre/ + +.vscode/* diff --git a/calibration.py b/calibration.py new file mode 100644 index 0000000..5f76edc --- /dev/null +++ b/calibration.py @@ -0,0 +1,128 @@ +import torch +import torchvision.utils as tu +import cv2 +import numpy as np +import os +import sys +import argparse +import src.models as models +from src.utils.geometric import ExtrinsicsCalculator, BoxRenderFlags +from src.utils.save_pointcloud import save_ply + + + +STATIC_IMAGE_SIZE = (180, 320) # (height,width) +STATIC_BOX_FLAG = BoxRenderFlags.LABEL_DOWN_AS_BACKGROUND +STATIC_DEVICE = 'cpu' + +def parse_arguments(args): + usage_text = ( + "Calibration script." + "Usage: python calibration.py [options]," + " with [options]:" + ) + parser = argparse.ArgumentParser(description=usage_text) + parser.add_argument("-d","--depth", type = str, help = "Path to depthmap", required = True) + parser.add_argument("-m","--model_path", type = str, help = "Path to saved model params", required = True) + parser.add_argument("-o","--save_path", type = str, help = "Path to save results", required = True) + parser.add_argument("-b","--box_path", type = str, help = "Path to box", default = r"data/asymmetric_box.obj") + parser.add_argument("-s","--scale", type = float, help = "Factor that converts depthmap to meters") + parser.add_argument("-i","--intrinsics", nargs=4, metavar=('fx', 'cx', 'fy', 'cy',), + help="camera instrinsic factors", type=float, + default=None) + return parser.parse_known_args(args) + +def align( + model : torch.nn.Module, + depthmap : torch.Tensor, + intrinsics : torch.Tensor, + box_path : str, + device : str, + save_path : str, + box_flag : BoxRenderFlags = STATIC_BOX_FLAG, + confidence : float = 0.75, +) -> None: + os.makedirs(save_path, exist_ok=True) + predictions = model(depthmap)[1] + _, nclasses, height, width = predictions.shape + + labels = predictions.argmax(dim = 1, keepdim = True) + one_hot = torch.nn.functional.one_hot(labels.squeeze(),num_classes = nclasses).permute(2,0,1).unsqueeze(0) + + extrinsics_calculator = ExtrinsicsCalculator(box_path, device, box_flag) + + extrinsics, _, pointclouds = extrinsics_calculator.forward(depthmap, one_hot, intrinsics) + extrinsics = extrinsics.squeeze().numpy().T + pointclouds = pointclouds[0].permute(1,2,0).reshape(-1,3).numpy() + save_ply(os.path.join(save_path, "original.ply"),pointclouds , scale = 1) + pcloud_homo = np.concatenate([pointclouds, np.ones((height * width, 1))], axis = 1) + transformed_pcloud = pcloud_homo.dot(extrinsics) + save_ply(os.path.join(save_path, "transformed.ply"),transformed_pcloud[:,:3], scale = 1) + np.savetxt(os.path.join(save_path, "extrinsics.txt"), extrinsics) + print(extrinsics) + + + +def loadModel( + path_to_model : str, + device : str +) -> torch.nn.Module: + print("Loading previously saved model from {}".format(path_to_model)) + checkpoint = torch.load(path_to_model) + model_params = { + 'width': 320, + 'height': 180, + 'ndf': 32, + 'upsample_type': "nearest", + } + + + model_name = checkpoint['model_name'] + if 'nclasses' in checkpoint: + nclasses = checkpoint['nclasses'] + + if 'ndf' in checkpoint: + model_params['ndf'] = checkpoint['ndf'] + + model_params['nclasses'] = nclasses + model = models.get_UNet_model(model_name, model_params) + model.load_state_dict(checkpoint['state_dict']) + model.to(device) + model.eval() + return model + +def loadData( + path_to_depthmap : str, + scale : float +) -> torch.Tensor: + depth_np = cv2.imread(path_to_depthmap, -1).astype(np.float32) + depth_np = cv2.resize(depth_np, STATIC_IMAGE_SIZE[::-1], interpolation=cv2.INTER_NEAREST) + depth_t = torch.from_numpy(depth_np).unsqueeze(0).unsqueeze(0) / scale + return depth_t + +if __name__ == "__main__": + args, _ = parse_arguments(sys.argv) + intrinsics = torch.FloatTensor([ + args.intrinsics[0], + 0.0, + args.intrinsics[1], + 0.0, + args.intrinsics[2], + args.intrinsics[3], + 0.0, + 0.0, + 1.0 + ]).view((3,3)).unsqueeze(0) + + model = loadModel(args.model_path,STATIC_DEVICE) + depthmap = loadData(args.depth, args.scale) + + align( + model, + depthmap, + intrinsics, + args.box_path, + STATIC_DEVICE, + args.save_path + ) + diff --git a/data/asymmetric_box.obj b/data/asymmetric_box.obj index 3cbd968..9ab5022 100644 --- a/data/asymmetric_box.obj +++ b/data/asymmetric_box.obj @@ -3,9 +3,9 @@ mtllib SemanticAsymmetricCalibrationNew.mtl o mid_bottom_front_2f v 50.000397 -165.000549 -0.000796 -v 50.000427 394.999421 -0.000583 -v 50.000397 -165.000396 -410.000305 -v 50.000427 394.999604 -410.000092 +v 50.000427 390 -0.000583 +v 50.000397 -165.000396 -405.000305 +v 50.000427 390.999604 -405.000092 vn 1.0000 -0.0000 -0.0000 usemtl material_0 s 1 @@ -13,19 +13,19 @@ f 3//1 4//1 2//1 f 2//1 1//1 3//1 o mid_bottom_back_2b v -279.999908 -165.000762 -0.000796 -v -279.999847 394.99936 -410.000092 -v -279.999908 -165.000762 -410.000305 -v -279.999847 394.99936 -0.000583 +v -279.999847 390 -405.000305 +v -279.999908 -165.000762 -405.000305 +v -279.999847 390 -0.000583 vn -1.0000 0.0000 0.0000 usemtl material_0 s 1 f 7//2 5//2 8//2 f 8//2 6//2 7//2 o mid_bottom_down_2d -v -279.999847 394.99936 -410.000092 -v 50.000397 -165.000396 -410.000305 -v -279.999908 -165.00061 -410.000305 -v 50.000427 394.999604 -410.000092 +v -279.999847 390 -405.000305 +v 50.000397 -165.000396 -405.000305 +v -279.999908 -165.00061 -405.000305 +v 50.000427 390 -405.000305 vn 0.0000 0.0000 -1.0000 usemtl material_0 s 1 @@ -33,9 +33,9 @@ f 9//3 10//3 11//3 f 10//3 9//3 12//3 o mid_bottom_up_2u v 50.000397 -165.000549 -0.000796 -v 50.000427 394.999421 -0.000583 +v 50.000427 390 -0.000583 v -279.999908 -165.000762 -0.000796 -v -279.999847 394.999177 -0.000583 +v -279.999847 390 -0.000583 vn 0.0000 -0.0000 1.0000 usemtl material_0 s 1 @@ -44,27 +44,27 @@ f 16//4 15//4 14//4 o mid_bottom_right_2r v 50.000397 -165.000549 -0.000796 v -279.999908 -165.000762 -0.000796 -v 50.000397 -165.000396 -410.000305 -v -279.999908 -165.00061 -410.000305 +v 50.000397 -165.000396 -405.000305 +v -279.999908 -165.00061 -405.000305 vn 0.0000 -1.0000 -0.0000 usemtl material_0 s 1 f 19//5 18//5 20//5 f 19//5 17//5 18//5 o mid_bottom_left_2l -v 50.000427 394.999421 -0.000583 -v -279.999847 394.99936 -410.000092 -v 50.000427 394.999604 -410.000092 -v -279.999847 394.999177 -0.000583 +v 50.000427 390 -0.000583 +v -279.999847 390 -405.000305 +v 50.000427 390 -405.000305 +v -279.999847 390 -0.000583 vn -0.0000 1.0000 0.0000 usemtl material_0 s 1 f 24//6 21//6 23//6 f 23//6 22//6 24//6 o mid_top_right_3r -v -509.999573 -165.000549 410.000061 -v 50.000427 -165.000549 410.000061 -v -509.999573 -165.000488 0.000554 +v -505 -165.000549 405.000061 +v 50.000427 -165.000549 405.000061 +v -505 -165.000488 0.000554 v 50.000427 -165.000488 0.000554 vn 0.0000 -1.0000 -0.0000 usemtl material_0 @@ -72,10 +72,10 @@ s 1 f 27//7 28//7 26//7 f 26//7 25//7 27//7 o mid_top_left_3l -v -509.999573 164.999451 410.000122 +v -505 164.999451 405.000122 v 50.000427 164.999512 0.000608 -v -509.999573 164.999512 0.000608 -v 50.000427 164.999451 410.000122 +v -505 164.999512 0.000608 +v 50.000427 164.999451 405.000122 vn 0.0000 1.0000 0.0000 usemtl material_0 s 1 @@ -83,8 +83,8 @@ f 31//8 29//8 32//8 f 32//8 30//8 31//8 o mid_top_down_3d v 50.000427 164.999512 0.000608 -v -509.999573 -165.000488 0.000554 -v -509.999573 164.999512 0.000608 +v -505 -165.000488 0.000554 +v -505 164.999512 0.000608 v 50.000427 -165.000488 0.000554 vn 0.0000 0.0000 -1.0000 usemtl material_0 @@ -92,150 +92,150 @@ s 1 f 33//9 34//9 35//9 f 34//9 33//9 36//9 o mid_top_up_3u -v -509.999573 -165.000549 410.000061 -v 50.000427 -165.000549 410.000061 -v -509.999573 164.999451 410.000122 -v 50.000427 164.999451 410.000122 +v -505 -165.000549 405.000061 +v 50.000427 -165.000549 405.000061 +v -505 164.999451 405.000122 +v 50.000427 164.999451 405.000122 vn 0.0000 -0.0000 1.0000 usemtl material_0 s 1 f 37//10 38//10 39//10 f 40//10 39//10 38//10 o mid_top_front_3f -v 50.000427 -165.000549 410.000061 +v 50.000427 -165.000549 405.000061 v 50.000427 164.999512 0.000608 v 50.000427 -165.000488 0.000554 -v 50.000427 164.999451 410.000122 +v 50.000427 164.999451 405.000122 vn 1.0000 0.0000 0.0000 usemtl material_0 s 1 f 44//11 41//11 43//11 f 43//11 42//11 44//11 o mid_top_back_3b -v -509.999573 -165.000549 410.000061 -v -509.999573 164.999451 410.000122 -v -509.999573 -165.000488 0.000554 -v -509.999573 164.999512 0.000608 +v -505 -165.000549 405.000061 +v -505 164.999451 405.000122 +v -505 -165.000488 0.000554 +v -505 164.999512 0.000608 vn -1.0000 -0.0000 0.0000 usemtl material_0 s 1 f 47//12 46//12 48//12 f 47//12 45//12 46//12 o bottom_right_1r -v -279.999969 -165.000427 -410.000336 -v 280.000031 -165.000427 -410.000336 -v -279.999969 -165.000351 -819.999817 -v 280.000031 -165.000351 -819.999817 +v -279.999969 -165.000427 -405.000305 +v 275 -165.000427 -405.000305 +v -279.999969 -165.000351 -810 +v 275 -165.000351 -810 vn 0.0000 -1.0000 -0.0000 usemtl material_0 s 1 f 51//13 52//13 50//13 f 50//13 49//13 51//13 o bottom_up_1u -v -279.999969 -165.000427 -410.000336 -v 280.000031 -165.000427 -410.000336 -v -279.999969 164.999573 -410.000275 -v 280.000031 164.999573 -410.000275 +v -279.999969 -165.000427 -405.000305 +v 275 -165.000427 -405.000305 +v -279.999969 164.999573 -405.000305 +v 275 164.999573 -405.000305 vn 0.0000 -0.0000 1.0000 usemtl material_0 s 1 f 53//14 54//14 55//14 f 56//14 55//14 54//14 o bottom_back_1b -v -279.999969 -165.000427 -410.000336 -v -279.999969 164.999573 -410.000275 -v -279.999969 -165.000351 -819.999817 -v -279.999969 164.999649 -819.999817 +v -279.999969 -165.000427 -405.000305 +v -279.999969 164.999573 -405.000305 +v -279.999969 -165.000351 -810 +v -279.999969 164.999649 -810 vn -1.0000 0.0000 0.0000 usemtl material_0 s 1 f 59//15 58//15 60//15 f 59//15 57//15 58//15 o bottom_front_1f -v 280.000031 -165.000427 -410.000336 -v 280.000031 164.999649 -819.999817 -v 280.000031 -165.000351 -819.999817 -v 280.000031 164.999573 -410.000275 +v 275 -165.000427 -405.000305 +v 275 164.999649 -810 +v 275 -165.000351 -810 +v 275 164.999573 -405.000305 vn 1.0000 0.0000 0.0000 usemtl material_0 s 1 f 64//16 61//16 63//16 f 63//16 62//16 64//16 o bottom_down_1d -v 280.000031 164.999649 -819.999817 -v -279.999969 -165.000351 -819.999817 -v -279.999969 164.999649 -819.999817 -v 280.000031 -165.000351 -819.999817 +v 275 164.999649 -810 +v -279.999969 -165.000351 -810 +v -279.999969 164.999649 -810 +v 275 -165.000351 -810 vn 0.0000 0.0000 -1.0000 usemtl material_0 s 1 f 65//17 66//17 67//17 f 66//17 65//17 68//17 o bottom_left_1l -v -279.999969 164.999573 -410.000275 -v 280.000031 164.999649 -819.999817 -v -279.999969 164.999649 -819.999817 -v 280.000031 164.999573 -410.000275 +v -279.999969 164.999573 -405.000305 +v 275 164.999649 -810 +v -279.999969 164.999649 -810 +v 275 164.999573 -405.000305 vn 0.0000 1.0000 0.0000 usemtl material_0 s 1 f 71//18 69//18 72//18 f 72//18 70//18 71//18 o top_front_4f -v 50.000427 -394.999451 819.999573 -v 50.000458 165.000504 819.999756 -v 50.000427 -394.999298 410.000092 -v 50.000458 165.000702 410.000305 +v 50.000427 -390 810 +v 50.000458 165.000504 810 +v 50.000427 -390 405 +v 50.000458 165.000702 405 vn 1.0000 -0.0000 -0.0000 usemtl material_0 s 1 f 75//19 76//19 74//19 f 74//19 73//19 75//19 o top_down_4d -v -279.999817 165.000458 410.000305 -v 50.000427 -394.999298 410.000092 -v -279.999878 -394.999512 410.000092 -v 50.000458 165.000702 410.000305 +v -279.999817 165.000458 405.000305 +v 50.000427 -390 405.000092 +v -279.999878 -390 405.000092 +v 50.000458 165.000702 405.000305 vn -0.0000 0.0000 -1.0000 usemtl material_0 s 1 f 77//20 78//20 79//20 f 78//20 77//20 80//20 o top_right_4r -v 50.000427 -394.999451 819.999573 -v -279.999878 -394.999664 819.999573 -v 50.000427 -394.999298 410.000092 -v -279.999878 -394.999512 410.000092 +v 50.000427 -390 810 +v -279.999878 -390 810 +v 50.000427 -390 405.000092 +v -279.999878 -390 405.000092 vn 0.0000 -1.0000 -0.0000 usemtl material_0 s 1 f 83//21 82//21 84//21 f 83//21 81//21 82//21 o top_up_4u -v 50.000427 -394.999451 819.999573 -v 50.000458 165.000504 819.999756 -v -279.999878 -394.999664 819.999573 -v -279.999817 165.000259 819.999756 +v 50.000427 -390 810 +v 50.000458 165.000504 810 +v -279.999878 -390 810 +v -279.999817 165.000259 810 vn 0.0000 -0.0000 1.0000 usemtl material_0 s 1 f 85//22 86//22 87//22 f 88//22 87//22 86//22 o top_back_4b -v -279.999878 -394.999664 819.999573 -v -279.999817 165.000458 410.000305 -v -279.999878 -394.999512 410.000092 -v -279.999817 165.000259 819.999756 +v -279.999878 -390 810 +v -279.999817 165.000458 405.000305 +v -279.999878 -390 405.000092 +v -279.999817 165.000259 810 vn -1.0000 0.0000 0.0000 usemtl material_0 s 1 f 91//23 89//23 92//23 f 92//23 90//23 91//23 o top_left_4l -v 50.000458 165.000504 819.999756 -v -279.999817 165.000458 410.000305 -v 50.000458 165.000702 410.000305 -v -279.999817 165.000259 819.999756 +v 50.000458 165.000504 810 +v -279.999817 165.000458 405.000305 +v 50.000458 165.000702 405.000305 +v -279.999817 165.000259 810 vn -0.0000 1.0000 0.0000 usemtl material_0 s 1 diff --git a/data/structure.py b/data/structure.py deleted file mode 100644 index 9030c2c..0000000 --- a/data/structure.py +++ /dev/null @@ -1,27 +0,0 @@ -from torch.utils.data.dataset import Dataset - -import os - - -class StructureDataLoadParameters(object): - def __init__(self, path, depth, normal, label, pose): - self.path = path - self.load_depth = depth - self.load_normal = normal - self.load_label = label - self.load_pose = pose - - -class StructureData(Dataset): - #360D Dataset# - def __init__(self, params): - super(StructureData, self).__init__() - files = os.listdir(params.path) - if params.load_depth: - depths = [f for f in files if f.endswith("_depth.exr")] - if params.load_normal: - normals = [f for f in files if f.endswith("_normal.exr")] - def __len__(self): - return self.length - - def __getitem__(self, idx): diff --git a/src/utils/__init__.py b/src/utils/__init__.py index 138c1a2..90e9607 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -3,4 +3,5 @@ from .visualization import * from .image_utils import * from .geometric import * -from .train_utils import * \ No newline at end of file +from .train_utils import * +from .save_pointcloud import * \ No newline at end of file diff --git a/src/utils/geometric.py b/src/utils/geometric.py index fab0bb7..1ac20dc 100644 --- a/src/utils/geometric.py +++ b/src/utils/geometric.py @@ -2,8 +2,13 @@ from src.utils import projections from src.io import box_model_loader, plywrite from src.dataset.rendering.box_renderer import BoxRenderFlags +from enum import Flag, auto #import .projections -output_path = "D:\\VCL\\Users\\vlad\\Datasets\\structureNetData\\output_structurenet" +class BoxRenderFlags (Flag): + LABEL_UP_AS_BACKGROUND = auto() + LABEL_DOWN_AS_BACKGROUND = auto() + LABEL_TOP_AND_BOTTOM_AS_BACKGROUND = LABEL_UP_AS_BACKGROUND | LABEL_DOWN_AS_BACKGROUND + import os.path as osp ''' Given two sets (set1 and set2 with dimensions b,c,N [b:batch,c:channels,N:spatial]) @@ -73,7 +78,7 @@ def computeNonRigidTransformation( R = torch.cat(list_R, dim = 0) scale = torch.cat(list_scale, dim = 0).unsqueeze(-1) - t = -torch.bmm(R,centroid_source)*scale + centroid_target + t = -torch.bmm(R,centroid_source) + centroid_target return R, t, scale ''' diff --git a/src/utils/save_pointcloud.py b/src/utils/save_pointcloud.py new file mode 100644 index 0000000..024391d --- /dev/null +++ b/src/utils/save_pointcloud.py @@ -0,0 +1,47 @@ +import numpy as np + +def save_ply(filename, tensor, scale, color=[0,0,0] , normals = None): + #w,h,c = tensor.shape + if len(tensor.shape) == 2: + tensor = np.expand_dims(tensor, 0) + + h,w,c = tensor.shape + x_coords = tensor[:, :, 0] * scale + y_coords = tensor[:, :, 1] * scale + z_coords = tensor[:, :, 2] * scale + if normals is not None: + if len(normals.shape) == 2: + normals = np.expand_dims(normals, 0) + nx_coords = normals[:, :, 0] + ny_coords = normals[:, :, 1] + nz_coords = normals[:, :, 2] + with open(filename, "w") as ply_file: + ply_file.write("ply\n") + ply_file.write("format ascii 1.0\n") + ply_file.write("element vertex {}\n".format(w * h)) + ply_file.write("property float x\n") + ply_file.write("property float y\n") + ply_file.write("property float z\n") + if normals is not None: + ply_file.write('property float nx\n') + ply_file.write('property float ny\n') + ply_file.write('property float nz\n') + ply_file.write("property uchar red\n") + ply_file.write("property uchar green\n") + ply_file.write("property uchar blue\n") + ply_file.write("end_header\n") + + if normals is None: + for x in np.arange(h): + for y in np.arange(w): + ply_file.write("{} {} {} {} {} {}\n".format(\ + x_coords[x, y], y_coords[x, y], z_coords[x, y],\ + color[0],color[1],color[2] + )) + else: + for x in np.arange(h): + for y in np.arange(w): + ply_file.write("{} {} {} {} {} {} {} {} {}\n".format(\ + x_coords[x, y], y_coords[x, y], z_coords[x, y],\ + nx_coords[x, y], ny_coords[x, y], nz_coords[x, y],\ + color[0],color[1],color[2])) \ No newline at end of file