def __init__(self, cfg, in_channels, optimizer=None):
super().__init__()
self.cfg = cfg
# Refinement
self.transform_pc = TransformPC(cfg)
self.feature_projection = FeatureProjection(cfg)
self.pc_encode = PointNet2(cfg)
self.displacement_net = LinearDisplacementNet(cfg)
self.optimizer = None if optimizer is None else optimizer(
self.parameters())
# emd loss
self.emd_dist = emd.emdModule()
if torch.cuda.is_available():
self.transform_pc = torch.nn.DataParallel(
self.transform_pc, device_ids=cfg.CONST.DEVICE).cuda()
self.feature_projection = torch.nn.DataParallel(
self.feature_projection, device_ids=cfg.CONST.DEVICE).cuda()
self.pc_encode = torch.nn.DataParallel(
self.pc_encode, device_ids=cfg.CONST.DEVICE).cuda()
self.displacement_net = torch.nn.DataParallel(
self.displacement_net, device_ids=cfg.CONST.DEVICE).cuda()
self.emd_dist = torch.nn.DataParallel(
self.emd_dist, device_ids=cfg.CONST.DEVICE).cuda()
self.cuda()
def build_train_loss(self):
# Set up loss functions
self.chamfer_dist = torch.nn.DataParallel(ChamferDistance().to(
self.gpu_ids[0]),
device_ids=self.gpu_ids)
self.chamfer_dist_mean = torch.nn.DataParallel(
ChamferDistanceMean().to(self.gpu_ids[0]), device_ids=self.gpu_ids)
self.emd_dist = torch.nn.DataParallel(emd.emdModule().to(
self.gpu_ids[0]),
device_ids=self.gpu_ids)
def __init__(self, input_pcs):
super().__init__()
self.pred_pcs = nn.Parameter(input_pcs)
self.emd_dist = emd.emdModule()
self.optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
if torch.cuda.is_available():
self.pred_pcs.cuda()
self.emd_dist.cuda()
self.cuda()
def __init__(self, cfg, optimizer=None, scheduler=None):
super().__init__()
self.cfg = cfg
# Graphx Reconstructor
self.reconstructor = Graphx_Rec(
cfg=cfg,
in_channels=3,
in_instances=cfg.GRAPHX.NUM_INIT_POINTS,
activation=nn.ReLU(),
)
self.optimizer = None if optimizer is None else optimizer(self.reconstructor.parameters())
self.scheduler = None if scheduler or optimizer is None else scheduler(self.optimizer)
# emd loss
self.emd_dist = emd.emdModule()
if torch.cuda.is_available():
# Reconstructor
self.reconstructor = torch.nn.DataParallel(self.reconstructor, device_ids=cfg.CONST.DEVICE).cuda()
# loss
self.emd_dist = torch.nn.DataParallel(self.emd_dist, device_ids=cfg.CONST.DEVICE).cuda()
self.cuda()
def test_rec_net(cfg):
# Enable the inbuilt cudnn auto-tuner to find the best algorithm to use
torch.backends.cudnn.benchmark = True
test_transforms = utils.data_transforms.Compose([
utils.data_transforms.ToTensor(),
])
dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[
cfg.DATASET.TEST_DATASET](cfg)
test_data_loader = torch.utils.data.DataLoader(
dataset=dataset_loader.get_dataset(utils.data_loaders.DatasetType.TEST,
test_transforms),
batch_size=cfg.TEST.BATCH_SIZE,
num_workers=1,
pin_memory=True,
shuffle=False)
# Set up networks
# The parameters here need to be set in cfg
net = GRAPHX_REC_MODEL(
cfg=cfg,
optimizer=lambda x: torch.optim.Adam(x,
lr=cfg.TRAIN.GRAPHX_LEARNING_RATE,
weight_decay=cfg.TRAIN.
GRAPHX_WEIGHT_DECAY),
scheduler=lambda x: MultiStepLR(
x, milestones=cfg.TRAIN.MILESTONES, gamma=cfg.TRAIN.GAMMA),
)
if torch.cuda.is_available():
net = torch.nn.DataParallel(net).cuda()
# Load weight
# Load weight for encoder, decoder
print('[INFO] %s Loading reconstruction weights from %s ...' %
(dt.now(), cfg.TEST.WEIGHT_PATH))
rec_checkpoint = torch.load(cfg.TEST.WEIGHT_PATH)
net.load_state_dict(rec_checkpoint['net'])
print('[INFO] Best reconstruction result at epoch %d ...' %
rec_checkpoint['epoch_idx'])
# Set up loss functions
emd_dist = emd.emdModule()
cd = ChamferLoss().cuda()
# Batch average meterics
cd_distances = utils.network_utils.AverageMeter()
emd_distances = utils.network_utils.AverageMeter()
# Switch models to evaluation mode
net.eval()
n_batches = len(test_data_loader)
# Testing loop
for sample_idx, (taxonomy_names, sample_names, rendering_images, model_azi,
model_ele, init_point_clouds,
ground_truth_point_clouds) in enumerate(test_data_loader):
with torch.no_grad():
# Only one image per sample
rendering_images = torch.squeeze(rendering_images, 1)
# Get data from data loader
rendering_images = utils.network_utils.var_or_cuda(
rendering_images)
model_azi = utils.network_utils.var_or_cuda(model_azi)
model_ele = utils.network_utils.var_or_cuda(model_ele)
init_point_clouds = utils.network_utils.var_or_cuda(
init_point_clouds)
ground_truth_point_clouds = utils.network_utils.var_or_cuda(
ground_truth_point_clouds)
#=================================================#
# Test the encoder, decoder #
#=================================================#
loss, pred_pc = net.module.valid_step(rendering_images,
init_point_clouds,
ground_truth_point_clouds)
# Compute CD, EMD
cd_distance = cd(pred_pc, ground_truth_point_clouds
) / cfg.TEST.BATCH_SIZE / cfg.CONST.NUM_POINTS
# compute reconstruction loss
emd_loss, _ = emd_dist(pred_pc,
ground_truth_point_clouds,
eps=0.005,
iters=50)
emd_distance = torch.sqrt(emd_loss).mean(1).mean()
# Append loss and accuracy to average metrics
cd_distances.update(cd_distance.item())
emd_distances.update(emd_distance.item())
print("Test on [%d/%d] data, CD: %.4f EMD %.4f" %
(sample_idx + 1, n_batches, cd_distance.item(),
emd_distance.item()))
# print result
print("Reconstruction result:")
print("CD result: ", cd_distances.avg)
print("EMD result", emd_distances.avg)
logname = cfg.TEST.RESULT_PATH
with open(logname, 'a') as f:
f.write('Reconstruction result: \n')
f.write("CD result: %.8f \n" % cd_distances.avg)
f.write("EMD result: %.8f \n" % emd_distances.avg)
def build_val_loss(self):
self.chamfer_dist_mean = ChamferDistanceMean().cuda()
self.emd_dist = emd.emdModule().cuda()
def build_val_loss(self):
# Set up loss functions
self.chamfer_dist = ChamferDistance().cuda()
self.chamfer_dist_mean = ChamferDistanceMean().cuda()
self.emd_dist = emd.emdModule().cuda()
def __init__(self,
cfg,
optimizer_G=None,
scheduler_G=None,
optimizer_D=None,
scheduler_D=None):
super().__init__()
self.cfg = cfg
# Graphx Generator
self.model_G = Graphx_Rec(
cfg=cfg,
in_channels=3,
in_instances=cfg.GRAPHX.NUM_INIT_POINTS,
activation=nn.ReLU(),
)
# Projection Discriminator
self.model_D = ProjectionD(
num_classes=cfg.DATASET.NUM_CLASSES,
img_shape=(cfg.RENDER.N_VIEWS + 3, cfg.RENDER.IMG_SIZE,
cfg.RENDER.IMG_SIZE),
)
# Renderer
self.renderer = ComputeDepthMaps(
projection=cfg.RENDER.PROJECTION,
eyepos_scale=cfg.RENDER.EYEPOS,
image_size=cfg.RENDER.IMG_SIZE,
).float()
# OptimizerG
self.optimizer_G = None if optimizer_G is None else optimizer_G(
self.model_G.parameters())
self.scheduler_G = None if scheduler_G or optimizer_G is None else scheduler_G(
self.optimizer_G)
# OptimizerD
self.optimizer_D = None if optimizer_D is None else optimizer_D(
self.model_D.parameters())
self.scheduler_D = None if scheduler_D or optimizer_D is None else scheduler_D(
self.optimizer_D)
# a dict store the losses for each step
self.loss = {}
# emd loss
self.emd_dist = emd.emdModule()
# GAN criterion
self.criterionD = torch.nn.MSELoss()
if torch.cuda.is_available():
# Generator
self.model_G = torch.nn.DataParallel(
self.model_G, device_ids=cfg.CONST.DEVICE).cuda()
# Discriminator
self.model_D = torch.nn.DataParallel(
self.model_D, device_ids=cfg.CONST.DEVICE).cuda()
# Renderer
self.renderer = torch.nn.DataParallel(
self.renderer, device_ids=cfg.CONST.DEVICE).cuda()
# loss
self.emd_dist = torch.nn.DataParallel(
self.emd_dist, device_ids=cfg.CONST.DEVICE).cuda()
self.criterionD = torch.nn.DataParallel(
self.criterionD, device_ids=cfg.CONST.DEVICE).cuda()
self.cuda()
def test_refine_net(cfg):
# Enable the inbuilt cudnn auto-tuner to find the best algorithm to use
torch.backends.cudnn.benchmark = True
test_transforms = utils.data_transforms.Compose([
utils.data_transforms.ToTensor(),
])
dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[cfg.DATASET.TEST_DATASET](cfg)
test_data_loader = torch.utils.data.DataLoader(dataset=dataset_loader.get_dataset(
utils.data_loaders.DatasetType.TEST, test_transforms),
batch_size=cfg.TEST.BATCH_SIZE,
num_workers=1,
pin_memory=True,
shuffle=False)
# Set up networks
# The parameters here need to be set in cfg
rec_net = Graphx_Rec(
cfg=cfg,
in_channels=3,
in_instances=cfg.GRAPHX.NUM_INIT_POINTS,
activation=nn.ReLU(),
)
# Refine network
refine_net = GRAPHX_REFINE_MODEL(
cfg=cfg,
in_channels=3,
optimizer=lambda x: torch.optim.Adam(x, lr=cfg.REFINE.LEARNING_RATE)
)
if torch.cuda.is_available():
rec_net = torch.nn.DataParallel(rec_net, device_ids=cfg.CONST.DEVICE).cuda()
refine_net = torch.nn.DataParallel(refine_net, device_ids=cfg.CONST.DEVICE).cuda()
# Load weight
# Load pretrained generator
print('[INFO] %s Recovering generator from %s ...' % (dt.now(), cfg.TEST.GENERATOR_WEIGHTS))
rec_net_dict = rec_net.state_dict()
pretrained_dict = torch.load(cfg.TEST.GENERATOR_WEIGHTS)
pretrained_weight_dict = pretrained_dict['net']
new_weight_dict = OrderedDict()
for k, v in pretrained_weight_dict.items():
if cfg.REFINE.GENERATOR_TYPE == 'REC':
name = k[21:] # remove module.reconstructor.
elif cfg.REFINE.GENERATOR_TYPE == 'GAN':
name = k[15:] # remove module.model_G.
if name in rec_net_dict:
new_weight_dict[name] = v
rec_net_dict.update(new_weight_dict)
rec_net.load_state_dict(rec_net_dict)
# Load weight
# Load weight for encoder, decoder
print('[INFO] %s Recovering refiner from %s ...' % (dt.now(), cfg.TEST.REFINER_WEIGHTS))
refine_checkpoint = torch.load(cfg.TEST.REFINER_WEIGHTS)
refine_net.load_state_dict(refine_checkpoint['net'])
print('[INFO] Best reconstruction result at epoch %d ...' % refine_checkpoint['epoch_idx'])
epoch_id = int(refine_checkpoint['epoch_idx'])
rec_net.eval()
refine_net.eval()
# Set up loss functions
emd_dist = emd.emdModule()
cd = ChamferLoss().cuda()
# Batch average meterics
cd_distances = utils.network_utils.AverageMeter()
emd_distances = utils.network_utils.AverageMeter()
n_batches = len(test_data_loader)
# Testing loop
for sample_idx, (taxonomy_names, sample_names, rendering_images, update_images,
model_azi, model_ele,
init_point_clouds, ground_truth_point_clouds) in enumerate(test_data_loader):
with torch.no_grad():
# Only one image per sample
rendering_images = torch.squeeze(rendering_images, 1)
update_images = torch.squeeze(update_images, 1)
# Get data from data loader
rendering_images = utils.network_utils.var_or_cuda(rendering_images)
update_images = utils.network_utils.var_or_cuda(update_images)
model_azi = utils.network_utils.var_or_cuda(model_azi)
model_ele = utils.network_utils.var_or_cuda(model_ele)
init_point_clouds = utils.network_utils.var_or_cuda(init_point_clouds)
ground_truth_point_clouds = utils.network_utils.var_or_cuda(ground_truth_point_clouds)
#=================================================#
# Test the encoder, decoder #
#=================================================#
# rec net give out a coarse point cloud
coarse_pc, _ = rec_net(rendering_images, init_point_clouds)
# refine net give out a refine result
loss, pred_pc = refine_net.module.valid_step(update_images, coarse_pc, ground_truth_point_clouds, model_azi, model_ele)
# Compute CD, EMD
cd_distance = cd(pred_pc, ground_truth_point_clouds) / cfg.TEST.BATCH_SIZE / cfg.CONST.NUM_POINTS
# compute reconstruction loss
emd_loss, _ = emd_dist(
pred_pc, ground_truth_point_clouds, eps=0.005, iters=50
)
emd_distance = torch.sqrt(emd_loss).mean(1).mean()
# Append loss and accuracy to average metrics
cd_distances.update(cd_distance.item())
emd_distances.update(emd_distance.item())
print("Test on [%d/%d] data, CD: %.4f EMD %.4f" % (sample_idx + 1, n_batches, cd_distance.item(), emd_distance.item()))
# print result
print("Reconstruction result:")
print("CD result: ", cd_distances.avg)
print("EMD result", emd_distances.avg)
logname = cfg.TEST.RESULT_PATH
with open(logname, 'a') as f:
f.write('Reconstruction result: \n')
f.write("CD result: %.8f \n" % cd_distances.avg)
f.write("EMD result: %.8f \n" % emd_distances.avg)
class Metrics(object):
ITEMS = [
{
"name": "F-Score",
"enabled": True,
"eval_func": "cls._get_f_score",
"is_greater_better": True,
"init_value": 0,
},
{
"name": "ChamferDistance",
"enabled": True,
"eval_func": "cls._get_chamfer_distance",
"eval_object": ChamferDistanceMean(),
"is_greater_better": False,
"init_value": 32767,
},
{
"name": "EMD",
"enabled": True,
"eval_func": "cls._get_emd",
"eval_object": emd.emdModule(),
"is_greater_better": False,
"init_value": 32767,
},
]
@classmethod
def get(cls, pred, gt):
_items = cls.items()
_values = [0] * len(_items)
for i, item in enumerate(_items):
eval_func = eval(item["eval_func"])
_values[i] = eval_func(pred, gt)
return _values
@classmethod
def items(cls):
return [i for i in cls.ITEMS if i["enabled"]]
@classmethod
def names(cls):
_items = cls.items()
return [i["name"] for i in _items]
@classmethod
def _get_f_score(cls, pred, gt, th=0.01):
"""References: https://github.com/lmb-freiburg/what3d/blob/master/util.py"""
pred = cls._get_open3d_ptcloud(pred)
gt = cls._get_open3d_ptcloud(gt)
dist1 = pred.compute_point_cloud_distance(gt)
dist2 = gt.compute_point_cloud_distance(pred)
recall = float(sum(d < th for d in dist2)) / float(len(dist2))
precision = float(sum(d < th for d in dist1)) / float(len(dist1))
return 2 * recall * precision / (
recall + precision) if recall + precision else 0
@classmethod
def _get_open3d_ptcloud(cls, tensor):
tensor = tensor.squeeze().cpu().numpy()
ptcloud = open3d.geometry.PointCloud()
ptcloud.points = open3d.utility.Vector3dVector(tensor)
return ptcloud
@classmethod
def _get_chamfer_distance(cls, pred, gt):
chamfer_distance = cls.ITEMS[1]["eval_object"]
return chamfer_distance(pred, gt).item() * 1000
@classmethod
def _get_emd(cls, pred, gt):
EMD = cls.ITEMS[2]["eval_object"]
dist, _ = EMD(pred, gt, eps=0.005, iters=50) # for val
# dist, _ = EMD(pred, gt, 0.002, 10000) # final test ?
emd = torch.sqrt(dist).mean(1).mean()
return emd.item() * 100
def __init__(self, metric_name, values):
self._items = Metrics.items()
self._values = [item["init_value"] for item in self._items]
self.metric_name = metric_name
if type(values).__name__ == "dict":
metric_indexes = {}
for idx, item in enumerate(self._items):
item_name = item["name"]
metric_indexes[item_name] = idx
for k, v in values.items():
if k not in metric_indexes:
logger.warn("Ignore Metric[Name=%s] due to disability." %
k)
continue
self._values[metric_indexes[k]] = v
elif type(values).__name__ == "list":
self._values = values
else:
raise Exception("Unsupported value type: %s" % type(values))
def state_dict(self):
_dict = {}
for i in range(len(self._items)):
item = self._items[i]["name"]
value = self._values[i]
_dict[item] = value
return _dict
def __repr__(self):
return str(self.state_dict())
def better_than(self, other):
if other is None:
return True
_index = -1
for i, _item in enumerate(self._items):
if _item["name"] == self.metric_name:
_index = i
break
if _index == -1:
raise Exception("Invalid metric name to compare.")
_metric = self._items[i]
_value = self._values[_index]
other_value = other._values[_index]
return _value > other_value if _metric[
"is_greater_better"] else _value < other_value
