-
Notifications
You must be signed in to change notification settings - Fork 95
/
Copy pathdraw-wireframe.py
executable file
·108 lines (88 loc) · 3.8 KB
/
draw-wireframe.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
#!/usr/bin/env python3
import os
import glob
import os.path as osp
import cv2
import numpy as np
import scipy.io
import matplotlib as mpl
import matplotlib.pyplot as plt
from lcnn.utils import parmap
GT = "data/wireframe/valid/*.npz"
WF = "/data/lcnn/wirebase/result/wireframe/wireframe_1_rerun-baseline_0.5_0.5/2/*.mat"
AFM = "/data/lcnn/wirebase/result/wireframe/afm/*.npz"
IMGS = "/data/lcnn/wirebase/Wireframe/v1.1/test/*.jpg"
def imshow(im):
sizes = im.shape
height = float(sizes[0])
width = float(sizes[1])
fig = plt.figure()
fig.set_size_inches(width / height, 1, forward=False)
ax = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0])
ax.set_axis_off()
fig.add_axes(ax)
plt.xlim([-0.5, sizes[1] - 0.5])
plt.ylim([sizes[0] - 0.5, -0.5])
plt.imshow(im)
def main():
gts = sorted(glob.glob(GT))
afm = sorted(glob.glob(AFM))
wf = sorted(glob.glob(WF))
img = sorted(glob.glob(IMGS))
prefix = "/data/lcnn/wirebase/myplot/"
os.makedirs(osp.join(prefix, "GT"), exist_ok=True)
os.makedirs(osp.join(prefix, "LSD"), exist_ok=True)
os.makedirs(osp.join(prefix, "AFM"), exist_ok=True)
os.makedirs(osp.join(prefix, "WF"), exist_ok=True)
os.makedirs(osp.join(prefix, "LL"), exist_ok=True)
def draw(args):
i, (wf_name, gt_name, afm_name, img_name) = args
img = cv2.imread(img_name, 0)
lsd = cv2.createLineSegmentDetector(cv2.LSD_REFINE_ADV)
lsd_line, _, _, lsd_score = lsd.detect(img)
lsd_line = lsd_line.reshape(-1, 2, 2)
lsd_score = lsd_score.flatten()
img = cv2.imread(img_name)[:, :, ::-1]
with np.load(gt_name) as fgt:
gt_line = fgt["lpos"][:, :, :2]
gt_line[:, :, 0] *= img.shape[0] / 128
gt_line[:, :, 1] *= img.shape[1] / 128
with np.load(afm_name) as fafm:
afm_line = fafm["lines"].reshape(-1, 2, 2)[:, :, ::-1]
wf_line = scipy.io.loadmat(wf_name)["lines"].reshape(-1, 2, 2)
wf_line = wf_line[:, :, ::-1]
plt.figure("GT")
imshow(img)
for a, b in gt_line - 0.5:
plt.plot([a[1], b[1]], [a[0], b[0]], color="orange", linewidth=0.5)
plt.scatter(a[1], a[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
plt.scatter(b[1], b[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
plt.savefig(osp.join(prefix, "GT", f"{i:05}"), dpi=3000, bbox_inches=0)
plt.close()
plt.figure("LSD")
imshow(img)
for a, b in lsd_line[:, :, ::-1] - 0.5:
plt.plot([a[1], b[1]], [a[0], b[0]], color="orange", linewidth=0.5)
plt.scatter(a[1], a[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
plt.scatter(b[1], b[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
plt.savefig(osp.join(prefix, "LSD", f"{i:05}"), dpi=3000, bbox_inches=0)
plt.close()
plt.figure("AFM")
imshow(img)
for a, b in afm_line - 0.5:
plt.plot([a[1], b[1]], [a[0], b[0]], color="orange", linewidth=0.5)
plt.scatter(a[1], a[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
plt.scatter(b[1], b[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
plt.savefig(osp.join(prefix, "AFM", f"{i:05}"), dpi=3000, bbox_inches=0)
plt.close()
plt.figure("WF")
imshow(img)
for a, b in wf_line - 0.5:
plt.plot([a[1], b[1]], [a[0], b[0]], color="orange", linewidth=0.5)
plt.scatter(a[1], a[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
plt.scatter(b[1], b[0], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
plt.savefig(osp.join(prefix, "WF", f"{i:05}"), dpi=3000, bbox_inches=0)
plt.close()
parmap(draw, enumerate(zip(wf, gts, afm, img)))
if __name__ == "__main__":
main()