def handle(allname):
fname, iname, gtname = allname
print("Processing", fname)
im = cv2.imread(iname)
with np.load(fname) as f:
lines = f["lines"]
scores = f["score"]
with np.load(gtname) as f:
gtlines = f["lpos"][:, :, :2]
gtlines[:, :, 0] *= im.shape[0] / 128
gtlines[:, :, 1] *= im.shape[1] / 128
for i in range(1, len(lines)):
if (lines[i] == lines[0]).all():
lines = lines[:i]
scores = scores[:i]
break
lines[:, :, 0] *= im.shape[0] / 128
lines[:, :, 1] *= im.shape[1] / 128
diag = (im.shape[0]**2 + im.shape[1]**2)**0.5
for threshold in thresholds:
nlines, nscores = postprocess(lines, scores, diag * threshold, 0,
False)
outdir = osp.join(prefix, f"{threshold:.3f}".replace(".", "_"))
os.makedirs(outdir, exist_ok=True)
npz_name = osp.join(outdir, osp.split(fname)[-1])
if args["--plot"]:
# plot gt
imshow(im[:, :, ::-1])
for (a, b) in gtlines:
plt.plot([a[1], b[1]], [a[0], b[0]],
c="orange",
linewidth=0.5)
plt.scatter(a[1], a[0], **PLTOPTS)
plt.scatter(b[1], b[0], **PLTOPTS)
plt.savefig(npz_name.replace(".npz", ".png"),
dpi=500,
bbox_inches=0)
thres = [0.96, 0.97, 0.98, 0.99]
for i, t in enumerate(thres):
imshow(im[:, :, ::-1])
for (a, b), s in zip(nlines[nscores > t],
nscores[nscores > t]):
plt.plot([a[1], b[1]], [a[0], b[0]],
c=c(s),
linewidth=0.5)
plt.scatter(a[1], a[0], **PLTOPTS)
plt.scatter(b[1], b[0], **PLTOPTS)
plt.savefig(npz_name.replace(".npz", f"_{i}.png"),
dpi=500,
bbox_inches=0)
nlines[:, :, 0] *= 128 / im.shape[0]
nlines[:, :, 1] *= 128 / im.shape[1]
np.savez_compressed(npz_name, lines=nlines, score=nscores)
def postprocess(self, threshold=0):
"""
Filters the lines to remove close duplicates in the image.
Arguments:
imshape -- image dimensions
threshold -- returned lines must have greater score (default 0)
Returns:
nlines -- filtered lines
nscores -- filtered scores
"""
diag = (self.imshape[0] ** 2 + self.imshape[1] ** 2) ** 0.5
# Multiply lines by image shape to get image point coordinates
nlines, nscores = postprocess(self.lines() * self.imshape[:2], self.scores(), diag * 0.01, 0, False)
return nlines[nscores > threshold], nscores[nscores > threshold]
def main():
args = docopt(__doc__)
config_file = args["<yaml-config>"] or "config/wireframe.yaml"
C.update(C.from_yaml(filename=config_file))
M.update(C.model)
pprint.pprint(C, indent=4)
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
device_name = "cpu"
os.environ["CUDA_VISIBLE_DEVICES"] = args["--devices"]
if torch.cuda.is_available():
device_name = "cuda"
torch.backends.cudnn.deterministic = True
torch.cuda.manual_seed(0)
print("Let's use", torch.cuda.device_count(), "GPU(s)!")
else:
print("CUDA is not available")
device = torch.device(device_name)
checkpoint = torch.load(args["<checkpoint>"], map_location=device)
# Load model
if os.path.isfile(C.io.vote_index):
vote_index = sio.loadmat(C.io.vote_index)['vote_index']
else:
vote_index = hough_transform(rows=128,
cols=128,
theta_res=3,
rho_res=1)
sio.savemat(C.io.vote_index, {'vote_index': vote_index})
vote_index = torch.from_numpy(vote_index).float().contiguous().to(device)
print('load vote_index', vote_index.shape)
model = hg(
depth=M.depth,
head=lambda c_in, c_out: MultitaskHead(c_in, c_out),
num_stacks=M.num_stacks,
num_blocks=M.num_blocks,
num_classes=sum(sum(M.head_size, [])),
vote_index=vote_index,
)
model = MultitaskLearner(model)
model = LineVectorizer(model)
model.load_state_dict(checkpoint["model_state_dict"])
model = model.to(device)
model.eval()
for imname in args["<images>"]:
print(f"Processing {imname}")
im = skimage.io.imread(imname)
if im.ndim == 2:
im = np.repeat(im[:, :, None], 3, 2)
im = im[:, :, :3]
im_resized = skimage.transform.resize(im, (512, 512)) * 255
image = (im_resized - M.image.mean) / M.image.stddev
image = torch.from_numpy(np.rollaxis(image, 2)[None].copy()).float()
with torch.no_grad():
input_dict = {
"image":
image.to(device),
"meta": [{
"junc":
torch.zeros(1, 2).to(device),
"jtyp":
torch.zeros(1, dtype=torch.uint8).to(device),
"Lpos":
torch.zeros(2, 2, dtype=torch.uint8).to(device),
"Lneg":
torch.zeros(2, 2, dtype=torch.uint8).to(device),
}],
"target": {
"jmap": torch.zeros([1, 1, 128, 128]).to(device),
"joff": torch.zeros([1, 1, 2, 128, 128]).to(device),
},
"mode":
"testing",
}
H = model(input_dict)["preds"]
lines = H["lines"][0].cpu().numpy() / 128 * im.shape[:2]
scores = H["score"][0].cpu().numpy()
for i in range(1, len(lines)):
if (lines[i] == lines[0]).all():
lines = lines[:i]
scores = scores[:i]
break
# postprocess lines to remove overlapped lines
diag = (im.shape[0]**2 + im.shape[1]**2)**0.5
nlines, nscores = postprocess(lines, scores, diag * 0.01, 0, False)
for i, t in enumerate([0.94, 0.95, 0.96, 0.97, 0.98, 0.99]):
plt.gca().set_axis_off()
plt.subplots_adjust(top=1,
bottom=0,
right=1,
left=0,
hspace=0,
wspace=0)
plt.margins(0, 0)
for (a, b), s in zip(nlines, nscores):
if s < t:
continue
plt.plot([a[1], b[1]], [a[0], b[0]],
c=c(s),
linewidth=2,
zorder=s)
plt.scatter(a[1], a[0], **PLTOPTS)
plt.scatter(b[1], b[0], **PLTOPTS)
plt.gca().xaxis.set_major_locator(plt.NullLocator())
plt.gca().yaxis.set_major_locator(plt.NullLocator())
plt.imshow(im)
plt.savefig(imname.replace(".png", f"-{t:.02f}.svg"),
bbox_inches="tight")
plt.show()
plt.close()
def main():
args = docopt(__doc__)
config_file = args["<yaml-config>"] or "config/wireframe.yaml"
C.update(C.from_yaml(filename=config_file))
M.update(C.model)
pprint.pprint(C, indent=4)
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
if M.backbone == "stacked_hourglass":
model = lcnn.models.hg(
depth=M.depth,
head=lambda c_in, c_out: MultitaskHead(c_in, c_out),
num_stacks=M.num_stacks,
num_blocks=M.num_blocks,
num_classes=sum(sum(M.head_size, [])),
)
else:
raise NotImplementedError
model = MultitaskLearner(model)
device_name = "cpu"
os.environ["CUDA_VISIBLE_DEVICES"] = args["--devices"]
if torch.cuda.is_available():
device_name = "cuda"
torch.backends.cudnn.deterministic = True
torch.cuda.manual_seed(0)
checkpoint = torch.load(args["<checkpoint>"])
print("Let's use", torch.cuda.device_count(), "GPU(s)!")
else:
checkpoint = torch.load(args["<checkpoint>"],
map_location=torch.device('cpu'))
print("CUDA is not available")
device = torch.device(device_name)
model.load_state_dict(checkpoint["model_state_dict"])
model = model.to(device)
model.eval()
print(f'evaluation batch size {M.batch_size_eval}')
loader = torch.utils.data.DataLoader(
WireframeDataset(args["<image-dir>"], split="valid"),
shuffle=False,
batch_size=M.batch_size_eval,
collate_fn=collate,
num_workers=C.io.num_workers if os.name != "nt" else 0,
pin_memory=True,
)
if os.path.exists(args["<output-dir>"]):
shutil.rmtree(args["<output-dir>"])
os.makedirs(args["<output-dir>"], exist_ok=False)
outdir = os.path.join(args["<output-dir>"], 'test_result')
os.mkdir(outdir)
for batch_idx, (image, target, iname) in enumerate(loader):
with torch.no_grad():
# predict given image
input_target = {
"center": torch.zeros_like(target['center']),
"corner": torch.zeros_like(target['corner']),
"corner_offset": torch.zeros_like(target['corner_offset']),
"corner_bin_offset":
torch.zeros_like(target['corner_bin_offset'])
}
input_dict = {
"image": recursive_to(image, device),
"target": recursive_to(input_target, device),
"mode": "validation",
}
network_start_time = time()
H = model(input_dict)["preds"]
network_end_time = time()
# plot gt & prediction
for i in range(len(iname)): #M.batch_size
if not args["--plot"]:
continue
im = image[i].cpu().numpy().transpose(1, 2, 0) # [512,512,3]
# im = im * M.image.stddev + M.image.mean
# plot&process gt
gt_im_info = [
im,
iname[i].split('.')[0] + '_gt.' + iname[i].split('.')[1]
]
gt_center = target["center"][i].cpu().numpy()
gt_corner = target["corner"][i].cpu().numpy()
gt_corner_offset = target["corner_offset"][i].cpu().numpy()
gt_corner_bin_offset = target["corner_bin_offset"][i].cpu(
).numpy()
feature_maps = [
gt_center, gt_corner, gt_corner_offset,
gt_corner_bin_offset
]
postprocess(gt_im_info,
feature_maps,
outdir,
NMS=False,
plot=True)
# plot&process pd
pd_im_info = [
im,
iname[i].split('.')[0] + '_pd.' + iname[i].split('.')[1]
]
pd_center = H["center"][i].cpu().numpy()
pd_corner = H["corner"][i].cpu().numpy()
pd_corner_offset = H["corner_offset"][i].cpu().numpy()
pd_corner_bin_offset = H["corner_bin_offset"][i].cpu().numpy()
feature_maps = [
pd_center, pd_corner, pd_corner_offset,
pd_corner_bin_offset
]
postprocess_start_time = time()
grouped_corners = postprocess(pd_im_info,
feature_maps,
outdir,
NMS=True,
plot=True)
postprocess_end_time = time()
print(
f'inference time is {postprocess_end_time-postprocess_start_time+network_end_time-network_start_time}, network cost:{network_end_time-network_start_time}, postprocessing cost:{postprocess_end_time-postprocess_start_time}'
)
# Evaluation:
# eval() # TBD
print('-----finished-----')
return
def main():
args = docopt(__doc__)
config_file = args["<yaml-config>"] or "config/wireframe.yaml"
C.update(C.from_yaml(filename=config_file))
M.update(C.model)
pprint.pprint(C, indent=4)
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
device_name = "cpu"
os.environ["CUDA_VISIBLE_DEVICES"] = args["--devices"]
if torch.cuda.is_available():
device_name = "cuda"
torch.backends.cudnn.deterministic = True
torch.cuda.manual_seed(0)
print("Let's use", torch.cuda.device_count(), "GPU(s)!")
else:
print("CUDA is not available")
device = torch.device(device_name)
checkpoint = torch.load(args["<checkpoint>"], map_location=device)
# Load model
model = lcnn.models.hg(
depth=M.depth,
head=lambda c_in, c_out: MultitaskHead(c_in, c_out),
num_stacks=M.num_stacks,
num_blocks=M.num_blocks,
num_classes=sum(sum(M.head_size, [])),
)
model = MultitaskLearner(model)
model = LineVectorizer(model)
model.load_state_dict(checkpoint["model_state_dict"])
model = model.to(device)
model.eval()
# for imname in args["<images>"]:
for root, dirs, files in os.walk(r'images'):
for root_file in files:
path = os.path.join(root, root_file)
# print(path)
for imname in glob.glob(path):
print(f"Processing {imname}")
im = skimage.io.imread(imname)
if im.ndim == 2:
im = np.repeat(im[:, :, None], 3, 2)
im = im[:, :, :3]
im_resized = skimage.transform.resize(im, (512, 512)) * 255
image = (im_resized - M.image.mean) / M.image.stddev
image = torch.from_numpy(np.rollaxis(image,
2)[None].copy()).float()
with torch.no_grad():
input_dict = {
"image":
image.to(device),
"meta": [{
"junc":
torch.zeros(1, 2).to(device),
"jtyp":
torch.zeros(1, dtype=torch.uint8).to(device),
"Lpos":
torch.zeros(2, 2, dtype=torch.uint8).to(device),
"Lneg":
torch.zeros(2, 2, dtype=torch.uint8).to(device),
}],
"target": {
"jmap": torch.zeros([1, 1, 128, 128]).to(device),
"joff": torch.zeros([1, 1, 2, 128,
128]).to(device),
},
"mode":
"testing",
}
H = model(input_dict)["preds"]
lines = H["lines"][0].cpu().numpy() / 128 * im.shape[:2]
scores = H["score"][0].cpu().numpy()
for i in range(1, len(lines)):
if (lines[i] == lines[0]).all():
lines = lines[:i]
scores = scores[:i]
break
# postprocess lines to remove overlapped lines
diag = (im.shape[0]**2 + im.shape[1]**2)**0.5
nlines, nscores = postprocess(lines, scores, diag * 0.01, 0,
False)
partExprotName = imname.split(".")[0]
exportName = partExprotName + ".txt"
with open(exportName, "w") as writeFile:
for i, t in enumerate([0.94]):
plt.gca().set_axis_off()
plt.subplots_adjust(top=1,
bottom=0,
right=1,
left=0,
hspace=0,
wspace=0)
plt.margins(0, 0)
for (a, b), s in zip(nlines, nscores):
if s < t:
continue
print(a[1], a[0], b[1], b[0], file=writeFile)
# plt.plot([a[1], b[1]], [a[0], b[0]], c=c(s), linewidth=2, zorder=s)
# plt.scatter(a[1], a[0], **PLTOPTS)
# plt.scatter(b[1], b[0], **PLTOPTS)
# plt.gca().xaxis.set_major_locator(plt.NullLocator())
# plt.gca().yaxis.set_major_locator(plt.NullLocator())
# plt.imshow(im)
# plt.savefig(imname.replace(".png", f"-{t:.02f}.svg"), bbox_inches="tight")
# plt.show()
# plt.close()
for new_root, new_dir, new_files in os.walk(r'images'):
# print(new_root)
for root_file1 in new_files:
path1 = os.path.join(new_root, root_file1)
for all_files in glob.glob(path1):
txtname1 = os.path.splitext(all_files)[1]
txtname0 = os.path.splitext(all_files)[0]
# print(txtname1, txtname0)
if txtname1 == '.txt':
# old_path = os.path.join(new_root, all_files)
new_path = 'results/'
shutil.move(path1, new_path)
