-
Notifications
You must be signed in to change notification settings - Fork 9
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Vlad St
committed
Feb 18, 2021
1 parent
28f943b
commit 47d0f79
Showing
7 changed files
with
270 additions
and
114 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -127,3 +127,5 @@ dmypy.json | |
|
||
# Pyre type checker | ||
.pyre/ | ||
|
||
.vscode/* |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
) | ||
|
Oops, something went wrong.