def main(args):
# save_folder = Path(args.save_folder)
# if not save_folder.exists():
# save_folder.mkdir(parents=True)
# net = utils.load_trained_model(HorizonNet, args.pth).to(device)
net = utils.load_trained_model(HorizonNet,
pretrained_model_path).to(device)
net.eval()
img_dirs = Path(args.img_dir)
img_dirs = list(img_dirs.glob("*"))
# print(img_file_paths)
with torch.no_grad():
for i, img_dir in enumerate(img_dirs):
# print(img_dir)
# if str(img_dir.stem) != "panel_421865_洋室":
# continue
# print("here")
img_files = list(img_dir.glob("*.jpg")) + list(
img_dir.glob("*.png"))
img_path = None
for img_file in img_files:
if "panel" in img_file.stem and "raw" not in img_file.stem:
img_path = img_file
if img_path is None:
continue
print("start processing {}".format(img_path))
# output_dir = save_folder/img_path.stem
if args.output_dir is None:
output_dir = img_dir
else:
output_dir = args.output_dir
# if not output_dir.exists():
# output_dir.mkdir(parents=True)
img_orig = preprocess(img_path)
img_orig.resize(
(1024,
512)).save(img_path.parent /
(img_path.stem + "_{}x{}.png".format(1024, 512)))
if img_orig.size != (1024, 512):
img = img_orig.resize((1024, 512), Image.BICUBIC)
else:
img = img_orig
img_ori = np.array(img)[..., :3].transpose([2, 0, 1]).copy()
x = torch.FloatTensor([img_ori / 255])
# Inferenceing corners
cor_id, z0, z1, vis_out, boundary = inference(
net, x, device, args.flip, args.rotate, args.visualize,
args.force_cuboid, args.min_v, args.r)
boundary = ((boundary / np.pi + 0.5) * 512).round().astype(int)
# print(boundary[0,511])
# print(boundary[1,511])
if vis_out is not None:
vis_path = output_dir / (img_path.stem + '_raw.png')
# add red line to highlight the coordinates
vis_out
vh, vw = vis_out.shape[:2]
Image.fromarray(vis_out)\
.resize((vw//2, vh//2), Image.LANCZOS)\
.save(vis_path)
make_3D_json_file(cor_id,
np.array(img_orig) / 255.0,
boundary,
output_dir,
img_path=img_path,
camera_height=args.camera_height)
try:
make_3d_files(cor_id,
np.array(img_orig) / 255.0,
output_dir,
img_path=img_path,
camera_height=args.camera_height,
ignore_ceiling=False,
write_obj_files=args.write_obj_files,
write_point_cloud=args.write_point_cloud)
# make_3d_files(cor_id, vis_out/255.0, output_dir, write_obj_files=True, write_point_cloud=True)
except Exception as e:
print(e)
print("error", img_path)
continue
# Prepare image to processed
paths = sorted(glob.glob(args.img_glob))
if len(paths) == 0:
print('no images found')
for path in paths:
assert os.path.isfile(path), '%s not found' % path
# Check target directory
if not os.path.isdir(args.output_dir):
print('Output directory %s not existed. Create one.' % args.output_dir)
os.makedirs(args.output_dir)
device = torch.device('cpu' if args.no_cuda else 'cuda')
# Loaded trained model
net = utils.load_trained_model(HorizonNet, args.pth).to(device)
net.eval()
# Inferencing
with torch.no_grad():
for i_path in tqdm(paths, desc='Inferencing'):
k = os.path.split(i_path)[-1][:-4]
# Load image
img_pil = Image.open(i_path)
if img_pil.size != (1024, 512):
img_pil = img_pil.resize((1024, 512), Image.BICUBIC)
img_ori = np.array(img_pil)[..., :3].transpose([2, 0, 1]).copy()
x = torch.FloatTensor([img_ori / 255])
# Inferenceing corners
def main(args):
output_dir_json = Path(args.output_dir_json)
if not output_dir_json.exists():
output_dir_json.mkdir(parents=True)
output_dir_pcd = Path(args.output_dir_pcd)
if not output_dir_pcd.exists():
output_dir_pcd.mkdir(parents=True)
output_dir_line = Path(args.output_dir_line)
if not output_dir_line.exists():
output_dir_line.mkdir(parents=True)
# net = utils.load_trained_model(HorizonNet, args.pth).to(device)
net = utils.load_trained_model(HorizonNet,
pretrained_model_path).to(device)
net.eval()
img_dirs = Path(args.img_dir)
img_files = list(img_dirs.glob("*.jpg"))
# print(img_file_paths)
with torch.no_grad():
for i, img_path in enumerate(img_files):
try:
print(img_path)
if "panel" not in img_path.stem or "raw" in img_path.stem:
continue
# try:
img_orig = preprocess(img_path)
# img_orig = Image.open(img_path)
# except:
# continue
if img_orig.size != (1024, 512):
img = img_orig.resize((1024, 512), Image.BICUBIC)
else:
img = img_orig
img_ori = np.array(img)[..., :3].transpose([2, 0, 1]).copy()
x = torch.FloatTensor([img_ori / 255])
# Inferenceing corners
cor_id, z0, z1, vis_out, boundary = inference(
net, x, device, args.flip, args.rotate, args.visualize,
args.force_cuboid, args.min_v, args.r)
boundary = ((boundary / np.pi + 0.5) * 512).round().astype(int)
# print(boundary[0,511])
# print(boundary[1,511])
if vis_out is not None:
vis_path = output_dir_line / (img_path.stem + '_raw.png')
# print(vis_path)
vh, vw = vis_out.shape[:2]
Image.fromarray(vis_out)\
.resize((vw//2, vh//2), Image.LANCZOS)\
.save(vis_path)
make_3D_json_file(cor_id,
np.array(img_orig) / 255.0,
boundary,
output_dir_json,
img_path=img_path,
camera_height=1.0)
# try:
make_3d_files(cor_id,
np.array(img_orig) / 255.0,
output_dir_pcd,
img_path=img_path,
write_obj_files=args.write_obj_files,
write_point_cloud=args.write_point_cloud)
# except:
# print("error", img_path)
# continue
#
except:
print("error", img_path)
with open("./3d_reconstruction/log", "a") as f:
f.write(img_path + "\n")
continue