os.mkdir(evaluation_path)
dump_to_json = []
progressbar = tqdm(range(len(test_loader)))
with torch.no_grad():
for batch_idx, (imgs, _, img_name) in enumerate(test_loader):
imgs = imgs.to(device)
seg_pred = net(imgs)
seg_pred = F.softmax(seg_pred, dim=1)
seg_pred = seg_pred.detach().cpu().numpy()
for b in range(len(seg_pred)):
seg = seg_pred[b]
exist = [1 for i in range(4)]
lane_coords = getLane.prob2lines_tusimple(seg, exist, smooth=False, thresh=0.5, resize_shape=(720, 1280), y_px_gap=10, pts=56)
for i in range(len(lane_coords)):
lane_coords[i] = sorted(lane_coords[i], key=lambda pair: pair[1])
path_tree = split_path(img_name[b])
save_dir, save_name = path_tree[-3:-1], path_tree[-1]
save_dir = os.path.join(out_path, *save_dir)
save_name = save_name[:-3] + "lines.txt"
save_name = os.path.join(save_dir, save_name)
if not os.path.exists(save_dir):
os.makedirs(save_dir, exist_ok=True)
with open(save_name, "w") as f:
for l in lane_coords:
for (x, y) in l:
print("{} {}".format(x, y), end=" ", file=f)
progressbar = tqdm(range(len(test_loader)))
with torch.no_grad():
for batch_idx, sample in enumerate(test_loader):
img = sample['img'].to(device)
img_name = sample['img_name']
seg_pred, exist_pred = net(img)[:2]
seg_pred = seg_pred.detach().cpu().numpy()
exist_pred = exist_pred.detach().cpu().numpy()
for b in range(len(seg_pred)):
seg = seg_pred[b]
exist = [1 if exist_pred[b, i] > 0.5 else 0 for i in range(4)]
lane_coords = getLane.prob2lines_tusimple(seg,
exist,
resize_shape=(720, 1280),
y_px_gap=10,
pts=56)
for i in range(len(lane_coords)):
lane_coords[i] = sorted(lane_coords[i],
key=lambda pair: pair[1])
path_tree = split_path(img_name[b])
save_dir, save_name = path_tree[-3:-1], path_tree[-1]
save_dir = os.path.join(out_path, *save_dir)
save_name = save_name[:-3] + "lines.txt"
save_name = os.path.join(save_dir, save_name)
if not os.path.exists(save_dir):
os.makedirs(save_dir, exist_ok=True)
with open(save_name, "w") as f:
img_vis = img.copy()
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = transform({'img': img})['img']
img = img.unsqueeze(0).to(device)
with torch.no_grad():
seg_pred, exist_pred = net(img)[:2]
seg_pred = F.softmax(seg_pred, dim=1)
seg_pred = seg_pred.cpu().numpy()
exist_pred = exist_pred.cpu().numpy()
b=0
seg = seg_pred[b]
exist = [1 if exist_pred[b, i] > 0.5 else 0 for i in range(4)]
if dataset_name == 'Tusimple':
lane_coords = getLane.prob2lines_tusimple(seg, exist, resize_shape=original_shape[::-1], y_px_gap=10, pts=56)
elif dataset_name == 'CULane':
lane_coords = getLane.prob2lines_CULane(seg, exist, resize_shape=original_shape[::-1], y_px_gap=20, pts=18)
for i in range(len(lane_coords)):
lane_coords[i] = sorted(lane_coords[i], key=lambda pair: pair[1])
for lane in lane_coords:
foo = np.int32([lane])[:,np.int32([lane])[0,:,0] >= 0]
cv2.polylines(img_vis, foo, isClosed=False, color=(0,0,255), thickness=2)
#show segmentation map
seg = seg.swapaxes(0,2).swapaxes(0,1)
seg = cv2.resize(seg, dsize=original_shape, interpolation=cv2.INTER_CUBIC)
lane_img = np.zeros_like(img_vis)
color = np.array([[255, 125, 0], [0, 255, 0], [0, 0, 255], [0, 255, 255]], dtype='uint8')