def create_montage(render_dir, output_cage):
"""render shapes inside a directory with thea"""
files = find_files(render_dir, "png")
output_dir = os.path.join(render_dir, "montage")
os.makedirs(output_dir, exist_ok=True)
source_files = [p for p in files if "Sa.png" in p]
pool = ThreadPool(processes=4)
results = []
for source in source_files:
dir_path = os.path.dirname(source)
sname, tname = os.path.basename(source).split("-")[:2]
target = source.replace("Sa", "Sb")
deformed = glob(
os.path.join(dir_path, "{}-{}-Sab*.png".format(sname, tname)))
deformed = " ".join(deformed)
if output_cage:
cage1 = glob(
os.path.join(dir_path, "{}-{}-cage1*.png".format(sname,
tname)))
cage1 = " ".join(cage1)
cage2 = glob(
os.path.join(dir_path, "{}-{}-cage2*.png".format(sname,
tname)))
cage2 = " ".join(cage2)
else:
cage1 = cage2 = ""
output_file = os.path.join(output_dir,
"{}-{}.png".format(sname, tname))
# -gravity Center -crop 480x480+0+0 +repage
results.append(
pool.apply_async(
call_proc,
("montage -geometry +2+0 -trim -tile x1 {} {} {} {} {} {}".
format(source, target, deformed, cage1, cage2,
output_file), )))
# Close the pool and wait for each running task to complete
pool.close()
pool.join()
for result in results:
out, err = result.get()
if len(err) > 0:
print("err: {}".format(err))
else:
return ("", "")
if __name__ == "__main__":
N_CORE = 8
N_POINT = 5000
print("Using %d of %d cores" % (N_CORE, multiprocessing.cpu_count()))
source_dir = sys.argv[1] # input directoy
output_dir = sys.argv[2] # output directory
###################################
# 1. gather source and target
###################################
source_files = find_files(source_dir, 'obj')
logger.info("Found {} source files".format(len(source_files)))
os.makedirs(output_dir, exist_ok=True)
###################################
# Sample
###################################
pool = ThreadPool(processes=N_CORE)
results = []
for input_file in source_files:
source_name = os.path.splitext(os.path.basename(input_file))[0]
my_out_dir = os.path.join(
output_dir, os.path.relpath(os.path.dirname(input_file),
source_dir))
os.makedirs(my_out_dir, exist_ok=True)
def create_comparison_montage(render_dirs,
labels,
output_dir,
output_cage=False):
files = find_files(render_dirs[0], "png")
source_files = [p for p in files if "Sb.png" in p]
logger.info("Found {} files".format(len(source_files)))
pool = ThreadPool(processes=4)
results = []
os.makedirs(output_dir, exist_ok=True)
for source in source_files:
sname, tname = os.path.basename(source).split("-")[:2]
output_file = os.path.join(output_dir,
"{}-{}.png".format(sname, tname))
images = [
glob(os.path.join(cur_dir, "{}-{}-Sab*.png".format(sname, tname)))
for cur_dir in render_dirs
]
if not all([len(im_found) > 0 for im_found in images]):
indices = [i for i, x in enumerate(images) if len(x) == 0]
logger.warn(
"", "Cannot find {} in {}".format(
"{}-{}-Sab*.png".format(sname, tname),
", ".join([render_dirs[i] for i in indices])))
continue
images = [im_found[0] for im_found in images]
cages = []
mylabels = labels[:]
if output_cage:
# find cages
for i, dir_img in enumerate(zip(render_dirs, images)):
cur_dir, image = dir_img
cage1 = glob(image.replace("Sab", "cage1"))
cage2 = glob(image.replace("Sab", "cage2"))
if len(cage1) > 0 and len(cage2) > 0:
cages.append((i, cage1[0], cage2[0]))
# insert cages to the correct position in images
cnt = 0
for offset, cage1, cage2 in cages:
images.insert(offset + cnt + 1, cage1)
images.insert(offset + cnt + 2, cage2)
mylabels.insert(offset + cnt + 1,
mylabels[offset + cnt] + "_cage1")
mylabels.insert(offset + cnt + 2,
mylabels[offset + cnt] + "_cage2")
cnt += 2
assert (len(images) == len(mylabels))
image_strs = " ".join([
"-label {} {}".format(l, i) for l, i in zip(mylabels, images)
])
else:
image_strs = " ".join([
"-label {} {}".format(l, i) for l, i in zip(mylabels, images)
])
num_cols = len(images) + 2
target = source.replace("Sa", "Sb")
results.append(
pool.apply_async(call_proc, (
"montage -geometry +0+0 -gravity Center -crop 420x450+0+0 +repage -tile {}x1 -label input {} -label target {} {} {}"
.format(num_cols, target, source, image_strs, output_file), )))
# Close the pool and wait for each running task to complete
pool.close()
pool.join()
for result in results:
out, err = result.get()
if len(err) > 0:
print("err: {}".format(err))
def create_two_row_comparison_montage(render_dirs,
labels,
output_dir,
output_cage=True):
files = find_files(render_dirs[0], "png")
source_files = [p for p in files if "Sa.png" in p]
logger.info("Found {} files".format(len(source_files)))
pool = ThreadPool(processes=4)
results = []
os.makedirs(output_dir, exist_ok=True)
# first concatenate cage1-cage2
for cur_dir in render_dirs:
cage1s = glob(os.path.join(cur_dir, "*cage1*.png"))
cage2s = [f.replace("cage1", "cage2") for f in cage1s]
for cage1, cage2 in zip(cage1s, cage2s):
if not (os.path.isfile(cage1) and os.path.isfile(cage2)):
continue
output_file = os.path.join(cage1.replace("cage1", "cages"))
results.append(
pool.apply_async(call_proc, (
"montage -geometry +0+0 -gravity Center -crop 400x400+0+0 +repage -tile 2x1 {} {} {}"
.format(cage1, cage2, output_file), )))
pool.close()
pool.join()
for result in results:
out, err = result.get()
if len(err) > 0:
print("err: {}".format(err))
results.clear()
pool = ThreadPool(processes=4)
for source in source_files:
sname, tname = os.path.basename(source).split("-")[:2]
output_file = os.path.join(output_dir,
"{}-{}.png".format(sname, tname))
images = [
glob(os.path.join(cur_dir, "{}-{}-Sab*.png".format(sname, tname)))
for cur_dir in render_dirs
]
if not all([len(im_found) > 0 for im_found in images]):
indices = [i for i, x in enumerate(images) if len(x) == 0]
logger.warn(
"", "Cannot find {} in {}".format(
"{}-{}-Sab*.png".format(sname, tname),
", ".join([render_dirs[i] for i in indices])))
images = [
"null:" if len(im_found) == 0 else im_found[0]
for im_found in images
]
cages = [
glob(os.path.join(cur_dir, "{}-{}-cages*.png".format(sname,
tname)))
for cur_dir in render_dirs
]
if not all([len(im_found) > 0 for im_found in cages]):
indices = [i for i, x in enumerate(cages) if len(x) == 0]
logger.warn(
"", "Cannot find {} in {}".format(
"{}-{}-cages*.png".format(sname, tname),
", ".join([render_dirs[i] for i in indices])))
cages = [
"null:" if len(im_found) == 0 else im_found[0]
for im_found in cages
]
mylabels = labels[:]
assert (len(images) == len(mylabels))
image_strs = " ".join(images)
cage_str = " ".join(
["-label {} {}".format(l, i) for l, i in zip(mylabels, cages)])
num_cols = len(images) + 1
target = source.replace("Sa", "Sb")
results.append(
pool.apply_async(
call_proc,
("montage -geometry \'420x400>+0+0\' -tile {}x2 {} "
"{} {} {} {}".format(num_cols, source, image_strs, target,
cage_str, output_file), )))
# Close the pool and wait for each running task to complete
pool.close()
pool.join()
for result in results:
out, err = result.get()
if len(err) > 0:
print("err: {}".format(err))
for cur_dir in render_dirs:
call_proc("rm {}".format(os.path.join(cur_dir, "*.cages*.png")))
def run(in_files, out_dir, ref_file):
ref_v, ref_f = read_trimesh(ref_file)
ref_v = torch.from_numpy(ref_v[:, :3]).float()
ref_f = torch.from_numpy(ref_f).long()
_, ref_area = compute_face_normals_and_areas(ref_v, ref_f)
ref_area = torch.sum(ref_area, dim=-1)
for in_file in in_files:
v, f = read_trimesh(in_file)
v = torch.from_numpy(v[:, :3]).float()
f = torch.from_numpy(f).long()
v, _, _ = center_bounding_box(v)
_, area = compute_face_normals_and_areas(v, f)
area = torch.sum(area, dim=-1)
ratio = torch.sqrt(ref_area/area)
ratio = ratio.unsqueeze(-1).unsqueeze(-1)
v = v * ratio
out_path = os.path.join(out_dir, os.path.basename(in_file))
pymesh.save_mesh_raw(out_path, v.numpy(), f.numpy())
print("saved to {}".format(out_path))
if __name__ == "__main__":
in_file = sys.argv[1]
ref_file = sys.argv[2]
out_dir = sys.argv[3]
if os.path.isdir(in_file):
in_file = find_files(in_file, ["ply", "obj"])
os.makedirs(out_dir, exist_ok=True)
run(in_file, out_dir, ref_file)
def evaluate_svr(result_dirs, resample, overwrite_pts=False):
""" ours is the first in the result dirs """
if isinstance(result_dirs, str):
result_dirs = [result_dirs]
########## initialize ############
eval_result = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: 1e10))) # eval_result[metric[folder[file]]]
avg_result = defaultdict(lambda: defaultdict(lambda: defaultdict(float))) # eval_result[metric[folder[file]]]
cotLap = CotLaplacian()
uniLap = UniformLaplacian()
if resample and not opt.mse:
for cur_dir in result_dirs:
pts_dir = os.path.join(cur_dir, "eval_pts")
os.makedirs(pts_dir, exist_ok=True)
result = svr_sample_pts(cur_dir, pts_dir, overwrite_pts)
if not result:
logger.warn("Failed to sample points in {}".format(cur_dir))
########## load results ###########
# find Sa.ply, Sb.ply and a list of Sab.ply
###################################
[print("dir{}: {}".format(i, name)) for i, name in enumerate(result_dirs)]
files = find_files(result_dirs[0], ["ply", "obj"])
target_files = [p for p in files if "Sb." in p]
target_names = np.unique(np.array([os.path.basename(p).split("-")[1] for p in target_files])).tolist()
logger.info("Found {} target files".format(len(target_names)))
########## evaluation ############
print("{}: {}".format("filename".ljust(70), " | ".join(["dir{}".format(i).rjust(20) for i in range(len(result_dirs))])))
print("{}: {}".format(" ".ljust(70), " | ".join(["CD/HD".rjust(20) for i in range(len(result_dirs))])))
cnt = 0
for target in target_names:
# 1. load ground truth
gt_path = glob(os.path.join(result_dirs[0], "*-{}-Sb.*".format(target)))[0]
try:
gt_shape, gt_face = read_trimesh(gt_path, clean=False)
if resample:
gt_pts_file = os.path.join(result_dirs[0], "eval_pts", "{}.pts".format(target))
if not os.path.isfile(gt_pts_file):
logger.warn("Cound\'t find {}. Skip to process the next.".format(gt_pts_file))
continue
gt_pts = load(gt_pts_file)
gt_pts = torch.from_numpy(gt_pts[:,:3].astype(np.float32)).unsqueeze(0).cuda()
ours_paths = glob(os.path.join(result_dirs[0], "*-{}-Sab.*".format(target)))
others_path = [glob( os.path.join(cur_dir, "{}.*".format(target)) ) for cur_dir in result_dirs[1:]]
# 2. evaluate ours, all *-{target}-Sab
if len(ours_paths) == 0:
logger.warn("Cound\'t find {}. Skip to process the next.".format(os.path.join(result_dirs[0], "*-{}-Sab.*".format(target))))
continue
for ours in ours_paths:
# load shape and points
output_shape, output_face = read_trimesh(ours, clean=False)
ours = os.path.basename(ours)
cur_dir = result_dirs[0]
if resample:
output_pts_file = os.path.join(cur_dir, "eval_pts", ours[:-4]+".pts")
if not os.path.isfile(output_pts_file):
logger.warn("Cound\'t find {}. Skip to process the next source.".format(output_pts_file))
continue
output_pts = load(output_pts_file)
output_pts = torch.from_numpy(output_pts[:,:3].astype(np.float32)).unsqueeze(0).cuda()
# compute chamfer
dist12, dist21, _, _ = nndistance(gt_pts, output_pts)
cd = torch.mean(torch.mean(dist12, dim=-1) + torch.mean(dist21, dim=-1)).item()
hd = max(torch.max(dist12).item(), torch.max(dist21).item())
else:
dist12, dist21, _, _ = nndistance(gt_shape, output_shape)
cd = torch.mean(torch.mean(dist12, dim=-1) + torch.mean(dist21, dim=-1)).item()
hd = max(torch.max(dist12).item(), torch.max(dist21).item())
eval_result[cur_dir]["CD"][target] = min(eval_result[cur_dir]["CD"][target], cd)
avg_result[cur_dir]["CD"]["avg"] += (cd - avg_result[cur_dir]["CD"]["avg"])/(avg_result[cur_dir]["CD"]["cnt"]+1)
avg_result[cur_dir]["CD"]["cnt"]+=1
eval_result[cur_dir]["HD"][target] = min(eval_result[cur_dir]["HD"][target], hd)
avg_result[cur_dir]["HD"]["avg"] += (hd - avg_result[cur_dir]["HD"]["avg"])/(avg_result[cur_dir]["HD"]["cnt"]+1)
avg_result[cur_dir]["HD"]["cnt"]+=1
# 3. evaluation others
for cur_dir in result_dirs[1:]:
result_path = glob(os.path.join(cur_dir, "{}.*".format(target)))
if len(result_path) == 0:
logger.warn("Cound\'t find {}. Skip to process the next.".format(result_path))
continue
result_path = result_path[0]
output_shape, output_face = read_trimesh(result_path, clean=False)
result_name = os.path.splitext(os.path.basename(result_path))[0]
if resample:
output_pts_file = os.path.join(cur_dir, "eval_pts", result_name+".pts")
if not os.path.isfile(output_pts_file):
logger.warn("Cound\'t find {}. Skip to process the next source.".format(output_pts_file))
continue
output_pts = load(output_pts_file)
output_pts = torch.from_numpy(output_pts[:,:3].astype(np.float32)).unsqueeze(0).cuda()
# compute chamfer
dist12, dist21, _, _ = nndistance(gt_pts, output_pts)
cd = torch.mean(torch.mean(dist12, dim=-1) + torch.mean(dist21, dim=-1)).item()
hd = max(torch.max(dist12).item(), torch.max(dist21).item())
else:
dist12, dist21, _, _ = nndistance(gt_shape, output_shape)
cd = torch.mean(torch.mean(dist12, dim=-1) + torch.mean(dist21, dim=-1)).item()
hd = max(torch.max(dist12).item(), torch.max(dist21).item())
eval_result[cur_dir]["CD"][target] = min(eval_result[cur_dir]["CD"][target], cd)
avg_result[cur_dir]["CD"]["avg"] += (cd - avg_result[cur_dir]["CD"]["avg"])/(avg_result[cur_dir]["CD"]["cnt"]+1)
avg_result[cur_dir]["CD"]["cnt"]+=1
eval_result[cur_dir]["HD"][target] = min(eval_result[cur_dir]["HD"][target], hd)
avg_result[cur_dir]["HD"]["avg"] += (hd - avg_result[cur_dir]["HD"]["avg"])/(avg_result[cur_dir]["HD"]["cnt"]+1)
avg_result[cur_dir]["HD"]["cnt"]+=1
print("{}: {}".format(target.ljust(70), " | ".join(
["{:8.4g}/{:8.4g}".format(
eval_result[cur_dir]["CD"][target],
eval_result[cur_dir]["HD"][target],
)
for cur_dir in result_dirs]
).ljust(30)))
except Exception as e:
traceback.print_exc(file=sys.stdout)
logger.warn("Failed to evaluation {}. Skip to process the next.".format(target))
print("{}: {}".format("AVG".ljust(70), " | ".join(
["{:8.4g}/{:8.4g}".format(
avg_result[cur_dir]["CD"]["avg"],
avg_result[cur_dir]["HD"]["avg"],
)
for cur_dir in result_dirs]
).ljust(30)))
########## write evaluation ############
for cur_dir in result_dirs:
for metric in eval_result[cur_dir]:
output_file = os.path.join(cur_dir, "eval_{}.txt".format(metric))
with open(output_file, "w") as eval_file:
for name, value in eval_result[cur_dir][metric].items():
if (name != "avg" and name != "cnt"):
eval_file.write("{} {:8.4g}\n".format(name, value))
eval_file.write("avg {:8.4g}".format(eval_result[cur_dir][metric]["avg"]))