def cd_margin_loss(x, y, thres=0.01, alpha=0.5):
dist1, dist2 = ChamferDistance()(x, y)
dist1, dist2 = dist1[dist1 > thres], dist2[dist2 > thres]
loss = torch.tensor(0).to(x)
if dist1.size(0) > 0:
loss += alpha * torch.mean(dist1)
if dist2.size(0) > 0:
loss += (1 - alpha) * torch.mean(dist2)
return loss
def __init__(self, config):
super(RecursiveDecoder, self).__init__()
self.conf = config
self.box_decoder = BoxDecoder(config.feature_size, config.hidden_size)
self.box_diff_decoder = BoxDiffDecoder(config.feature_size,
config.hidden_size)
self.child_decoder = ConcatChildDecoder(config.feature_size,
config.hidden_size,
config.max_child_num)
self.sample_decoder = SampleDecoder(config.feature_size,
config.hidden_size)
self.leaf_classifier = LeafClassifier(config.feature_size,
config.hidden_size)
self.bceLoss = nn.BCEWithLogitsLoss(reduction='none')
self.chamferLoss = ChamferDistance()
self.ceLoss = nn.CrossEntropyLoss(reduction='none')
self.node_diff_feature_extractor = NodeDiffFeatureExtractor(
config.feature_size, config.hidden_size)
self.node_diff_classifier = NodeDiffClassifier(config.feature_size,
config.hidden_size)
self.add_child_decoder = AddChildDecoder(config.feature_size,
config.hidden_size,
config.max_child_num)
self.register_buffer('unit_cube',
torch.from_numpy(load_pts('cube.pts')))
def __init__(self, config):
super(RecursiveDecoder, self).__init__()
self.conf = config
self.node_decoder = NodeDecoder(config.geo_feat_size, config.feature_size, config.num_point)
self.child_decoder = GNNChildDecoder(
node_feat_size=config.feature_size,
hidden_size=config.hidden_size,
max_child_num=config.max_child_num,
edge_symmetric_type=config.edge_symmetric_type,
num_iterations=config.num_dec_gnn_iterations,
edge_type_num=len(config.edge_types))
self.sample_decoder = SampleDecoder(config.feature_size, config.hidden_size)
self.leaf_classifier = LeafClassifier(config.feature_size, config.hidden_size)
self.bceLoss = nn.BCEWithLogitsLoss(reduction='none')
self.chamferLoss = ChamferDistance()
self.semCELoss = nn.CrossEntropyLoss(reduction='none')
self.mseLoss = nn.MSELoss(reduction='none')
self.register_buffer('unit_cube', torch.from_numpy(load_pts('cube.pts')))
def cd(x, y):
# x = x.permute(0, 2, 1)
#y = y.permute(0, 2, 1)
d1, d2 = ChamferDistance()(x, y)
loss=torch.sum(d1, dim=1) + torch.sum(d2, dim=1)
print(loss)
return torch.mean(loss,dim=0)
def __init__(self, feat_len, num_point):
super(PartDecoder, self).__init__()
self.num_point = num_point
self.mlp1 = nn.Linear(feat_len, feat_len)
self.mlp2 = nn.Linear(feat_len, feat_len)
self.mlp3 = nn.Linear(feat_len, num_point * 3)
self.bn1 = nn.BatchNorm1d(feat_len)
self.bn2 = nn.BatchNorm1d(feat_len)
self.chamferLoss = ChamferDistance()
def get_simplification_loss(self,
ref_pc,
samp_pc,
pc_size,
gamma=1,
delta=0):
if self.skip_projection or not self.training:
return torch.tensor(0).to(ref_pc)
# ref_pc and samp_pc are B x N x 3 matrices
cost_p1_p2, cost_p2_p1 = ChamferDistance()(samp_pc, ref_pc)
max_cost = torch.max(cost_p1_p2, dim=1)[0] # furthest point
max_cost = torch.mean(max_cost)
cost_p1_p2 = torch.mean(cost_p1_p2)
cost_p2_p1 = torch.mean(cost_p2_p1)
loss = cost_p1_p2 + max_cost + (gamma + delta * pc_size) * cost_p2_p1
return loss
def __init__(self, args):
super(PCRNet, self).__init__()
self.emb_dims = args.emb_dims
self.cycle = args.cycle
self.iterations = args.iterations
if args.emb_nn == 'pointnet':
self.emb_nn = PointNet(emb_dims=self.emb_dims)
elif args.emb_nn == 'dgcnn':
self.emb_nn = DGCNN(emb_dims=self.emb_dims)
else:
raise Exception('Not implemented')
if args.pointer == 'identity':
self.pointer = Identity()
else:
raise Exception("Not implemented")
if args.head == 'mlp':
self.head = MLPHead(args=args)
else:
raise Exception('Not implemented')
self.chamfer = ChamferDistance()
matplotlib.use('agg')
import matplotlib.pyplot as plt
import torch.nn as nn
from tqdm import tqdm
from voxel import *
from torch.utils.data import DataLoader
from PIL import Image
from torchvision.transforms import Normalize as norm
from torchvision import transforms
preprocess = transforms.Compose(
[transforms.Resize((224, 224)),
transforms.ToTensor()])
from chamfer_distance import ChamferDistance
chamfer_dist = ChamferDistance()
from tri_distance import TriDistance
tri_dist = TriDistance()
# loads the initial mesh and stores vertex, face, and adjacency matrix information
def load_initial(obj='386.obj'):
# load obj file
obj = ObjLoader(obj)
labels = np.array(obj.vertices)
features = torch.FloatTensor(labels).cuda()
faces = torch.LongTensor(np.array(obj.faces) - 1).cuda()
points = torch.rand([1000, 3]).cuda() - .5
verts = features.clone()
def main():
# Parse the training argument
parser = argparse.ArgumentParser()
parser.add_argument(
'--training',
help=
'Decide whether to train the model or just run testing on previously saved model.'
)
args = parser.parse_args()
is_training = args.training
if (is_training is None) or (is_training == 'True'):
is_training = True
else:
is_training = False
print('is_training mode = ', is_training)
chamferDist = ChamferDistance()
# Decide on GPU or CPU
if torch.cuda.is_available():
gpu_or_cpu = torch.device('cuda')
else:
gpu_or_cpu = torch.device('cpu')
# Training Configuration
# image_root = "./../../../datasets/cs253-wi20-public/ShapeNetRendering/"
# point_cloud_root = "./../../../datasets/cs253-wi20-public/ShapeNet_pointclouds/"
image_root = "/datasets/cs253-wi20-public/ShapeNetRendering/"
point_cloud_root = "/datasets/cs253-wi20-public/ShapeNet_pointclouds/"
num_epochs = 1000
batch_size = 64
shuffle = True
num_workers = 8
use_2048 = True
img_size = 227 # I don't know why, but this has to be 227!
learning_rate = 1e-4
num_points = 2048
transform = transforms.Compose([
transforms.Resize(img_size, interpolation=2),
transforms.CenterCrop(img_size),
transforms.ToTensor()
])
# Checkpoint
use_checkpoint = False
# Split and Get data. Override the saved files if you change the ratios.
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1
split_data(train_ratio, val_ratio, test_ratio, overrideFiles=False)
path_train = 'train_data.txt'
path_val = 'val_data.txt'
path_test = 'test_data.txt'
train_data = read_from_file(path_train)
val_data = read_from_file(path_val)
test_data = read_from_file(path_test)
# Data loader
train_data_loader = get_loader(image_root, point_cloud_root, train_data,
use_2048, transform, batch_size, shuffle,
num_workers)
val_data_loader = get_loader(image_root, point_cloud_root, val_data,
use_2048, transform, batch_size, shuffle,
num_workers)
test_data_loader = get_loader(image_root, point_cloud_root, test_data,
use_2048, transform, batch_size, shuffle,
num_workers)
print('Len of train loader = ', len(train_data_loader))
# create model
print("model building...")
model = pic2points(num_points=num_points)
model.to(device=gpu_or_cpu)
if is_training:
# Train
print('Starting training...')
train_losses, val_loss, best_model = train(
model,
train_data_loader,
val_data_loader,
chamferDist,
model_name="Baseline_DL_Vis",
num_epochs=num_epochs,
lr=learning_rate,
use_checkpoint=use_checkpoint)
else:
best_model = torch.load('best-Baseline_DL_Vis.pt')
print('Loaded previously saved model.')
model = best_model.cuda()
model.eval()
# Compute chamfer distance on Pix3D dataset.
img_path = "/datasets/cs253-wi20-public/pix3d/"
pc_path = "/datasets/cs253-wi20-public/pix_pointclouds/"
objects = ['table', 'sofa']
test_dataset = TestDataset(img_path, pc_path, objects)
test_data_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=1,
shuffle=True,
num_workers=8)
print('Starting testing on Pix3D dataset...')
total_test_loss = 0.
# Get loss on training data.
with torch.no_grad():
for i, (image, point_cloud) in enumerate(test_data_loader):
image, point_cloud = Variable(image), Variable(point_cloud)
# print(image.size())
if (image.size(1) != 3):
continue
# print('reaching.')
image, point_cloud = image.float().to(
device=gpu_or_cpu), point_cloud.float().to(device=gpu_or_cpu)
pred = model(image)
dist1, dist2 = chamferDist(pred, point_cloud)
loss = (torch.mean(dist1)) + (torch.mean(dist2))
# emd_cost = torch.sum(dist(pred.cuda().double(), points.cuda().double()))
total_test_loss += loss.item()
# print(total_test_loss)
# break
if i % 100 == 0:
print('Batch ' + str(i) + ' finished.')
print('Chamfer distance on Pix3D dataset = ',
total_test_loss / len(test_data_loader))
data = data.to(device)
optimizer.zero_grad()
decoded, _ = model(data)
dist1, dist2 = criterion(decoded.reshape(-1, 2048, 3),
data.y.reshape(-1, 2048, 3))
loss = (torch.mean(dist1)) + (torch.mean(dist2))
loss.backward()
total_loss += loss.item() * data.num_graphs
optimizer.step()
return total_loss / len(dataset)
if __name__ == '__main__':
dataset = Completion3D('../data/Completion3D',
split='train',
categories='Airplane')
print(dataset[0])
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SaNet().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
print(model)
print('Training started:')
criterion = ChamferDistance()
for epoch in range(1, 401):
loss = train()
print('Epoch {:03d}, Loss: {:.4f}'.format(epoch, loss))
if epoch % 10 == 0:
torch.save(model.state_dict(),
'./trained/SA_net_Ch' + '{}'.format(epoch) + '.pt')
def dist_cd(pc1, pc2): chamfer_dist = ChamferDistance() dist1, dist2 = chamfer_dist(pc1, pc2) return torch.mean(dist1) + torch.mean(dist2)
def dist_cd(pc1, pc2):
chamfer_dist = ChamferDistance()
dist1, dist2 = chamfer_dist(pc1, pc2)
return (torch.mean(torch.sqrt(dist1)) + torch.mean(torch.sqrt(dist2)))/2
def compute_chamfer_loss(gt_points, reconstruct_points):
chamfer_dist = ChamferDistance()
dist1, dist2 = chamfer_dist(gt_points, reconstruct_points)
loss = torch.mean(dist1, dim=-1) + torch.mean(dist2, dim=-1)
return loss
def __init__(self):
super(ChamfersDistance, self).__init__()
self.chamfer_dist = ChamferDistance()
def main(args):
"""
for each mesh, render 4-view-normals and get mesh-voxels, also compute and save {3D-IoU, Normal errors}
"""
# flags for visual sanity check
visualCheck_0 = False
# init.
rn = ColoredRenderer()
rn.camera = ProjectPointsOrthogonal(rt=np.array([0,0,0]), t=np.array([0,0,2]), f=np.array([consts.dim_h*2,consts.dim_h*2]), c=np.array([consts.dim_w,consts.dim_h]), k=np.zeros(5))
rn.frustum = {'near': 0.5, 'far': 25, 'height': consts.dim_h*2, 'width': consts.dim_w*2}
chamfer_dist = ChamferDistance()
# get training/test data indices
training_inds, testing_inds = get_training_test_indices(args=args,shuffle=False)
meshRefinedPath_list = compute_split_range(testing_inds=testing_inds,args=args)
# for each mesh, render 4-view-normals and get mesh-voxels, also compute and save {3D-IoU, Normal errors}
frameIdx = [0, 0, 0]
frameIdx[0] = int( meshRefinedPath_list[0].split("/")[-1].split("_meshRefined")[0])
frameIdx[2] = int(meshRefinedPath_list[-1].split("/")[-1].split("_meshRefined")[0])+1
count = 0
timeStart = time.time()
for meshPath in meshRefinedPath_list:
# init.
frameIdx[1] = int(meshPath.split("/")[-1].split("_meshRefined")[0])
evalMetricsPath = "%s/%06d_evalMetrics.json" % (args.resultsDir, frameIdx[1])
evalMetricsPath_Next = "%s/%06d_evalMetrics.json" % (args.resultsDir, frameIdx[1]+1)
evalMetricsPath_additional = "%s/%06d_evalMetrics_additional.json" % (args.resultsDir, frameIdx[1])
evalMetricsPath_Next_additional = "%s/%06d_evalMetrics_additional.json" % (args.resultsDir, frameIdx[1]+1)
if os.path.exists(evalMetricsPath) and os.path.exists(evalMetricsPath_Next) and os.path.exists(evalMetricsPath_additional) and os.path.exists(evalMetricsPath_Next_additional):
continue
# ----- compute point based distance metrics -----
if True:
# note that the losses have been multiplied by "scale"
chamfer_dis, estV_2_gtM_dis, estMesh = compute_point_based_metrics(args=args,estMeshPath=meshPath,preFix=frameIdx[1],chamfer_dist=chamfer_dist,scale=10000.)
# ----- save the additional eval metrics into .json of args.resultsDir dir -----
if True:
evalMetrics_additional = {"chamfer_dis" : chamfer_dis,
"estV_2_gtM_dis": estV_2_gtM_dis}
with open(evalMetricsPath_additional, "w") as outfile:
json.dump(evalMetrics_additional, outfile)
if visualCheck_0:
print("check eval metrics additional json results...")
print(evalMetrics_additional)
os.system("cp %s ./examples/%06d_evalPrepare_metrics_additional.json" % (evalMetricsPath_additional,frameIdx[1]))
pdb.set_trace()
# ----- render front-view-normal & compute normal errors of [cos-dis, l2-dis] -----
if not args.only_compute_additional_metrics:
normal_errors = compute_n_save_normal_erros(args=args,estMeshPath=meshPath,rn=rn,preFix=frameIdx[1],estMesh=estMesh)
assert(len(normal_errors) == 1 and len(normal_errors[0]) == 2)
# ----- save the eval metrics into .json of args.resultsDir dir -----
if not args.only_compute_additional_metrics:
evalMetrics = {"norm_cos_dis_ft": np.array([normal_errors[0][0]]).tolist(),
"norm_l2_dis_ft": np.array([normal_errors[0][1]]).tolist()}
with open(evalMetricsPath, 'w') as outfile:
json.dump(evalMetrics, outfile)
visualCheck = False
if visualCheck:
print("check eval metrics json results...")
print(evalMetrics)
os.system("cp %s ./examples/%06d_evalPrepare_metrics.json" % (evalMetricsPath,frameIdx[1]))
pdb.set_trace()
# compute timing info
count += 1
hrsPassed = (time.time()-timeStart) / 3600.
hrsEachIter = hrsPassed / count
numItersRemain = len(meshRefinedPath_list) - count
hrsRemain = numItersRemain * hrsEachIter # hours that remain
minsRemain = hrsRemain * 60. # minutes that remain
# log
expName = args.resultsDir.split("/")[-1]
print("Exp. %s inference: split %d/%d | frameIdx %06d-%06d-%06d | remains %.3f m(s) ......" % (expName,args.splitIdx,args.splitNum,frameIdx[0],frameIdx[1],frameIdx[2],minsRemain))
def cd(x, y):
x = x.permute(0, 2, 1)
y = y.permute(0, 2, 1)
d1, d2 = ChamferDistance()(x, y)
return torch.sum(d1, dim=1) + torch.sum(d2, dim=1)
def compute_gen_sd_numbers(in_dir, data_path, object_list, shapediff_topk,
shapediff_metric, self_is_neighbor, tot_shape):
chamfer_loss = ChamferDistance()
unit_cube = torch.from_numpy(utils.load_pts('cube.pts'))
def box_dist(box_feature, gt_box_feature):
pred_box_pc = utils.transform_pc_batch(unit_cube, box_feature)
pred_reweight = utils.get_surface_reweighting_batch(
box_feature[:, 3:6], unit_cube.size(0))
gt_box_pc = utils.transform_pc_batch(unit_cube, gt_box_feature)
gt_reweight = utils.get_surface_reweighting_batch(
gt_box_feature[:, 3:6], unit_cube.size(0))
dist1, dist2 = chamfer_loss(gt_box_pc, pred_box_pc)
loss1 = (dist1 * gt_reweight).sum(dim=1) / (gt_reweight.sum(dim=1) +
1e-12)
loss2 = (dist2 *
pred_reweight).sum(dim=1) / (pred_reweight.sum(dim=1) + 1e-12)
loss = (loss1 + loss2) / 2
return loss
def struct_dist(gt_node, pred_node):
if gt_node.is_leaf:
if pred_node.is_leaf:
return 0
else:
return len(pred_node.boxes()) - 1
else:
if pred_node.is_leaf:
return len(gt_node.boxes()) - 1
else:
gt_sem = set([node.label for node in gt_node.children])
pred_sem = set([node.label for node in pred_node.children])
intersect_sem = set.intersection(gt_sem, pred_sem)
gt_cnodes_per_sem = dict()
for node_id, gt_cnode in enumerate(gt_node.children):
if gt_cnode.label in intersect_sem:
if gt_cnode.label not in gt_cnodes_per_sem:
gt_cnodes_per_sem[gt_cnode.label] = []
gt_cnodes_per_sem[gt_cnode.label].append(node_id)
pred_cnodes_per_sem = dict()
for node_id, pred_cnode in enumerate(pred_node.children):
if pred_cnode.label in intersect_sem:
if pred_cnode.label not in pred_cnodes_per_sem:
pred_cnodes_per_sem[pred_cnode.label] = []
pred_cnodes_per_sem[pred_cnode.label].append(node_id)
matched_gt_idx = []
matched_pred_idx = []
matched_gt2pred = np.zeros((100), dtype=np.int32)
for sem in intersect_sem:
gt_boxes = torch.cat([
gt_node.children[cid].get_box_quat()
for cid in gt_cnodes_per_sem[sem]
],
dim=0)
pred_boxes = torch.cat([
pred_node.children[cid].get_box_quat()
for cid in pred_cnodes_per_sem[sem]
],
dim=0)
num_gt = gt_boxes.size(0)
num_pred = pred_boxes.size(0)
if num_gt == 1 and num_pred == 1:
cur_matched_gt_idx = [0]
cur_matched_pred_idx = [0]
else:
gt_boxes_tiled = gt_boxes.unsqueeze(dim=1).repeat(
1, num_pred, 1)
pred_boxes_tiled = pred_boxes.unsqueeze(dim=0).repeat(
num_gt, 1, 1)
dmat = box_dist(gt_boxes_tiled.view(-1, 10),
pred_boxes_tiled.view(-1, 10)).view(
-1, num_gt, num_pred)
_, cur_matched_gt_idx, cur_matched_pred_idx = utils.linear_assignment(
dmat)
for i in range(len(cur_matched_gt_idx)):
matched_gt_idx.append(
gt_cnodes_per_sem[sem][cur_matched_gt_idx[i]])
matched_pred_idx.append(
pred_cnodes_per_sem[sem][cur_matched_pred_idx[i]])
matched_gt2pred[gt_cnodes_per_sem[sem][
cur_matched_gt_idx[i]]] = pred_cnodes_per_sem[sem][
cur_matched_pred_idx[i]]
struct_diff = 0.0
for i in range(len(gt_node.children)):
if i not in matched_gt_idx:
struct_diff += len(gt_node.children[i].boxes())
for i in range(len(pred_node.children)):
if i not in matched_pred_idx:
struct_diff += len(pred_node.children[i].boxes())
for i in range(len(matched_gt_idx)):
gt_id = matched_gt_idx[i]
pred_id = matched_pred_idx[i]
struct_diff += struct_dist(gt_node.children[gt_id],
pred_node.children[pred_id])
return struct_diff
# create dataset and data loader
data_features = [
'object', 'name', 'neighbor_diffs', 'neighbor_objs', 'neighbor_names'
]
dataset = PartNetShapeDiffDataset(data_path, object_list, data_features,
shapediff_topk, shapediff_metric,
self_is_neighbor)
tot_gen = 100
bar = ProgressBar()
quality = 0.0
coverage = 0.0
for i in bar(range(tot_shape)):
obj, obj_name, neighbor_diffs, neighbor_objs, neighbor_names = dataset[
i]
mat1 = np.zeros((shapediff_topk, tot_gen), dtype=np.float32)
mat2 = np.zeros((shapediff_topk, tot_gen), dtype=np.float32)
for j in range(tot_gen):
gen_obj = PartNetDataset.load_object(
os.path.join(in_dir, obj_name, 'obj2-%03d.json' % j))
for ni in range(shapediff_topk):
sd = struct_dist(neighbor_objs[ni].root, gen_obj.root)
mat1[ni, j] = sd / len(neighbor_objs[ni].root.boxes())
mat2[ni, j] = sd / len(gen_obj.root.boxes())
quality += mat2.min(axis=0).mean()
coverage += mat1.min(axis=1).mean()
np.save(os.path.join(in_dir, obj_name, 'sd_mat1_stats.npy'), mat1)
np.save(os.path.join(in_dir, obj_name, 'sd_mat2_stats.npy'), mat2)
quality /= tot_shape
coverage /= tot_shape
print('mean sd quality: ', quality)
print('mean sd coverage: ', coverage)
print('q + c: %.5f' % (quality + coverage))
with open(
os.path.join(in_dir,
'neighbor_%s_sd_stats.txt' % shapediff_metric),
'w') as fout:
fout.write('mean sd quality: %f\n' % quality)
fout.write('mean sd coverage: %f\n' % coverage)
fout.write('q + c: %.5f\n' % (quality + coverage))
sampler = models.groundtruthSampler(camNum=opt.camNum,
fov=opt.fov,
imHeight=opt.imageHeight,
imWidth=opt.imageWidth,
isNoRenderError=opt.isNoRenderError)
# Define the model and optimizer
lr_scale = 1
pointNet = PointNetRefinePoint()
if opt.isFineTune:
pointNet.load_state_dict(
torch.load('%s/pointNet_%d.pth' %
(opt.experiment, opt.epochIdFineTune)))
lr_scale = np.power(0.5, int((opt.epochIdFineTune + 1) / 2.0))
chamferDist = ChamferDistance()
if opt.cuda:
pointNet = pointNet.cuda()
chamferDist = chamferDist.cuda()
opPointNet = torch.optim.Adam(pointNet.parameters(),
lr=1e-4 * lr_scale,
betas=(0.5, 0.999))
j = 0
pointErrsNpList = np.ones([1, 2], dtype=np.float32)
normalErrsNpList = np.ones([1, 2], dtype=np.float32)
meanAngleErrsNpList = np.ones([1, 2], dtype=np.float32)
medianAngleErrsNpList = np.ones([1, 2], dtype=np.float32)
for epoch in list(range(opt.epochIdFineTune + 1, opt.nepoch)):
def compute_gen_cd_numbers(in_dir, data_path, object_list, shapediff_topk,
shapediff_metric, self_is_neighbor, tot_shape):
chamfer_loss = ChamferDistance()
data_features = [
'object', 'name', 'neighbor_diffs', 'neighbor_objs', 'neighbor_names'
]
dataset = PartNetShapeDiffDataset(data_path, object_list, data_features,
shapediff_topk, shapediff_metric,
self_is_neighbor)
tot_gen = 100
bar = ProgressBar()
quality = 0.0
coverage = 0.0
for i in bar(range(tot_shape)):
obj, obj_name, neighbor_diffs, neighbor_objs, neighbor_names = dataset[
i]
mat = np.zeros((shapediff_topk, tot_gen), dtype=np.float32)
gt_pcs = []
for ni in range(shapediff_topk):
obbs_np = torch.cat([
item.view(1, -1)
for item in neighbor_objs[ni].boxes(leafs_only=True)
],
dim=0).cpu().numpy()
mesh_v, mesh_f = utils.gen_obb_mesh(obbs_np)
pc_sample = utils.sample_pc(mesh_v, mesh_f)
gt_pcs.append(np.expand_dims(pc_sample, axis=0))
gt_pcs = np.concatenate(gt_pcs, axis=0)
gt_pcs = torch.from_numpy(gt_pcs).float().cuda()
for i in range(tot_gen):
obj = PartNetDataset.load_object(
os.path.join(in_dir, obj_name, 'obj2-%03d.json' % i))
obbs_np = torch.cat(
[item.view(1, -1) for item in obj.boxes(leafs_only=True)],
dim=0).cpu().numpy()
mesh_v, mesh_f = utils.gen_obb_mesh(obbs_np)
gen_pc = utils.sample_pc(mesh_v, mesh_f)
gen_pc = np.tile(np.expand_dims(gen_pc, axis=0),
[shapediff_topk, 1, 1])
gen_pc = torch.from_numpy(gen_pc).float().cuda()
d1, d2 = chamfer_loss(gt_pcs.cuda(), gen_pc)
mat[:, i] = (d1.sqrt().mean(dim=1) +
d2.sqrt().mean(dim=1)).cpu().numpy() / 2
quality += mat.min(axis=0).mean()
coverage += mat.min(axis=1).mean()
np.save(os.path.join(in_dir, obj_name, 'cd_stats.npy'), mat)
quality /= tot_shape
coverage /= tot_shape
print('mean cd quality: %.5f' % quality)
print('mean cd coverage: %.5f' % coverage)
print('q + c: %.5f' % (quality + coverage))
with open(
os.path.join(in_dir,
'neighbor_%s_cd_stats.txt' % shapediff_metric),
'w') as fout:
fout.write('mean cd quality: %.5f\n' % quality)
fout.write('mean cd coverage: %.5f\n' % coverage)
fout.write('q + c: %.5f\n' % (quality + coverage))
def cd_loss(x, y, alpha=0.5):
dist1, dist2 = ChamferDistance()(x, y)
return alpha * torch.mean(dist1) + (1 - alpha) * torch.mean(dist2)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.abspath(os.path.join(BASE_DIR, '../../')))
sys.path.append(os.path.abspath(os.path.join(BASE_DIR, '../../dataloaders')))
import shapenet_part_loader
import shapenet_core13_loader
import shapenet_core55_loader
from model import PointCapsNet
import segmentation as seg
import open3d as o3d
import matplotlib.pyplot as plt
from chamfer_distance import ChamferDistance
CD = ChamferDistance()
## MONKEY PATCHING
PointCloud = o3d.geometry.PointCloud
Vector3dVector = o3d.utility.Vector3dVector
draw_geometries = o3d.visualization.draw_geometries
viz = o3d.visualization.Visualizer()
image_id = 0
USE_CUDA = True
def show_points(points_tensor):
#print("showing tensor of shape", points_tensor.size())
prc_r_all=points_tensor.transpose(1, 0).contiguous().data.cpu()
prc_r_all_point=PointCloud()
def main():
chamferDist = ChamferDistance()
# Decide on GPU or CPU
if torch.cuda.is_available():
gpu_or_cpu = torch.device('cuda')
else:
gpu_or_cpu = torch.device('cpu')
image_root = "/datasets/cs253-wi20-public/ShapeNetRendering/"
point_cloud_root = "/datasets/cs253-wi20-public/ShapeNet_pointclouds/"
num_epochs = 20
batch_size = 64
shuffle = True
num_workers = 8
use_2048 = True
img_size = 227 # I don't know why, but this has to be 227!
learning_rate = 5e-4
num_points = 2500
transform = transforms.Compose([
transforms.Resize(img_size, interpolation=2),
transforms.CenterCrop(img_size),
transforms.ToTensor()
])
path_test = 'test_data.txt'
test_data = read_from_file(path_test)
test_data_loader = get_loader(image_root, point_cloud_root, test_data,
use_2048, transform, batch_size, shuffle,
num_workers)
model = torch.load('best-Baseline_FixedDL.pt').to(device=gpu_or_cpu)
model.eval()
for i, (image, point_cloud) in enumerate(test_data_loader):
image, point_cloud = Variable(image, requires_grad=False), Variable(
point_cloud, requires_grad=False)
image, point_cloud = image.float().to(
device=gpu_or_cpu), point_cloud.float().to(device=gpu_or_cpu)
pred = model(image)
dist1, dist2 = chamferDist(pred, point_cloud)
loss = (torch.mean(dist1)) + (torch.mean(dist2))
print('pred size = ', pred.size())
pred = pred.to('cpu')
out = []
for p in pred:
out.append(p.detach().numpy())
print('type(out) = ', type(out))
# Visualize the prediction
print('asdf')
Visualize(out).ShowRandom()
break
vid += v.shape[0]
all_v = np.vstack(all_v)
all_f = np.vstack(all_f)
return all_v, all_f
def sample_pc(v, f, n_points=2048):
mesh = trimesh.Trimesh(vertices=v, faces=f - 1)
points, __ = trimesh.sample.sample_surface(mesh=mesh, count=n_points)
return points
# load unit cube pc
from chamfer_distance import ChamferDistance
chamferLoss = ChamferDistance()
unit_cube = torch.from_numpy(load_pts('cube.pts'))
def boxLoss(box_feature, gt_box_feature):
global unit_cube
if unit_cube.device != box_feature.device:
unit_cube = unit_cube.to(box_feature.device)
pred_box_pc = transform_pc_batch(unit_cube, box_feature)
pred_reweight = get_surface_reweighting_batch(box_feature[:, 3:6],
unit_cube.size(0))
gt_box_pc = transform_pc_batch(unit_cube, gt_box_feature)
gt_reweight = get_surface_reweighting_batch(gt_box_feature[:, 3:6],
unit_cube.size(0))
dist1, dist2 = chamferLoss(gt_box_pc, pred_box_pc)
import os
import sys
from argparse import ArgumentParser
import numpy as np
import torch
from progressbar import ProgressBar
from chamfer_distance import ChamferDistance
from config import add_eval_args
from data import PartNetDataset, PartNetShapeDiffDataset, Tree
from eval_utils import compute_recon_numbers
import utils
import shutil
sys.setrecursionlimit(5000) # this code uses recursion a lot for code simplicity
chamfer_loss = ChamferDistance()
parser = ArgumentParser()
parser = add_eval_args(parser)
parser.add_argument('--start_id', type=int, default=0)
parser.add_argument('--end_id', type=int, default=-1)
parser.add_argument('--baseline_dir', type=str, help='structurenet baseline result directory')
eval_conf = parser.parse_args()
# load train config
conf = torch.load(os.path.join(eval_conf.ckpt_path, eval_conf.exp_name, 'conf.pth'))
eval_conf.category = conf.category
eval_conf.data_path = conf.data_path
if hasattr(conf, 'self_is_neighbor'):
eval_conf.self_is_neighbor = conf.self_is_neighbor
