def main():
args = parse_args()
img_path = args.img_path
weight_path = args.weight_path
_set = "IMAGENET"
mean = IMG_MEAN[_set]
std = IMG_STD[_set]
transform_img = Resize((800, 288))
transform_x = Compose(ToTensor(), Normalize(mean=mean, std=std))
transform = Compose(transform_img, transform_x)
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # RGB for net model input
img = transform_img({'img': img})['img']
x = transform_x({'img': img})['img']
x.unsqueeze_(0)
net = LaneNet(pretrained=False, embed_dim=4, delta_v=.5, delta_d=3.)
save_dict = torch.load(weight_path, map_location='cpu')
net.load_state_dict(save_dict['net'])
net.eval()
output = net(x)
embedding = output['embedding']
embedding = embedding.detach().cpu().numpy()
embedding = np.transpose(embedding[0], (1, 2, 0))
binary_seg = output['binary_seg']
bin_seg_prob = binary_seg.detach().cpu().numpy()
bin_seg_pred = np.argmax(bin_seg_prob, axis=1)[0]
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
seg_img = np.zeros_like(img)
lane_seg_img = embedding_post_process(embedding, bin_seg_pred,
args.band_width, 4)
color = np.array([[255, 125, 0], [0, 255, 0], [0, 0, 255], [0, 255, 255]],
dtype='uint8')
for i, lane_idx in enumerate(np.unique(lane_seg_img)):
if lane_idx == 0:
continue
seg_img[lane_seg_img == lane_idx] = color[i - 1]
img = cv2.addWeighted(src1=seg_img, alpha=0.8, src2=img, beta=1., gamma=0.)
cv2.imwrite("demo/demo_result.jpg", img)
if args.visualize:
cv2.imshow("", img)
cv2.waitKey(0)
cv2.destroyAllWindows()
mean=(0.485, 0.456, 0.406)
std=(0.229, 0.224, 0.225)
transform = Compose(Resize(resize_shape), ToTensor(),
Normalize(mean=mean, std=std))
dataset_name = exp_cfg['dataset'].pop('dataset_name')
Dataset_Type = getattr(dataset, dataset_name)
test_dataset = Dataset_Type(Dataset_Path['Tusimple'], "test", transform)
test_loader = DataLoader(test_dataset, batch_size=32, collate_fn=test_dataset.collate, num_workers=4)
net = LaneNet(pretrained=True, **exp_cfg['model'])
save_name = os.path.join(exp_dir, exp_dir.split('/')[-1] + '_best.pth')
save_dict = torch.load(save_name, map_location='cpu')
print("\nloading", save_name, "...... From Epoch: ", save_dict['epoch'])
net.load_state_dict(save_dict['net'])
net = torch.nn.DataParallel(net.to(device))
net.eval()
# ------------ test ------------
out_path = os.path.join(exp_dir, "coord_output")
evaluation_path = os.path.join(exp_dir, "evaluate")
if not os.path.exists(out_path):
os.mkdir(out_path)
if not os.path.exists(evaluation_path):
os.mkdir(evaluation_path)
dump_to_json = []
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']