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()