def main():
global args
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
state = torch.load(args.model_path)['state_dict']
model.load_state_dict(state)
model.eval()
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
train_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)])
f = open(args.img_list)
ni = 0
count_gt = [0, 0, 0, 0]
total_iou = [0, 0, 0, 0]
total_iou_big = [0, 0, 0, 0]
for line in f:
line = line.strip()
arrl = line.split(" ")
#gtlb = cv2.imread(args.gtpng_dir + arrl[1], -1)
gtlb = Image.open(args.gtpng_dir + arrl[1])
#print gtlb.shape
image = tv.datasets.folder.default_loader(args.img_dir + arrl[0])
image = image.resize((833, 705), Image.ANTIALIAS)
#image = cv2.imread(args.img_dir + arrl[0], -1)
#image = cv2.resize(image, (732,704), interpolation = cv2.INTER_NEAREST)
image = train_transform(image)
image = image.unsqueeze(0)
image = Variable(image.cuda(0), volatile=True)
output = model(image)
print output.size()
#output = F.log_softmax(output, dim=1)
prob = output.data[0].max(0)[1].cpu().numpy() #.transpose(1,0)
#prob.transpose((1,0))
gtlb = np.array(
gtlb.resize((prob.shape[1], prob.shape[0]), Image.ANTIALIAS))
#print gtlb.shape
#print gtlb[100:110,100:110]
#print prob[100:110,100:110]
#print output.max(),type(output)
iou = iou_one_frame(gtlb, prob)
print('IoU of ' + str(ni) + ' ' + arrl[0] + ': ' + str(iou))
for i in range(0, 4):
if iou[i] >= 0:
count_gt[i] = count_gt[i] + 1
total_iou[i] = total_iou[i] + iou[i]
mean_iou = np.divide(total_iou, count_gt)
print('Image numer: ' + str(ni))
print('Mean IoU of four lanes: ' + str(mean_iou))
print(
'Overall evaluation: ' + str(mean_iou[0] * 0.2 + mean_iou[1] * 0.3 +
mean_iou[2] * 0.3 + mean_iou[3] * 0.2))
f.close()
def main():
global args
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
state = torch.load(args.model_path)['state_dict']
model.load_state_dict(state)
model.eval()
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
f = open(args.img_list)
ni = 0
for line in f:
line = line.strip()
arrl = line.split(" ")
image = cv2.imread(args.img_dir + arrl[0])
image = cv2.resize(image, (833, 705)).astype(np.float32)
image -= [104, 117, 123]
image = image.transpose(2, 0, 1)
#image = cv2.imread(args.img_dir + arrl[0], -1)
#image = cv2.resize(image, (732,704), interpolation = cv2.INTER_NEAREST)
print image.shape
image = torch.from_numpy(image).unsqueeze(0)
image = Variable(image.float().cuda(0), volatile=True)
output = model(image)
#prob = output.data[0].max(0)[1].cpu().numpy()
#print prob.max(),prob.shape
output = torch.nn.functional.softmax(output[0], dim=0)
prob = output.data.cpu().numpy()
# prob = output.data[0].max(0)[1].cpu().numpy()
# print output.size()
#print output.max(),type(output)
print prob[1].max(), prob.shape
prob = (prob[1] >= 0.4)
# output = torch.nn.functional.softmax(output[0],dim=0)
# prob = output.data.cpu().numpy()
# print prob[1].max(),prob[2].max(),prob[3].max(),prob[4].max(),prob.shape
probAll = np.zeros((prob.shape[0], prob.shape[1], 3), dtype=np.float)
probAll[:, :, 0] += prob # line 1
probAll[:, :, 1] += prob # line 2
probAll[:, :, 2] += prob # line 3
probAll = np.clip(probAll * 255, 0, 255)
test_img = cv2.imread(args.img_dir + arrl[0], -1)
probAll = cv2.resize(probAll, (1280, 720),
interpolation=cv2.INTER_NEAREST)
test_img = cv2.resize(test_img, (1280, 720))
ni = ni + 1
test_img = np.clip(test_img + probAll, 0, 255).astype('uint8')
cv2.imwrite(args.img_dir + 'prob/test_' + str(ni) + '_lane.png',
test_img)
print('write img: ' + str(ni + 1))
f.close()
def main():
global args
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
state = torch.load(args.model_path)['state_dict']
model.load_state_dict(state)
model.eval()
mean = [104, 117, 123]
f = open(args.img_list)
ni = 0
count_gt = [0]
total_iou = [0]
total_iou_big = [0]
for line in f:
line = line.strip()
arrl = line.split(" ")
#gtlb = cv2.imread(args.gtpng_dir + arrl[1], -1)
gtlb = cv2.imread(args.gtpng_dir + arrl[1], -1)
#print gtlb.shape
image = cv2.imread(args.img_dir + arrl[0])
image = cv2.resize(image, (833, 705)).astype(np.float32)
image -= mean
image = image.transpose(2, 0, 1)
#image = cv2.imread(args.img_dir + arrl[0], -1)
#image = cv2.resize(image, (732,704), interpolation = cv2.INTER_NEAREST)
print image.shape
image = torch.from_numpy(image).unsqueeze(0)
image = Variable(image.float().cuda(0), volatile=True)
t = time.time()
output = model(image)
print time.time() - t
#output = F.log_softmax(output, dim=1)
output = torch.nn.functional.softmax(output[0], dim=0)
prob = output.data.cpu().numpy()
#prob.transpose((1,0))
gtlb = cv2.resize(gtlb, (prob.shape[2], prob.shape[1]),
interpolation=cv2.INTER_NEAREST)
gtlb = np.clip(gtlb, 0, 1)
#print gtlb.shape
#print gtlb[100:110,100:110]
#print prob[100:110,100:110]
#print output.max(),type(output)
iou = iou_one_frame(gtlb, prob)
print('IoU of ' + str(ni) + ' ' + arrl[0] + ': ' + str(iou))
for i in range(0, 1):
if iou >= 0:
count_gt[i] = count_gt[i] + 1
total_iou[i] = total_iou[i] + iou
mean_iou = np.divide(total_iou, count_gt)
print('Image numer: ' + str(ni))
print('Mean IoU of four lanes: ' + str(mean_iou))
f.close()
def main():
global args, best_prec
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
#model.apply(weights_init)
#params = torch.load('checkpoints/old.pth.tar')
#model.load_state_dict(params['state_dict'])
if not os.path.exists(args.resume):
os.makedirs(args.resume)
print("Saving everything to directory %s." % (args.resume))
# define loss function (criterion) and optimizer
criterion = cross_entropy2d
#criterion = torch.nn.DataParallel(criterion).cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
cudnn.benchmark = True
# data transform
train_data = MyDataset(
'/mnt/lustre/share/dingmingyu/cityscapes/instance_list.txt',
args.dir_path, args.new_width, args.new_length, args.label_width,
args.label_length)
train_loader = DataLoader(dataset=train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
for epoch in range(args.start_epoch, args.epochs):
print 'epoch: ' + str(epoch + 1)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
# remember best prec and save checkpoint
if (epoch + 1) % args.save_freq == 0:
checkpoint_name = "%03d_%s" % (epoch + 1, "checkpoint.pth.tar")
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, checkpoint_name, args.resume)
def main():
global args
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
state_dict = torch.load('lane_torch.pth')
model_dict = model.state_dict()
pretrained_dict = {
'm' + k: v
for k, v in state_dict.items() if 'm' + k in model_dict
}
model_dict.update(pretrained_dict)
print model_dict['m45.weight'].shape
model.load_state_dict(model_dict)
model = torch.nn.DataParallel(model).cuda()
model.eval()
mean = [103.939, 116.779, 123.68]
f = open(args.img_list)
ni = 0
count_gt = [0, 0, 0, 0]
total_iou = [0, 0, 0, 0]
total_iou_big = [0, 0, 0, 0]
for line in f:
line = line.strip()
arrl = line.split(" ")
#gtlb = cv2.imread(args.gtpng_dir + arrl[1], -1)
gtlb = Image.open(args.gtpng_dir + arrl[1])
#print gtlb.shape
image = tv.datasets.folder.default_loader(args.img_dir + arrl[0])
image = image.resize((836, 705), Image.ANTIALIAS)
image = np.array(image).astype('float')
image -= mean
image = image.transpose(2, 0, 1)
#image = cv2.imread(args.img_dir + arrl[0], -1)
#image = cv2.resize(image, (732,704), interpolation = cv2.INTER_NEAREST)
print image.shape
image = torch.from_numpy(image).unsqueeze(0)
image = Variable(image.float().cuda(0), volatile=True)
output = model(image)
print output.size()
#output = F.log_softmax(output, dim=1)
prob = output.data[0].max(0)[1].cpu().numpy() #.transpose(1,0)
#prob.transpose((1,0))
gtlb = np.array(
gtlb.resize((prob.shape[1], prob.shape[0]), Image.ANTIALIAS))
#print gtlb.shape
#print gtlb[100:110,100:110]
#print prob[100:110,100:110]
#print output.max(),type(output)
iou = iou_one_frame(gtlb, prob)
print('IoU of ' + str(ni) + ' ' + arrl[0] + ': ' + str(iou))
for i in range(0, 4):
if iou[i] >= 0:
count_gt[i] = count_gt[i] + 1
total_iou[i] = total_iou[i] + iou[i]
mean_iou = np.divide(total_iou, count_gt)
print('Image numer: ' + str(ni))
print('Mean IoU of four lanes: ' + str(mean_iou))
print(
'Overall evaluation: ' + str(mean_iou[0] * 0.2 + mean_iou[1] * 0.3 +
mean_iou[2] * 0.3 + mean_iou[3] * 0.2))
f.close()
def main():
global args
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
state = torch.load(args.model_path)['state_dict']
model.load_state_dict(state)
model.eval()
mean = [104, 117, 123]
f = open(args.img_list)
ni = 0
count_gt = [0, 0, 0, 0]
total_iou = [0, 0, 0, 0]
total_iou_big = [0, 0, 0, 0]
count_len = 0
count_true = 0
for line in f:
#if ni>1:
# break
line = line.strip()
arrl = line.split(" ")
#gtlb = cv2.imread(args.gtpng_dir + arrl[1], -1)
gtlb = cv2.imread(args.gtpng_dir + arrl[1], -1)
#print gtlb.shape
gt_num_list = list(np.unique(gtlb))
gt_num_list.remove(0)
n_objects_gt = len(gt_num_list)
image = cv2.imread(args.img_dir + arrl[0])
image = cv2.resize(image, (833, 705)).astype(np.float32)
image -= mean
image = image.transpose(2, 0, 1)
#image = cv2.imread(args.img_dir + arrl[0], -1)
#image = cv2.resize(image, (732,704), interpolation = cv2.INTER_NEAREST)
#print image.shape
image = torch.from_numpy(image).unsqueeze(0)
image = Variable(image.float().cuda(0), volatile=True)
output, embedding, n_objects_predictions = model(image)
output = torch.nn.functional.softmax(output[0], dim=0)
prob = output.data.cpu().numpy()
embedding = embedding.data.cpu().numpy()
n_objects_predictions = n_objects_predictions * 4
n_objects_predictions = torch.round(n_objects_predictions).int()
n_objects_predictions = int(n_objects_predictions.data.cpu())
print n_objects_predictions, '~~~~~~~~~~~~~~~~~~~~~~~`'
if n_objects_predictions == n_objects_gt:
count_true += 1
count_len += 1
if not n_objects_predictions:
continue
prob[prob >= args.prethreshold] = 1.0
prob[prob < args.prethreshold] = 0
embedding = embedding[0, :, :, :].transpose((1, 2, 0))
#print prob.shape
mylist = []
indexlist = []
for i in range(embedding.shape[0]):
for j in range(embedding.shape[1]):
if prob[1][i][j] > 0:
mylist.append(embedding[i, j, :])
indexlist.append((i, j))
if len(mylist) < 4:
continue
mylist = np.array(mylist)
# bandwidth = estimate_bandwidth(mylist, quantile=0.3, n_samples=100, n_jobs = 8)
# print bandwidth
# estimator = MeanShift(bandwidth=1, bin_seeding=True)
estimator = KMeans(n_clusters=n_objects_predictions)
#estimator = AffinityPropagation(preference=-0.4, damping = 0.5)
estimator.fit(mylist)
for i in range(4):
print len(estimator.labels_[estimator.labels_ == i])
#print len(np.unique(estimator.labels_)),'~~~~~~~~~~~~~~~~'
new_prob = np.zeros((embedding.shape[0], embedding.shape[1]),
dtype=int)
for index, item in enumerate(estimator.labels_):
if item <= 4:
new_prob[indexlist[index][0]][indexlist[index][1]] = item + 1
gtlb = cv2.resize(gtlb, (prob.shape[2], prob.shape[1]),
interpolation=cv2.INTER_NEAREST)
iou = iou_one_frame(gtlb, new_prob, args.prethreshold)
print('IoU of ' + str(ni) + ' ' + arrl[0] + ': ' + str(iou))
for i in range(0, 4):
if iou[i] >= 0:
count_gt[i] = count_gt[i] + 1
total_iou[i] = total_iou[i] + iou[i]
ni += 1
mean_iou = np.divide(total_iou, count_gt)
print('len_acc' + str(count_true * 1.0 / count_len))
print('Image numer: ' + str(ni))
print('Mean IoU of four lanes: ' + str(mean_iou))
print(
'Overall evaluation: ' + str(mean_iou[0] * 0.2 + mean_iou[1] * 0.3 +
mean_iou[2] * 0.3 + mean_iou[3] * 0.2))
f.close()
from torch.autograd import Variable
import torchvision
import sys
sys.path.append('/mnt/lustre/dingmingyu/software/core/python')
import caffe
caffe.set_mode_gpu()
caffe.mpi_init()
import os
from pytorch2caffe import pytorch2caffe, plot_graph
import lanenet
state_dict = torch.load('lane_torch.pth')
model = lanenet.Net()
model_dict = model.state_dict()
pretrained_dict = {
'm' + k: v
for k, v in state_dict.items() if 'm' + k in model_dict
}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model.eval()
print(model)
input_var = Variable(torch.rand(1, 3, 705, 836))
output_var = model(input_var)
output_dir = 'demo'
def main():
global args
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
state = torch.load(args.model_path)['state_dict']
model.load_state_dict(state)
model.eval()
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
train_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)])
f = open(args.img_list)
ni = 0
for line in f:
line = line.strip()
arrl = line.split(" ")
image = tv.datasets.folder.default_loader(args.img_dir + arrl[0])
image = image.resize((833, 705), Image.ANTIALIAS)
#image = cv2.imread(args.img_dir + arrl[0], -1)
#image = cv2.resize(image, (732,704), interpolation = cv2.INTER_NEAREST)
image = train_transform(image)
print image.shape
image = image.unsqueeze(0)
image = Variable(image.cuda(0), volatile=True)
output = model(image)
#output = F.log_softmax(output, dim=1)
output = torch.nn.functional.softmax(output[0], dim=0)
prob = output.data.cpu().numpy()
print output.size()
print prob[:, 1, 1]
#print output.max(),type(output)
prob[prob >= 0.3] = 1
prob[prob < 0.3] = 0
probAll = np.zeros((prob.shape[1], prob.shape[2], 3), dtype=np.float)
prob1 = prob[1]
prob2 = prob[2]
prob3 = prob[3]
prob4 = prob[4]
probAll[:, :, 0] += prob1 # line 1
probAll[:, :, 1] += prob2 # line 2
probAll[:, :, 2] += prob3 # line 3
probAll[:, :, 0] += prob4 # line 4
probAll[:, :, 1] += prob4 # line 4
probAll[:, :, 2] += prob4 # line 4
probAll = np.clip(probAll * 255, 0, 255)
test_img = cv2.imread(args.img_dir + arrl[0], -1)
probAll = cv2.resize(probAll, (833, 705),
interpolation=cv2.INTER_NEAREST)
test_img = cv2.resize(test_img, (833, 705))
ni = ni + 1
test_img = np.clip(test_img + probAll, 0, 255).astype('uint8')
cv2.imwrite(args.img_dir + 'prob/test_' + str(ni) + '_lane.png',
test_img)
print('write img: ' + str(ni + 1))
f.close()
def main():
global args, best_prec
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
#model.apply(weights_init)
#params = torch.load('checkpoints/001_single_module.pth.tar') # finetune from single channel
#bs, n , h, w = params['state_dict']['module.bottom.conv1.cbr_unit.0.weight'].size()
#params['state_dict']['module.bottom.conv1.cbr_unit.0.weight'] = torch.cat((params['state_dict']['module.bottom.conv1.cbr_unit.0.weight'],torch.randn((bs,1,h,w)).float().cuda()),1)
#model.load_state_dict(params['state_dict'])
params = torch.load('checkpoints/001_checkpoint.pth.tar')
model.load_state_dict(params['state_dict'])
args.start_epoch = params['epoch']
model_vgg = lanenet.VGG()
model_vgg = torch.nn.DataParallel(model_vgg).cuda()
if not os.path.exists(args.resume):
os.makedirs(args.resume)
print("Saving everything to directory %s." % (args.resume))
# define loss function (criterion) and optimizer
criterion = cross_entropy2d
#criterion_mse = torch.nn.DataParallel(torch.nn.MSELoss()).cuda() #this DataParallel is useless, there will still be gather but not parallel.
criterion_mse = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.load_state_dict(params['optimizer'])
#optimizer = nn.DataParallel(optimizer) # no use
cudnn.benchmark = True
# data transform
train_data = MyDataset('/mnt/lustre/share/dingmingyu/new_list_lane.txt',
args.dir_path, args.new_width, args.new_length,
args.label_width, args.label_length)
train_loader = DataLoader(dataset=train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
for epoch in range(args.start_epoch, args.epochs):
print 'epoch: ' + str(epoch + 1)
# train for one epoch
train(train_loader, model, model_vgg, criterion, criterion_mse,
optimizer, epoch)
# evaluate on validation set
# remember best prec and save checkpoint
if (epoch + 1) % args.save_freq == 0:
checkpoint_name = "%03d_%s" % (epoch + 1, "checkpoint.pth.tar")
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, checkpoint_name, args.resume)
def main():
global args
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
state_dict = torch.load('lane_torch.pth')
model_dict = model.state_dict()
pretrained_dict = {
'm' + k: v
for k, v in state_dict.items() if 'm' + k in model_dict
}
model_dict.update(pretrained_dict)
print model_dict['m45.weight'].shape
model.load_state_dict(model_dict)
model = torch.nn.DataParallel(model).cuda()
model.eval()
mean = [103.939, 116.779, 123.68]
f = open(args.img_list)
ni = 0
for line in f:
line = line.strip()
arrl = line.split(" ")
image = tv.datasets.folder.default_loader(args.img_dir + arrl[0])
image = image.resize((836, 705), Image.ANTIALIAS)
image = np.array(image).astype('float')
image -= mean
image = image.transpose(2, 0, 1)
#image = cv2.imread(args.img_dir + arrl[0], -1)
#image = cv2.resize(image, (732,704), interpolation = cv2.INTER_NEAREST)
print image.shape
image = torch.from_numpy(image).unsqueeze(0)
image = Variable(image.float().cuda(0), volatile=True)
output = model(image)
#output = F.log_softmax(output, dim=1)
output = torch.nn.functional.softmax(output[0], dim=0)
prob = output.data.cpu().numpy()
print output.size()
print prob[:, 1, 1]
#print output.max(),type(output)
prob[prob >= 0.3] = 1
prob[prob < 0.3] = 0
probAll = np.zeros((prob.shape[1], prob.shape[2], 3), dtype=np.float)
prob1 = prob[1]
prob2 = prob[2]
prob3 = prob[3]
prob4 = prob[4]
probAll[:, :, 0] += prob1 # line 1
probAll[:, :, 1] += prob2 # line 2
probAll[:, :, 2] += prob3 # line 3
probAll[:, :, 0] += prob4 # line 4
probAll[:, :, 1] += prob4 # line 4
probAll[:, :, 2] += prob4 # line 4
probAll = np.clip(probAll * 255, 0, 255)
test_img = cv2.imread(args.img_dir + arrl[0], -1)
probAll = cv2.resize(probAll, (833, 705),
interpolation=cv2.INTER_NEAREST)
test_img = cv2.resize(test_img, (833, 705))
ni = ni + 1
test_img = np.clip(test_img + probAll, 0, 255).astype('uint8')
cv2.imwrite(args.img_dir + 'prob/test_' + str(ni) + '_lane.png',
test_img)
print('write img: ' + str(ni + 1))
f.close()
def main():
global args, best_prec
args = parser.parse_args()
print("Build model ...")
params = torch.load(
'/mnt/lustre/dingmingyu/workspace/Pytorch/checkpoints_size/old.pth.tar'
)
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(params['state_dict'])
if not os.path.exists(args.resume):
os.makedirs(args.resume)
print("Saving everything to directory %s." % (args.resume))
# define loss function (criterion) and optimizer
criterion = cross_entropy2d
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
cudnn.benchmark = True
# data transform
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
train_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)])
gt_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x * 255),
transforms.Lambda(lambda x: numpy.clip(x, 0, 4))
])
train_data = MyDataset('/mnt/lustre/share/dingmingyu/new_list_lane.txt',
args.dir_path, args.new_width, args.new_length,
args.label_width, args.label_length,
train_transform, gt_transform)
train_loader = DataLoader(dataset=train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
for epoch in range(args.start_epoch, args.epochs):
print 'epoch: ' + str(epoch + 1)
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
# remember best prec and save checkpoint
if (epoch + 1) % args.save_freq == 0:
checkpoint_name = "%03d_%s" % (epoch + 1, "checkpoint.pth.tar")
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, checkpoint_name, args.resume)
def main():
global args
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
state = torch.load(args.model_path)['state_dict']
model.load_state_dict(state)
model.eval()
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
f = open(args.img_list)
ni = 0
for line in f:
line = line.strip()
arrl = line.split(" ")
image = cv2.imread(args.img_dir + arrl[0])
image = cv2.resize(image, (833, 705)).astype(np.float32)
image -= [104, 117, 123]
image = image.transpose(2, 0, 1)
#image = cv2.imread(args.img_dir + arrl[0], -1)
#image = cv2.resize(image, (732,704), interpolation = cv2.INTER_NEAREST)
#print image.shape
image = torch.from_numpy(image).unsqueeze(0)
image = Variable(image.float().cuda(0), volatile=True)
output, embedding, n_objects_predictions = model(image)
#prob = output.data[0].max(0)[1].cpu().numpy()
#print prob.max(),prob.shape
output = torch.nn.functional.softmax(output[0], dim=0)
prob = output.data.cpu().numpy()
# prob = output.data[0].max(0)[1].cpu().numpy()
# print output.size()
#print output.max(),type(output)
print prob[1].max(), prob.shape
prob = (prob[1] >= 0.4)
n_objects_predictions = n_objects_predictions * 4
n_objects_predictions = torch.round(n_objects_predictions).int()
n_objects_predictions = int(n_objects_predictions.data.cpu())
embedding = embedding.data.cpu().numpy()
embedding = embedding[0, :, :, :].transpose((1, 2, 0))
mylist = []
indexlist = []
for i in range(embedding.shape[0]):
for j in range(embedding.shape[1]):
if prob[i][j] > 0:
mylist.append(embedding[i, j, :])
indexlist.append((i, j))
mylist = np.array(mylist)
#bandwidth = estimate_bandwidth(mylist, quantile=0.3, n_samples=100, n_jobs = 8)
#print bandwidth
estimator = MeanShift(bandwidth=1.5, bin_seeding=True)
#estimator = KMeans(n_clusters = n_objects_predictions)
estimator.fit(mylist)
print len(np.unique(estimator.labels_)), '~~~~~~~~~~~~~~~~'
for i in range(4):
print len(estimator.labels_[estimator.labels_ == i]), ' ',
probAll = np.zeros((prob.shape[0], prob.shape[1], 3), dtype=np.float)
# probAll[:,:,0] += prob # line 1
# probAll[:,:,1] += prob # line 2
# probAll[:,:,2] += prob # line 3
for index, item in enumerate(estimator.labels_):
x = indexlist[index][0]
y = indexlist[index][1]
if item < 3:
probAll[x, y, item] += prob[x, y] # line 1
else:
probAll[x, y, :] += 1
probAll = np.clip(probAll * 255, 0, 255)
test_img = cv2.imread(args.img_dir + arrl[0], -1)
probAll = cv2.resize(probAll, (1280, 720),
interpolation=cv2.INTER_NEAREST)
test_img = cv2.resize(test_img, (1280, 720))
ni = ni + 1
test_img = np.clip(test_img + probAll, 0, 255).astype('uint8')
cv2.imwrite(args.img_dir + 'prob/test_' + str(ni) + '_lane.png',
test_img)
print('write img: ' + str(ni + 1))
f.close()
def main():
global args
args = parser.parse_args()
print("Build model ...")
model = lanenet.Net()
model = torch.nn.DataParallel(model).cuda()
state = torch.load(args.model_path)['state_dict']
model.load_state_dict(state)
model.eval()
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
f = open(args.img_list)
ni = 0
for line in f:
if ni > 100:
break
line = line.strip()
arrl = line.split(" ")
image = cv2.imread(arrl[0]).astype(np.float32)
img = image[:, :, :3]
img -= [104, 117, 123]
h, w = img.shape[:2]
m = max(w, h)
ratio = 112.0 / m
new_w, new_h = int(ratio * w), int(ratio * h)
assert new_w > 0 and new_h > 0
img = cv2.resize(img, (new_w, new_h))
W, H = 112, 112
top = (H - new_h) // 2
bottom = (H - new_h) // 2
if top + bottom + new_h < H:
bottom += 1
left = (W - new_w) // 2
right = (W - new_w) // 2
if left + right + new_w < W:
right += 1
img = cv2.copyMakeBorder(img,
top,
bottom,
left,
right,
cv2.BORDER_CONSTANT,
value=[0, 0, 0])
image = image.transpose(2, 0, 1)
#image = cv2.imread(args.img_dir + arrl[0], -1)
#image = cv2.resize(image, (732,704), interpolation = cv2.INTER_NEAREST)
print image.shape
image = torch.from_numpy(image).unsqueeze(0)
image = Variable(image.float().cuda(0), volatile=True)
if 'car' in arrl[0]:
output = model(image)[1]
print('car', str(ni))
else:
output = model(image)[0]
print('person', str(ni))
#print output.size()
#prob = output.data[0].max(0)[1].cpu().numpy()
#print prob.max(),prob.shape
output = torch.nn.functional.softmax(output[0], dim=0)
prob = output.data.cpu().numpy()
# prob = output.data[0].max(0)[1].cpu().numpy()
# print output.size()
#print output.max(),type(output)
#print prob,prob.shape
prob = (prob[1] >= 0.6)
# output = torch.nn.functional.softmax(output[0],dim=0)
# prob = output.data.cpu().numpy()
# print prob[1].max(),prob[2].max(),prob[3].max(),prob[4].max(),prob.shape
probAll = np.zeros((prob.shape[0], prob.shape[1], 3), dtype=np.float)
probAll[:, :, 0] += prob # line 1
probAll[:, :, 1] += prob # line 2
probAll[:, :, 2] += prob # line 3
probAll = np.clip(probAll * 255, 0, 255)
probAll = cv2.resize(probAll, (112, 112),
interpolation=cv2.INTER_NEAREST)
ni = ni + 1
test_img = np.clip(probAll, 0, 255).astype('uint8')
image = cv2.imread(arrl[0])
cv2.imwrite(args.img_dir + 'test_' + str(ni) + '_ori.png', image)
cv2.imwrite(args.img_dir + 'test_' + str(ni) + '_lane.png', test_img)
print('write img: ' + str(ni + 1))
f.close()
