def save_result(self, save_path):
save_img_path = os.path.join(save_path, 'imgs/')
if not os.path.exists(save_img_path):
os.makedirs(save_img_path)
save_pc_path = os.path.join(save_path, 'pcs/')
if not os.path.exists(save_pc_path):
os.makedirs(save_pc_path)
# copy the imgs from cache
copy_tree(cfg.INFERENCE.CACHE_IMAGE_PATH, save_img_path)
# copy the pcs from cache
copy_tree(cfg.INFERENCE.CACHE_POINT_CLOUD_PATH, save_pc_path)
# save edited base point clouds
output_point_cloud_ply(np.array([GM.base_point_cloud.positions]),
['final_base'], save_pc_path)
# save edited work plane point clouds
self.work_plane_sketch_manager.save_work_plane_generate_point_clouds(
save_pc_path)
"""
# save counter info for the user study
self.save_counter(save_path)
reset_counter()
"""
# reset geometry manager
self.reset_manager()
self.update()
def save_merge_point_clouds(self, save_pc_path):
merge_pcs = list(GM.base_point_cloud.positions)
for work_plane_id, work_plane in enumerate(GM.work_planes):
for pc_id, pc in enumerate(work_plane.generate_point_clouds):
merge_pcs.append(list(pc.positions))
print(np.array(merge_pcs).shape)
output_point_cloud_ply(np.array([merge_pcs]), ['final_merge'], save_pc_path)
def symmetric_optimize(self):
self.makeCurrent()
if len(GM.work_planes) == 0:
ori_pcs = GM.base_point_cloud.positions
else:
ori_pcs_list = []
ori_pcs_list.append(GM.base_point_cloud.positions)
for work_plane in GM.work_planes:
for point_cloud in work_plane.generate_point_clouds:
pc = point_cloud.positions
ori_pcs_list.append(pc)
ori_pcs = np.concatenate(ori_pcs_list)
optimize_pc = symmetric_optimize_net(cfg.INFERENCE.OPTIMIZE_PATH,
ori_pcs, [0., 0., 0.],
[0., 1., 0.])
output_point_cloud_ply(np.array([optimize_pc]), ['optimize_pc'],
cfg.INFERENCE.CACHE_POINT_CLOUD_PATH)
def inference_net(cfg, upload_image=False):
# Enable the inbuilt cudnn auto-tuner to find the best algorithm to use
torch.backends.cudnn.benchmark = True
inference_transforms = utils.data_transforms.Compose([
utils.data_transforms.ToTensor(),
])
# 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(),
)
if torch.cuda.is_available():
rec_net = torch.nn.DataParallel(rec_net,
device_ids=cfg.CONST.DEVICE).cuda()
# Load pretrained generator
print('[INFO] %s Recovering generator from %s ...' %
(dt.now(), cfg.INFERENCE.GENERATOR_WEIGHTS))
rec_net_dict = rec_net.state_dict()
pretrained_dict = torch.load(cfg.INFERENCE.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)
rec_net.eval()
input_image = []
if upload_image:
# upload image from the file system
# TO DO
raise NotImplementedError
else:
# user draw image on the UI
screenshot_image_mask = get_green_mask(
cfg.INFERENCE.SCREENSHOT_IMAGE_PATH)
screenshot_image_mask = cv2.bitwise_not(screenshot_image_mask)
screenshot_image_mask = cv2.resize(screenshot_image_mask, (224, 224),
interpolation=cv2.INTER_AREA)
# convert all the pixel > 0 to 255
screenshot_image_mask[screenshot_image_mask < 255] = 0
sketch_image = cv2.cvtColor(screenshot_image_mask, cv2.COLOR_GRAY2RGB)
cv2.imwrite(cfg.INFERENCE.SKETCH_IMAGE_PATH, sketch_image)
input_image_path = cfg.INFERENCE.SKETCH_IMAGE_PATH
input_image = cv2.imread(input_image_path, cv2.IMREAD_COLOR).astype(
np.float32) / 255.
samples = []
samples.append(input_image)
samples = np.array(samples).astype(np.float32)
input_images = inference_transforms(rendering_images=samples)
# load init point clouds
init_point_cloud_np = init_pointcloud_loader(cfg.GRAPHX.NUM_INIT_POINTS)
init_point_clouds = np.array([init_point_cloud_np])
init_point_clouds = torch.from_numpy(init_point_clouds)
with torch.no_grad():
# generate rough point cloud
coarse_pc, _ = rec_net(input_images, init_point_clouds)
output_point_cloud_ply(coarse_pc.detach().cpu().numpy(), ['coarse'],
cfg.INFERENCE.CACHE_POINT_CLOUD_PATH)
return coarse_pc[0].detach().cpu().numpy()
def evaluate_refine_net(cfg):
# Enable the inbuilt cudnn auto-tuner to find the best algorithm to use
torch.backends.cudnn.benchmark = True
eval_transforms = utils.data_transforms.Compose([
utils.data_transforms.ToTensor(),
])
dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[
cfg.DATASET.TEST_DATASET](cfg)
eval_data_loader = torch.utils.data.DataLoader(
dataset=dataset_loader.get_dataset(utils.data_loaders.DatasetType.TEST,
eval_transforms),
batch_size=cfg.EVALUATE.BATCH_SIZE,
num_workers=1,
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.EVALUATE.GENERATOR_WEIGHTS))
rec_net_dict = rec_net.state_dict()
pretrained_dict = torch.load(cfg.EVALUATE.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.EVALUATE.REFINER_WEIGHTS))
refine_checkpoint = torch.load(cfg.EVALUATE.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()
# Testing loop
for sample_idx, (taxonomy_names, sample_names, rendering_images,
update_images, model_x, model_y, init_point_clouds,
ground_truth_point_clouds) in enumerate(eval_data_loader):
print("evaluate sample: ", sample_idx)
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_x = utils.network_utils.var_or_cuda(model_x)
model_y = utils.network_utils.var_or_cuda(model_y)
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 network #
#=================================================#
# 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_x,
model_y)
img_dir = cfg.EVALUATE.OUTPUT_FOLDER
azi = model_x[0].detach().cpu().numpy() * 180. / np.pi
ele = model_y[0].detach().cpu().numpy() * 180. / np.pi + 90.
sample_name = sample_names[0]
taxonomy_name = taxonomy_names[0]
# Save sample dir
output_dir = cfg.EVALUATE.OUTPUT_FOLDER
sample_dir = os.path.join(output_dir, str(sample_idx))
if not os.path.exists(sample_dir):
os.makedirs(sample_dir)
sample_name_dir = os.path.join(sample_dir, sample_name)
if not os.path.exists(sample_name_dir):
os.makedirs(sample_name_dir)
# Rendering image
sketch_path = os.path.join(sample_name_dir, 'sketch')
if not os.path.exists(sketch_path):
os.makedirs(sketch_path)
src_sketch_img_path = os.path.join(
cfg.DATASETS.SHAPENET.RENDERING_PATH %
(taxonomy_name, sample_name, 1))
copyfile(src_sketch_img_path,
os.path.join(sketch_path, 'sketch.png'))
# Predict Pointcloud
p_pc = pred_pc[0].detach().cpu().numpy()
rendering_views = utils.point_cloud_visualization_old.get_point_cloud_image(
p_pc,
os.path.join(sample_name_dir, 'rec results'),
sample_idx,
cfg.EVALUATE.VERSION_ID,
"",
view=[azi, ele])
# Groundtruth Pointcloud
gt_pc = ground_truth_point_clouds[0].detach().cpu().numpy()
rendering_views = utils.point_cloud_visualization_old.get_point_cloud_image(
gt_pc,
os.path.join(sample_name_dir, 'gt'),
sample_idx,
cfg.EVALUATE.VERSION_ID,
"",
view=[azi, ele])
# save ply file
pc_dir = os.path.join(sample_name_dir, 'pc')
if not os.path.exists(pc_dir):
os.makedirs(pc_dir)
output_point_cloud_ply(np.array([p_pc]), ['pred_pc'], pc_dir)
if sample_idx == 200:
break
def save_work_plane_generate_point_clouds(self, save_pc_path):
for work_plane_id, work_plane in enumerate(GM.work_planes):
for pc_id, pc in enumerate(work_plane.generate_point_clouds):
output_point_cloud_ply(np.array([pc.positions]), ['part' + str(work_plane_id) + '_' + str(pc_id)], save_pc_path)
def evaluate_rec_net(cfg):
# Enable the inbuilt cudnn auto-tuner to find the best algorithm to use
torch.backends.cudnn.benchmark = True
eval_transforms = utils.data_transforms.Compose([
utils.data_transforms.ToTensor(),
])
dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[cfg.DATASET.TEST_DATASET](cfg)
eval_data_loader = torch.utils.data.DataLoader(dataset=dataset_loader.get_dataset(
utils.data_loaders.DatasetType.TEST, eval_transforms),
batch_size=cfg.EVALUATE.BATCH_SIZE,
num_workers=1,
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.EVALUATE.WEIGHT_PATH))
rec_checkpoint = torch.load(cfg.EVALUATE.WEIGHT_PATH)
net.load_state_dict(rec_checkpoint['net'])
print('[INFO] Best reconstruction result at epoch %d ...' % rec_checkpoint['epoch_idx'])
epoch_id = int(rec_checkpoint['epoch_idx'])
net.eval()
# Testing loop
for sample_idx, (taxonomy_names, sample_names, rendering_images,
model_azi, model_ele,
init_point_clouds, ground_truth_point_clouds) in enumerate(eval_data_loader):
print("evaluate sample: ", sample_idx)
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)
loss, pred_pc = net.module.valid_step(rendering_images, init_point_clouds, ground_truth_point_clouds)
azi = model_azi[0].detach().cpu().numpy()*180./np.pi
ele = model_ele[0].detach().cpu().numpy()*180./np.pi + 90.
taxonomy_name = taxonomy_names[0]
sample_name = sample_names[0]
# Save sample dir
output_dir = cfg.EVALUATE.OUTPUT_FOLDER
sample_dir = os.path.join(output_dir, str(sample_idx))
if not os.path.exists(sample_dir):
os.makedirs(sample_dir)
sample_name_dir = os.path.join(sample_dir, sample_name)
if not os.path.exists(sample_name_dir):
os.makedirs(sample_name_dir)
# Rendering image
sketch_path = os.path.join(sample_name_dir, 'sketch')
if not os.path.exists(sketch_path):
os.makedirs(sketch_path)
src_sketch_img_path = os.path.join(cfg.DATASETS.SHAPENET.RENDERING_PATH % (taxonomy_name, sample_name, 1))
copyfile(src_sketch_img_path, os.path.join(sketch_path, 'sketch.png'))
# Predict Pointcloud
p_pc = pred_pc[0].detach().cpu().numpy()
rendering_views = utils.point_cloud_visualization_old.get_point_cloud_image(p_pc,
os.path.join(sample_name_dir, 'rec results'),
sample_idx,
cfg.EVALUATE.VERSION_ID,
"",
view=[azi, ele])
# Groundtruth Pointcloud
gt_pc = ground_truth_point_clouds[0].detach().cpu().numpy()
rendering_views = utils.point_cloud_visualization_old.get_point_cloud_image(gt_pc,
os.path.join(sample_name_dir, 'gt'),
sample_idx,
cfg.EVALUATE.VERSION_ID,
"",
view=[azi, ele])
# save ply file
pc_dir = os.path.join(sample_name_dir, 'pc')
if not os.path.exists(pc_dir):
os.makedirs(pc_dir)
output_point_cloud_ply(np.array([p_pc]), ['pred_pc'], pc_dir)
if sample_idx == 200:
break
def refine_net(cfg, coarse_pc, azi, ele):
# Enable the inbuilt cudnn auto-tuner to find the best algorithm to use
torch.backends.cudnn.benchmark = True
refine_transforms = utils.data_transforms.Compose([
utils.data_transforms.ToTensor(),
])
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():
refine_net = torch.nn.DataParallel(refine_net, device_ids=cfg.CONST.DEVICE).cuda()
# Load weight for refiner
print('[INFO] %s Recovering refiner from %s ...' % (dt.now(), cfg.INFERENCE.REFINER_WEIGHTS))
refine_checkpoint = torch.load(cfg.INFERENCE.REFINER_WEIGHTS)
pretrained_dict = refine_checkpoint['net']
refine_net.load_state_dict(pretrained_dict)
refine_net.eval()
refine_image_mask = get_red_mask(cfg.INFERENCE.REFINE_IMAGE_PATH)
refine_image_mask = cv2.bitwise_not(refine_image_mask)
refine_image_mask = cv2.resize(refine_image_mask, (224, 224), interpolation=cv2.INTER_AREA)
# find the empty parts by connected component
ret, thresh = cv2.threshold(refine_image_mask, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
num_labels, labels, status, _ = cv2.connectedComponentsWithStats(thresh, connectivity=4)
empty_img = np.zeros( (224, 224, 3), np.uint8 )
for i in range(0, num_labels):
# check back ground
s = status[i]
if not (s[0] == 0 and s[1] == 0 and s[2] == 224 and s[3] == 224):
mask = labels == i
empty_img[:, :, 0][mask] = 109
empty_img[:, :, 1][mask] = 0
empty_img[:, :, 2][mask] = 216
cv2.imwrite(cfg.INFERENCE.EMPTY_IMAGE_PATH, empty_img)
sketch_img = cv2.imread(cfg.INFERENCE.SKETCH_IMAGE_PATH)
empty_img = cv2.imread(cfg.INFERENCE.EMPTY_IMAGE_PATH)
sketch_empty_img = combine_sketch_empty(sketch_img, empty_img)
cv2.imwrite(cfg.INFERENCE.UPDATE_IMAGE_PATH, sketch_empty_img)
# load update image
update_image = cv2.imread(cfg.INFERENCE.UPDATE_IMAGE_PATH, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.
samples = []
samples.append(update_image)
samples = np.array(samples).astype(np.float32)
update_images = refine_transforms(rendering_images=samples)
# set view
azi = float(azi)
rec_radian_azi = azi*np.pi/180.
batch_azi = []
batch_azi.append(rec_radian_azi)
batch_azi = np.array(batch_azi).astype(np.float32)
batch_azi = torch.from_numpy(batch_azi)
ele = float(ele)
rec_radian_ele = (ele - 90.)*np.pi/180.
batch_ele = []
batch_ele.append(rec_radian_ele)
batch_ele = np.array(batch_ele).astype(np.float32)
batch_ele = torch.from_numpy(batch_ele)
coarse_pc = np.array([coarse_pc])
coarse_pc = torch.from_numpy(coarse_pc).cuda()
with torch.no_grad():
# generate refine point cloud
refine_pc = refine_net.module.refine(update_images, coarse_pc, batch_azi, batch_ele)
output_point_cloud_ply(refine_pc.detach().cpu().numpy(), ['refine'], cfg.INFERENCE.CACHE_POINT_CLOUD_PATH)
return refine_pc[0].detach().cpu().numpy()
