def json2npz(src_path, dst_path, cfg, plot=False):
split = 'test'
os.makedirs(dst_path, exist_ok=True)
json_file = os.path.join(src_path, f'{split}.json')
with open(json_file, 'r') as f:
dataset = json.load(f)
for data in tqdm(dataset, desc=split):
filename = data['filename']
lines = np.asarray(data['lines'])
image = cv2.imread(os.path.join(src_path, 'image', filename))
if cfg.type == 'fisheye':
coeff = {'K': np.asarray(data['K']), 'D': np.asarray(data['D'])}
camera = cam.Fisheye(coeff)
pts_list = camera.interp_line(lines, resolution=0.01)
lines = bez.fit_line(pts_list, order=6)[0]
elif cfg.type == 'spherical':
image_size = (image.shape[1], image.shape[0])
camera = cam.Spherical(image_size)
lines = camera.truncate_line(lines)
lines = camera.remove_line(lines, thresh=10.0)
pts_list = camera.interp_line(lines, resolution=0.01)
lines = bez.fit_line(pts_list, order=6)[0]
prefix = os.path.join(dst_path, filename.split('.')[0])
save_npz(prefix, image, lines.copy(), cfg.heatmap_size)
if plot:
bez.insert_line(image, lines, color=[0, 0, 255])
bez.insert_point(image, lines, color=[255, 0, 0], thickness=4)
cv2.namedWindow('image', 0)
cv2.imshow('image', image)
cv2.waitKey()
def call_back(data):
filename = data['filename']
lines0 = np.asarray(data['lines'])
image0 = cv2.imread(os.path.join(src_path, 'image', filename))
if cfg.type == 'pinhole':
camera = cam.Pinhole()
elif cfg.type == 'fisheye':
coeff = {'K': np.asarray(data['K']), 'D': np.asarray(data['D'])}
camera = cam.Fisheye(coeff)
else:
image_size = (image0.shape[1], image0.shape[0])
camera = cam.Spherical(image_size)
if split == 'train':
for i in range(len(tfs)):
image, lines = tfs[i](image0, lines0)
if cfg.type == 'spherical':
lines = camera.truncate_line(lines)
lines = camera.remove_line(lines, thresh=10.0)
pts_list = camera.interp_line(lines)
lines = bez.fit_line(pts_list, order=2)[0]
centers = lines[:, 1]
lines = bez.fit_line(pts_list, order=cfg.order)[0]
prefix = os.path.join(dst_path, split,
filename.split('.')[0] + f'_{i}')
save_npz(prefix, image, lines.copy(), centers, cfg)
if plot:
bez.insert_line(image, lines, color=[0, 0, 255])
bez.insert_point(image,
lines,
color=[255, 0, 0],
thickness=2)
cv2.namedWindow('image', 0)
cv2.imshow('image', image)
cv2.waitKey()
else:
image, lines = image0.copy(), lines0.copy()
if cfg.type == 'spherical':
lines = camera.truncate_line(lines)
lines = camera.remove_line(lines, thresh=10.0)
pts_list = camera.interp_line(lines)
lines = bez.fit_line(pts_list, order=2)[0]
centers = lines[:, 1]
lines = bez.fit_line(pts_list, order=cfg.order)[0]
prefix = os.path.join(dst_path, split, filename.split('.')[0])
save_npz(f'{prefix}', image, lines.copy(), centers, cfg)
if plot:
bez.insert_line(image, lines, color=[0, 0, 255])
bez.insert_point(image, lines, color=[255, 0, 0], thickness=2)
cv2.namedWindow('image', 0)
cv2.imshow('image', image)
cv2.waitKey()
def call_back(data):
filename = data['filename']
lines0 = np.asarray(data['lines'])
image0 = cv2.imread(os.path.join(src_path, 'image', filename))
coeff = {'K': np.asarray(data['K']), 'D': np.asarray(data['D'])}
camera = cam.Fisheye(coeff)
if split == 'train':
for i in range(len(tfs)):
image, lines = tfs[i](image0, lines0)
pts_list = camera.interp_line(lines, resolution=0.01)
lines = bez.fit_line(pts_list, order=2)[0]
centers = lines[:, 1]
lines = bez.fit_line(pts_list, order=cfg.order)[0]
prefix = os.path.join(dst_path, split, filename.split('.')[0] + f'_{i}')
save_npz(prefix, image, lines.copy(), centers, cfg)
if plot:
bez.insert_line(image, lines, color=[0, 0, 255])
bez.insert_point(image, lines, color=[255, 0, 0], thickness=2)
cv2.namedWindow('image', 0)
cv2.imshow('image', image)
cv2.waitKey()
else:
image, lines = image0.copy(), lines0.copy()
pts_list = camera.interp_line(lines, coeff, resolution=0.01)
lines = bez.fit_line(pts_list, order=2)[0]
centers = lines[:, 1]
lines = bez.fit_line(pts_list, order=cfg.order)[0]
prefix = os.path.join(dst_path, split, filename.split('.')[0])
save_npz(f'{prefix}', image, lines.copy(), centers, cfg)
if plot:
bez.insert_line(image, lines, color=[0, 0, 255])
bez.insert_point(image, lines, color=[255, 0, 0], thickness=2)
cv2.namedWindow('image', 0)
cv2.imshow('image', image)
cv2.waitKey()
def call_back(data):
filename = data['filename']
lines0 = np.asarray(data['lines'])
image0 = cv2.imread(os.path.join(src_path, 'image', filename))
if split == 'train':
for i in range(len(tfs)):
image, lines = tfs[i](image0, lines0)
centers = lines.mean(axis=1)
prefix = os.path.join(dst_path, split,
filename.split('.')[0] + f'_{i}')
save_npz(prefix, image, lines.copy(), centers, cfg)
if plot:
bez.insert_line(image, lines, color=[0, 0, 255])
bez.insert_point(image,
lines,
color=[255, 0, 0],
thickness=2)
cv2.namedWindow('image', 0)
cv2.imshow('image', image)
cv2.waitKey()
else:
image, lines = image0.copy(), lines0.copy()
centers = lines.mean(axis=1)
prefix = os.path.join(dst_path, split, filename.split('.')[0])
save_npz(f'{prefix}', image, lines.copy(), centers, cfg)
if plot:
bez.insert_line(image, lines, color=[0, 0, 255])
bez.insert_point(image, lines, color=[255, 0, 0], thickness=2)
cv2.namedWindow('image', 0)
cv2.imshow('image', image)
cv2.waitKey()
def save_npz(prefix, image, lines, centers, cfg):
n_pts = cfg.order + 1
image_size = (image.shape[1], image.shape[0])
heatmap_size = cfg.heatmap_size
sx, sy = heatmap_size[0] / image_size[0], heatmap_size[1] / image_size[1]
lines_mask = lines[:, 0, 1] > lines[:, -1, 1]
lines[lines_mask] = lines[lines_mask, ::-1]
lines[:, :, 0] = np.clip(lines[:, :, 0] * sx, 0, heatmap_size[0] - 1e-4)
lines[:, :, 1] = np.clip(lines[:, :, 1] * sy, 0, heatmap_size[1] - 1e-4)
centers[:, 0] = np.clip(centers[:, 0] * sx, 0, heatmap_size[0] - 1e-4)
centers[:, 1] = np.clip(centers[:, 1] * sy, 0, heatmap_size[1] - 1e-4)
jmap = np.zeros((1, ) + heatmap_size[::-1], dtype=np.float32)
joff = np.zeros((2, ) + heatmap_size[::-1], dtype=np.float32)
cmap = np.zeros((1, ) + heatmap_size[::-1], dtype=np.float32)
coff = np.zeros((2, ) + heatmap_size[::-1], dtype=np.float32)
lvec = np.zeros((
(n_pts // 2) * 2,
2,
) + heatmap_size[::-1],
dtype=np.float32)
lmap = np.zeros((1, ) + heatmap_size[::-1], dtype=np.float32)
juncs = np.concatenate((lines[:, 0], lines[:, -1]))
juncs = np.round(juncs, 3)
juncs = np.unique(juncs, axis=0)
lpos = lines.copy()
lneg = []
if n_pts % 2 == 1:
lines = np.delete(lines, n_pts // 2, axis=1)
def to_int(x):
return tuple(map(int, x))
for c, pts in zip(centers, lines):
v0, v1 = pts[0], pts[-1]
cint = to_int(c)
vint0 = to_int(v0)
vint1 = to_int(v1)
jmap[0, vint0[1], vint0[0]] = 1
jmap[0, vint1[1], vint1[0]] = 1
joff[:, vint0[1], vint0[0]] = v0 - vint0 - 0.5
joff[:, vint1[1], vint1[0]] = v1 - vint1 - 0.5
cmap[0, cint[1], cint[0]] = 1
coff[:, cint[1], cint[0]] = c - cint - 0.5
lvec[:, :, cint[1], cint[0]] = pts - c
lvec = lvec.reshape((-1, ) + heatmap_size[::-1])
lmap[0] = bez.insert_line(lmap[0], lpos, color=255) / 255.0
np.savez_compressed(f'{prefix}.npz',
junc=juncs,
lpos=lpos,
lneg=lneg,
jmap=jmap,
joff=joff,
cmap=cmap,
coff=coff,
lvec=lvec,
lmap=lmap)
cv2.imwrite(f'{prefix}.png', image)
def save_npz(prefix, image, lines, centers, cfg):
n_pts = cfg.order + 1
image_size = (image.shape[1], image.shape[0])
heatmap_size = cfg.heatmap_size
sx, sy = heatmap_size[0] / image_size[0], heatmap_size[1] / image_size[1]
lines_mask = lines[:, 0, 1] > lines[:, -1, 1]
lines[lines_mask] = lines[lines_mask, ::-1]
lines[:, :, 0] = np.clip(lines[:, :, 0] * sx, 0, heatmap_size[0] - 1e-4)
lines[:, :, 1] = np.clip(lines[:, :, 1] * sy, 0, heatmap_size[1] - 1e-4)
centers[:, 0] = np.clip(centers[:, 0] * sx, 0, heatmap_size[0] - 1e-4)
centers[:, 1] = np.clip(centers[:, 1] * sy, 0, heatmap_size[1] - 1e-4)
jmap = np.zeros((1, ) + heatmap_size[::-1], dtype=np.float32)
joff = np.zeros((2, ) + heatmap_size[::-1], dtype=np.float32)
cmap = np.zeros((1, ) + heatmap_size[::-1], dtype=np.float32)
coff = np.zeros((2, ) + heatmap_size[::-1], dtype=np.float32)
lvec = np.zeros((
(n_pts // 2) * 2,
2,
) + heatmap_size[::-1],
dtype=np.float32)
lmap = np.zeros((1, ) + heatmap_size[::-1], dtype=np.float32)
juncs = dict()
lpos = lines.copy()
lneg, scores = [], []
linds = []
if n_pts % 2 == 1:
lines = np.delete(lines, n_pts // 2, axis=1)
def to_int(x):
return tuple(map(int, x))
def add_junc(juncs, junc):
ind = len(juncs)
if junc in juncs:
ind = juncs[junc]
else:
juncs[junc] = ind
return ind
for c, pts in zip(centers, lines):
v0, v1 = pts[0], pts[-1]
lind = [add_junc(juncs, tuple(v0)), add_junc(juncs, tuple(v1))]
linds.append(lind)
cint = to_int(c)
vint0 = to_int(v0)
vint1 = to_int(v1)
jmap[0, vint0[1], vint0[0]] = 1
jmap[0, vint1[1], vint1[0]] = 1
joff[:, vint0[1], vint0[0]] = v0 - vint0 - 0.5
joff[:, vint1[1], vint1[0]] = v1 - vint1 - 0.5
cmap[0, cint[1], cint[0]] = 1
coff[:, cint[1], cint[0]] = c - cint - 0.5
lvec[:, :, cint[1], cint[0]] = pts - c
lvec = lvec.reshape((-1, ) + heatmap_size[::-1])
lmap[0] = bez.insert_line(lmap[0], lpos, color=255) / 255.0
juncs = np.asarray([np.array(junc) for junc in list(juncs.keys())])
linds = np.asarray(linds)
llmap = zoom(lmap[0], [0.5, 0.5])
lineset = set([frozenset(l) for l in linds])
for i0, i1 in combinations(range(len(juncs)), 2):
if frozenset([i0, i1]) not in lineset:
v0, v1 = juncs[i0], juncs[i1]
lneg.append([v0, v1])
v0, v1 = v0 / 2.0, v1 / 2.0
rr, cc, value = skimage.draw.line_aa(int(v0[1]), int(v0[0]),
int(v1[1]), int(v1[0]))
scores.append(np.average(np.minimum(value, llmap[rr, cc])))
lpos = np.asarray(lpos)
lneg = np.asarray(lneg)
scores = np.asarray(scores)
indices = np.argsort(-scores)[:2000]
lneg = lneg[indices]
np.savez_compressed(f'{prefix}.npz',
junc=juncs,
lpos=lpos,
lneg=lneg,
jmap=jmap,
joff=joff,
cmap=cmap,
coff=coff,
lvec=lvec,
lmap=lmap)
cv2.imwrite(f'{prefix}.png', image)