def validate(args): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, split=args.split, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes valloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4) running_metrics = runningScore(n_classes) # Setup Model model = get_model(args.model_path[:args.model_path.find('_')], n_classes) state = convert_state_dict(torch.load(args.model_path)['model_state']) model.load_state_dict(state) model.eval() for i, (images, labels) in tqdm(enumerate(valloader)): model.cuda() images = Variable(images.cuda(), volatile=True) labels = Variable(labels.cuda(), volatile=True) outputs = model(images) pred = outputs.data.max(1)[1].cpu().numpy() gt = labels.data.cpu().numpy() running_metrics.update(gt, pred) score, class_iou = running_metrics.get_scores() for k, v in score.items(): print(k, v) for i in range(n_classes): print(i, class_iou[i])
def train(args): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset, config_file=args.config_file) loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes # must use 1 worker for AWS sagemaker without ipc="host" or larger shared memory size trainloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=1, shuffle=True) # Setup Model model = get_model(args.arch, n_classes) # Setup log dir / logging if not os.path.exists(args.log_dir): os.makedirs(args.log_dir) configure(args.log_dir) model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count())) model.cuda() optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.9, weight_decay=5e-4) step = 0 for epoch in range(args.n_epoch): start_time = time.time() for i, (images, labels) in enumerate(trainloader): images = Variable(images.cuda()) labels = Variable(labels.cuda()) optimizer.zero_grad() outputs = model(images) loss = cross_entropy2d(outputs, labels) loss.backward() optimizer.step() log_value('Loss', loss.data[0], step) step += 1 if (i + 1) % 20 == 0: print("Epoch [%d/%d] Loss: %.4f" % (epoch + 1, args.n_epoch, loss.data[0]), flush=True) end_time = time.time() print('Epoch run time: %s' % (end_time - start_time)) torch.save( model, args.log_dir + "{}_{}_{}_{}.pt".format( args.arch, args.dataset, args.feature_scale, epoch))
def get_dataset_loader(): args=edict() args.dataset_name = 'cityscapes' args.config_path = os.path.join('/home/yzbx/git/gnu/pytorch-semseg', 'config.json') args.img_rows=224 args.img_cols=224 args.img_norm=True args.batch_size=2 data_loader = get_loader(args.dataset_name) data_path = get_data_path(args.dataset_name, args.config_path) t_loader = data_loader(data_path, is_transform=True, split='train', img_size=( args.img_rows, args.img_cols), augmentations=None, img_norm=args.img_norm) v_loader = data_loader(data_path, is_transform=True, split='val', img_size=( args.img_rows, args.img_cols), img_norm=args.img_norm) n_classes = t_loader.n_classes print('class number is',n_classes) trainloader = data.DataLoader( t_loader, batch_size=args.batch_size, num_workers=8, shuffle=True) valloader = data.DataLoader( v_loader, batch_size=args.batch_size, num_workers=8) return trainloader,valloader
def test(args, cfg): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_file_name = os.path.split(args.model_path)[1] model_name = model_file_name[:model_file_name.find("_")] IMG_Path = Path(args.img_path) IMG_File = natsort.natsorted(list(IMG_Path.glob("*.png")), alg=natsort.PATH) IMG_Str = [] for i in IMG_File: IMG_Str.append(str(i)) # Setup image print("Read Input Image from : {}".format(args.img_path)) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset, config_file=cfg) loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm) n_classes = loader.n_classes # Setup Model model = get_model(cfg['model'], n_classes) state = convert_state_dict(torch.load(args.model_path)["model_state"]) # state=torch.load(args.model_path)["model_state"] model.load_state_dict(state) model.eval() model.to(device) for j in tqdm(range(len(IMG_Str))): img_path = IMG_Str[j] img = misc.imread(img_path) # img = img[:, :, ::-1] img = img.astype(np.float64) # img -= loader.mean if args.img_norm: img = img.astype(float) / 255.0 # NHWC -> NCHW img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() images = img.to(device) outputs = model(images) outputs_probability = F.softmax(outputs) data = outputs_probability.data data_max = data.max(1) prob = data_max[0] prob_img_format = np.squeeze(prob.cpu().numpy(), axis=0) avg_prob = np.mean(prob_img_format) print("Confidence Score for %s: \n%f" % (img_path, avg_prob)) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0) decoded = loader.decode_segmap(pred) out_path = "test_out/test_confidence/out/" + Path(img_path).name decoded_bgr = cv.cvtColor(decoded, cv.COLOR_RGB2BGR) # misc.imsave(out_path, decoded) cv.imwrite(out_path, decoded_bgr)
def test(args, cfg): os.environ["CUDA_VISIBLE_DEVICES"] = "1" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_file_name = os.path.split(args.model_path)[1] model_name = model_file_name[:model_file_name.find("_")] IMG_Path = Path(args.img_path) IMG_File = natsort.natsorted(list(IMG_Path.glob("*.tif")), alg=natsort.PATH) IMG_Str = [] for i in IMG_File: IMG_Str.append(str(i)) # Setup image print("Read Input Image from : {}".format(args.img_path)) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset, config_file=cfg) loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm) n_classes = loader.n_classes v_loader = data_loader( data_path, is_transform=True, split=cfg['data']['val_split'], img_size=(cfg['data']['img_rows'], cfg['data']['img_cols']), ) valloader = data.DataLoader(v_loader, batch_size=cfg['training']['batch_size'], num_workers=cfg['training']['n_workers']) # Setup Model model = get_model(cfg['model'], n_classes) state = convert_state_dict(torch.load(args.model_path)["model_state"]) # state=torch.load(args.model_path)["model_state"] model.load_state_dict(state) model.eval() model.to(device) with torch.no_grad(): for i_val, (img_path, images_val, labels_val) in tqdm(enumerate(valloader)): img_name = img_path[0] images_val = images_val.to(device) outputs = model(images_val) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0) decoded = loader.decode_segmap(pred) out_path = "test_out/CAN_res50_4band_data07/" + Path( img_name).stem + ".png" decoded_bgr = cv.cvtColor(decoded, cv.COLOR_RGB2BGR) # misc.imsave(out_path, decoded) cv.imwrite(out_path, decoded_bgr)
def process_args_from_loaded_cfg(cfg, args): """ process arguments from loaded CFG. :param cfg: :param args: :return: """ # If calling this, it will use CFG for sure. # if 'runet' in cfg['model']['arch']: # args.hidden_size = cfg['model']['hidden_size'] # r = re.compile('-h\d+-') r_d = re.compile('\d+') _h = r.findall(cfg['logdir']) if len(_h) > 0: res = int(r_d.findall(_h[0])[0]) if res > 1: # avoid the represent h=init args.hidden_size = res else: res = args.hidden_size if cfg['model'].get('hidden_size'): assert cfg['model']['hidden_size'] == res else: cfg['model']['hidden_size'] = res args.gate = cfg['model'].get('gate') or args.gate if args.is_recurrent is not None: args.is_recurrent = cfg['training']['loss']['name'] \ in ['multi_step_cross_entropy'] out_path = args.out_path or os.path.join( 'results', cfg['data']['dataset'], os.path.basename(os.path.dirname(cfg['logdir']))) cfg['eval_out_path'] = out_path # Process for unet_level if cfg['model']['arch'] == "runet": if cfg['model'].get('unet_level'): args.unet_level = cfg['model']['unet_level'] else: unet_level = args.hidden_size // 32 args.unet_level = unet_level cfg['model']['unet_level'] = args.unet_level cfg['model']['recurrent_level'] = args.recurrent_level if not os.path.exists(cfg['data']['path']): cfg['data']['path'] = get_data_path( cfg['data']['path'], config_file=f"configs/dataset/{cfg['data']['dataset'].replace('_', '')}.yml") IPython.embed() return cfg, args
def setup(self, pre_encode=False): sbd_path = get_data_path('sbd') voc_path = get_data_path('pascal') target_path = self.root + '/SegmentationClass/pre_encoded/' if not os.path.exists(target_path): os.makedirs(target_path) # Load SBD train set sbd_train_list = tuple(open(sbd_path + 'dataset/train.txt', 'r')) sbd_train_list = [id_.rstrip() for id_ in sbd_train_list] # Load SBD val set sbd_val_list = tuple(open(sbd_path + 'dataset/val.txt', 'r')) sbd_val_list = [id_.rstrip() for id_ in sbd_val_list] # Join everything self.files[ 'train_aug'] = self.files['train'] + sbd_train_list + sbd_val_list # Remove duplicates and intersection with Pascal VOC validation set self.files['train_aug'] = list( set(self.files['train_aug']) - set(self.files['val'])) self.files['train_aug'].sort() if pre_encode: print("Pre-encoding segmentation masks...") lbl_dir = os.path.join(sbd_path, 'dataset', 'cls') lbl_list = [f for f in os.listdir(lbl_dir) if f.endswith('.mat')] for i in tqdm(lbl_list): lbl_path = os.path.join(lbl_dir, i) lbl = io.loadmat( lbl_path)['GTcls'][0]['Segmentation'][0].astype(np.int32) lbl = m.toimage(lbl, high=lbl.max(), low=lbl.min()) m.imsave(target_path + os.path.splitext(i)[0] + '.png', lbl) for i in tqdm(self.files['trainval']): lbl_path = self.root + '/SegmentationClass/' + i + '.png' lbl = self.encode_segmap(m.imread(lbl_path)) lbl = m.toimage(lbl, high=lbl.max(), low=lbl.min()) m.imsave(target_path + i + '.png', lbl)
def validate(args): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, split=args.split, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes valloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4) # Setup Model model = torch.load(args.model_path) model.eval() if torch.cuda.is_available(): model.cuda(args.gpu) gts, preds = [], [] for i, (images, labels) in tqdm(enumerate(valloader)): if torch.cuda.is_available(): images = Variable(images.cuda(args.gpu)) labels = Variable(labels.cuda(args.gpu)) else: images = Variable(images) labels = Variable(labels) outputs = model(images) pred = outputs.data.max(1)[1].cpu().numpy() pred = np.squeeze(pred) pred = cv2.resize(pred, labels.size()[1:][::-1], interpolation=cv2.INTER_NEAREST) pred = np.expand_dims(pred, axis=0) gt = labels.data.cpu().numpy() for gt_, pred_ in zip(gt, pred): gts.append(gt_) preds.append(pred_) score, class_iou = scores(gts, preds, n_class=n_classes) for k, v in score.items(): print k, v for i in range(n_classes): print i, class_iou[i]
def test(args): model_file_name = os.path.split(args.model_path)[1] model_name = model_file_name[:model_file_name.find('_')] print("Building " + model_name) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, phase='test') im_paths = loader.im_paths() n_classes = loader.n_classes testloader = data.DataLoader(loader, batch_size=1, num_workers=1, shuffle=False) # Setup Model model = get_model(model_name, n_classes) state = torch.load(args.model_path)['model_state'] model.load_state_dict(state) model.eval() model.cuda() # Run test for KITTI Road dataset for i, (image, tr_image, lidar, tr_lidar) in enumerate(testloader): im_name_splits = im_paths[i].split('/')[-1].split('.')[0].split('_') task = im_name_splits[0] print('processing %d-th image' % i) t0 = time.time() orig_h, orig_w = image.shape[1:3] with torch.no_grad(): tr_image = Variable(tr_image.cuda()) tr_lidar = Variable(tr_lidar.cuda()) outputs = model([tr_image, tr_lidar]) outputs = outputs.cpu().numpy().transpose((2, 3, 1, 0)).squeeze() outputs = cv2.resize(outputs, (orig_w, orig_h)) outputs = outputs[:, :, 1] print('Time({:d}'.format(i) + ') {0:.3f}'.format(time.time() - t0)) output_fg = outputs * 255. output_fg[output_fg > 255] = 255 output_fg = output_fg.astype(np.uint8) cv2.imwrite( './outputs/results/' + im_name_splits[0] + '_road_' + im_name_splits[1] + '.png', output_fg) print('write to ./outputs/results/' + im_name_splits[0] + '_road_' + im_name_splits[1] + '.png')
def validate(): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, img_size=args.img_size) n_classes = loader.n_classes n_channels = loader.n_channels valloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) # Setup Model model = torch.load(args.model_path) model.eval() if torch.cuda.is_available(): model.cuda(0) gts, preds = [], [] for i, (images, labels) in tqdm(enumerate(valloader)): if i >= args.max_samples: break if torch.cuda.is_available(): images = Variable(images.cuda(0)) labels = Variable(labels.cuda(0)) else: images = Variable(images) labels = Variable(labels) outputs = model(images) pred = np.squeeze((torch.max(outputs.data, 1, keepdim=True))[1].cpu().numpy()) gt = np.squeeze(labels.data.cpu().numpy()) for gt_, pred_ in zip(gt, pred): gts.append(gt_) preds.append(pred_) score, class_iou = scores(gts, preds, n_class=n_classes) for k, v in score.items(): print(k, v) for i in range(n_classes): print(i, class_iou[i])
def test(args): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_file_name = os.path.split(args.model_path)[1] model_name = model_file_name[:model_file_name.find("_")] # Setup image print("Read Input Image from : {}".format(args.img_path)) img = misc.imread(args.img_path) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm) n_classes = loader.n_classes img = img[:, :, ::-1] img = img.astype(np.float64) img -= loader.mean if args.img_norm: img = img.astype(float) / 255.0 # NHWC -> NCHW img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() # Setup Model model_dict = {"arch": model_name} model = get_model(model_dict, n_classes, version=args.dataset) state = convert_state_dict(torch.load(args.model_path)["model_state"]) model.load_state_dict(state) model.eval() model.to(device) images = img.to(device) outputs = model(images) if args.mask_path: print("Read Image Mask from : {}".format(args.mask_path)) mask = torch.load(args.mask_path) mask = mask.to(device) outputs = to_super_to_pixels(outputs, mask) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0) decoded = loader.decode_segmap(pred) print("Classes found: ", np.unique(pred)) misc.imsave(args.out_path, decoded) print("Segmentation Mask Saved at: {}".format(args.out_path))
def test(args): # Setup image print("Read Input Image from : ", args.img_path) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True) n_classes = loader.n_classes img = img[:, :, ::-1] img = img.astype(np.float64) img -= loader.mean # print(loader.img_size[0], loader.img_size[1]) # img = misc.imresize(img, (loader.img_size[0], loader.img_size[1])) img = img.astype(float) / 255.0 # NHWC -> NCWH img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() # Setup Model model = torch.load(args.model_path) model.eval() if torch.cuda.is_available(): model.cuda(0) images = Variable(img.cuda(0)) else: images = Variable(img) outputs = model(images) log('the size of outputs is '.format(outputs.size())) log('the size of outputs is ', outputs.data.max(1)[1].size()) pred = np.squeeze(outputs.data.cpu().numpy(), axis=0) pred = pred[0, :, :] > pred[1, :, :] pred = pred * 255 pred = outputs.data.cpu().numpy() log('The prediction shape is {}'.format(pred.shape)) # print('the size of pred is ', pred.shape()) # fig = plt.figure() # plt.imshow(pred[1,:,:]) # plt.title('testing result') # plt.show() pred = pred * 255 misc.imsave(args.out_path, pred[0, :, :]) print("Segmentation Mask Saved at: ", args.out_path)
def test(args): # Setup image print("Read Input Image from : {}".format(args.img_path)) orig_img = misc.imread(args.img_path) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True) n_classes = loader.n_classes img = orig_img[:, :, ::-1] img = img.astype(np.float64) img -= loader.mean img = misc.imresize(img, (loader.img_size[0], loader.img_size[1])) img = img.astype(float) / 255.0 # NHWC -> NCWH img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() # Setup Model model = get_model(args.arch, n_classes) model.load_state_dict(torch.load(args.model_path)['state_dict']) model = torch.nn.DataParallel(model, device_ids=range( torch.cuda.device_count())).cuda() model.eval() if torch.cuda.is_available(): model.cuda(0) images = Variable(img.cuda(0)) else: images = Variable(img) outputs = model(images) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0) decoded = loader.decode_segmap(pred) if args.alpha_blend: orig_img = misc.imresize(orig_img, (loader.img_size[0], loader.img_size[1])) out_img = ALPHA * orig_img + (1 - ALPHA) * decoded else: out_img = decoded print(np.unique(pred)) misc.imsave(args.out_path, out_img) print("Segmentation Mask Saved at: {}".format(args.out_path))
def validate(args): data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, split=args.split, is_transform=True, img_size=None, img_norm=args.img_norm) n_classes = loader.n_classes valloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4) running_metrics = runningScore(n_classes) # Setup Model from pytorchgo.model.deeplab_resnet import Res_Deeplab model = Res_Deeplab(NoLabels=n_classes, pretrained=False) state = torch.load("/home/hutao/MS_DeepLab_resnet_trained_VOC.pth") model.load_state_dict(state) model.eval() model.cuda() print "pra" for i, (images, labels) in tqdm(enumerate(valloader), desc="validation"): start_time = timeit.default_timer() img_large = torch.Tensor(np.zeros((1, 3, 513, 513))) img_large[:, :, :images.shape[2], :images.shape[3]] = images output = model(Variable(img_large, volatile=True).cuda()) output = output.data.max(1)[1].cpu().numpy() pred = output[:, :images.shape[2], :images.shape[3]] gt = labels.numpy() if args.measure_time: elapsed_time = timeit.default_timer() - start_time print('Inference time (iter {0:5d}): {1:3.5f} fps'.format( i + 1, pred.shape[0] / elapsed_time)) running_metrics.update(gt, pred) score, class_iou = running_metrics.get_scores() for k, v in score.items(): print(k, v) for i in range(n_classes): print(i, class_iou[i])
def validate(args): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, split=args.split, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes valloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4) # Setup Model model = torch.load(args.model_path) model.eval() gts, preds = [], [] for i, (images, labels) in tqdm(enumerate(valloader)): if torch.cuda.is_available(): model = torch.nn.DataParallel(model, device_ids=range( torch.cuda.device_count())) images = Variable(images.cuda(0)) labels = Variable(labels.cuda(0)) else: images = Variable(images) labels = Variable(labels) outputs = model(images) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=1) gt = labels.data.cpu().numpy() for gt_, pred_ in zip(gt, pred): gts.append(gt_) preds.append(pred_) score, class_iou = scores(gts, preds, n_class=n_classes) for k, v in score.items(): print(k, v) for i in range(n_classes): print(i, class_iou[i])
def test(args): args.img_path = '/home/shehabk/dataSets/CKPLUS/unet_256/A/test/S077_007_00000029.png' args.dataset = 'ckplus' args.out_path = '/home/shehabk/Desktop/' + os.path.basename(args.img_path) # args.arch = 'segnet' args.model_path = 'unet2_ckplus_1_99.pkl' # Setup image print("Read Input Image from : {}".format(args.img_path)) img = misc.imread(args.img_path,mode = 'RGB') data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True) n_classes = loader.n_classes img = img[:, :, ::-1] img = img.astype(np.float64) img -= loader.mean img = misc.imresize(img, (loader.img_size[0], loader.img_size[1])) img = img.astype(float) / 255.0 # NHWC -> NCWH img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() # Setup Model model = torch.load(args.model_path) model.eval() if torch.cuda.is_available(): model.cuda(0) images = Variable(img.cuda(0)) else: images = Variable(img) outputs = model(images) pred = np.squeeze(outputs.data.max(1 , keepdim=True )[1].cpu().numpy(), axis=1) decoded = loader.decode_segmap(pred[0]) print(np.unique(pred)) misc.imsave(args.out_path, decoded) print("Segmentation Mask Saved at: {}".format(args.out_path))
def test(args, cfg): os.environ["CUDA_VISIBLE_DEVICES"] = "0" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") data_loader = get_loader(cfg['data']['dataset']) data_path = get_data_path(cfg['data']['dataset'], config_file=cfg) loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm) n_classes = loader.n_classes t_loader = data_loader( data_path, is_transform=True, split='test', img_size=(cfg['data']['img_rows'], cfg['data']['img_cols']), ) testloader = data.DataLoader(t_loader, batch_size=1, num_workers=cfg['training']['n_workers']) # Setup Model model = get_model(cfg['model'], n_classes) state = convert_state_dict(torch.load(args.model_path)["model_state"]) # state=torch.load(args.model_path)["model_state"] model.load_state_dict(state) model.eval() model.to(device) with torch.no_grad(): for i_val, (img_path, image_src, image_dst) in tqdm(enumerate(testloader)): img_name = img_path[0] image_src = image_src.to(device) image_dst = image_dst.to(device) outputs = model(image_src, image_dst) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0) decoded = loader.decode_segmap(pred) out_path = "test_out/changenet_change_det/" + Path( img_name).stem + ".png" decoded_bgr = cv.cvtColor(decoded, cv.COLOR_RGB2BGR) # misc.imsave(out_path, decoded) cv.imwrite(out_path, decoded_bgr)
def validate(args): os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, split=args.split, is_transform=True, img_size=(args.img_rows, args.img_cols)) n_classes = loader.n_classes valloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4) running_metrics = runningScore(n_classes) # Setup Model model = get_model(args.arch, n_classes) checkpoint = torch.load(args.model_path) state = convert_state_dict(checkpoint['model_state']) model.load_state_dict(state) print("Loaded checkpoint '{}' (epoch {})".format(args.model_path, checkpoint['epoch'])) model.eval() for i, (images, labels) in tqdm(enumerate(valloader)): model.cuda() images = Variable(images.cuda(), volatile=True) labels = Variable(labels.cuda(), volatile=True) outputs = model(images) pred = outputs.data.max(1)[1].cpu().numpy() gt = labels.data.cpu().numpy() running_metrics.update(gt, pred) score, class_iou = running_metrics.get_scores() for k, v in score.items(): print(k, v) for i in range(n_classes): print(i, classes[i], class_iou[i]) print('\t'.join([str(class_iou[i]) for i in range(n_classes)]))
def test(args): # Setup image print("Read Input Image from : {}".format(args.img_path)) img = misc.imread(args.img_path) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset, config_file=args.config_file) loader = data_loader(data_path, is_transform=True) n_classes = loader.n_classes img = img[:, :, ::-1] img = img.astype(np.float64) img -= loader.mean img = misc.imresize(img, (loader.img_size[0], loader.img_size[1])) img = img.astype(float) / 255.0 # NHWC -> NCWH img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() # Setup Model model = torch.load(args.model_path) model.eval() #if torch.cuda.is_available(): # model.cuda(0) # images = Variable(img.cuda(0)) #else: images = Variable(img) import time start_time = time.clock() outputs = model(images) print('Time: {time}'.format(time=(time.clock() - start_time))) pred = outputs[0].cpu().data.numpy() if args.label: pred = pred[args.label] else: pred = pred.argmax(0) misc.imsave(args.out_path, pred) print("Segmentation Mask Saved at: {}".format(args.out_path))
def test(): # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, img_size=args.img_size) n_classes = loader.n_classes n_channels = loader.n_channels valloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, shuffle=True) # Setup Model model = torch.load(args.model_path) model.eval() if torch.cuda.is_available(): model.cuda(0) for i, (images, labels) in enumerate(tqdm(valloader)): if torch.cuda.is_available(): images = Variable(images.cuda(0)) labels = Variable(labels.cuda(0)) else: images = Variable(images) labels = Variable(labels) outputs = model(images) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=1) gt = labels.data.cpu().numpy() for gt_, pred_ in zip(gt, pred): gt_path = args.out_dir + "gt{}.png".format(i) pred_path = args.out_dir + "pred{}.png".format(i) decoded_gt = loader.decode_segmap(gt_) decoded_pred = loader.decode_segmap(pred_) misc.imsave(gt_path, decoded_gt) misc.imsave(pred_path, decoded_pred)
def test(args): # Setup image print("Read Input Image from : {}".format(args.img_path)) img = misc.imread(args.img_path) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True) n_classes = loader.n_classes img = img[:, :, ::-1] img = img.astype(np.float64) img -= loader.mean img = misc.imresize(img, (loader.img_size[0], loader.img_size[1])) img = img.astype(float) / 255.0 # NHWC -> NCWH img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() # Setup Model model = torch.load(args.model_path) model.eval() if torch.cuda.is_available(): model.cuda(0) images = Variable(img.cuda(0)) else: images = Variable(img) outputs = model(images) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=1) decoded = loader.decode_segmap(pred[0]) print(np.unique(pred)) misc.imsave(args.out_path, decoded) print("Segmentation Mask Saved at: {}".format(args.out_path))
def train(args): # Setup Augmentations data_aug= Compose([RandomRotate(10), RandomHorizontallyFlip()]) # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) t_loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols), augmentations=data_aug) v_loader = data_loader(data_path, is_transform=True, split='val', img_size=(args.img_rows, args.img_cols)) n_classes = t_loader.n_classes trainloader = data.DataLoader(t_loader, batch_size=args.batch_size, num_workers=8, shuffle=True) valloader = data.DataLoader(v_loader, batch_size=args.batch_size, num_workers=8) # Setup Metrics running_metrics = runningScore(n_classes) # Setup visdom for visualization if args.visdom: vis = visdom.Visdom() loss_window = vis.line(X=torch.zeros((1,)).cpu(), Y=torch.zeros((1)).cpu(), opts=dict(xlabel='minibatches', ylabel='Loss', title='Training Loss', legend=['Loss'])) # Setup Model model = get_model(args.arch, n_classes) model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count())) model.cuda() # Check if model has custom optimizer / loss if hasattr(model.module, 'optimizer'): optimizer = model.module.optimizer else: optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) if hasattr(model.module, 'loss'): print('Using custom loss') loss_fn = model.module.loss else: loss_fn = cross_entropy2d if args.resume is not None: if os.path.isfile(args.resume): print("Loading model and optimizer from checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['optimizer_state']) print("Loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) else: print("No checkpoint found at '{}'".format(args.resume)) best_iou = -100.0 for epoch in range(args.n_epoch): model.train() for i, (images, labels) in enumerate(trainloader): images = Variable(images.cuda()) labels = Variable(labels.cuda()) optimizer.zero_grad() outputs = model(images) loss = loss_fn(input=outputs, target=labels) loss.backward() optimizer.step() if args.visdom: vis.line( X=torch.ones((1, 1)).cpu() * i, Y=torch.Tensor([loss.data[0]]).unsqueeze(0).cpu(), win=loss_window, update='append') if (i+1) % 20 == 0: print("Epoch [%d/%d] Loss: %.4f" % (epoch+1, args.n_epoch, loss.data[0])) model.eval() for i_val, (images_val, labels_val) in tqdm(enumerate(valloader)): images_val = Variable(images_val.cuda(), volatile=True) labels_val = Variable(labels_val.cuda(), volatile=True) outputs = model(images_val) pred = outputs.data.max(1)[1].cpu().numpy() gt = labels_val.data.cpu().numpy() running_metrics.update(gt, pred) score, class_iou = running_metrics.get_scores() for k, v in score.items(): print(k, v) running_metrics.reset() if score['Mean IoU : \t'] >= best_iou: best_iou = score['Mean IoU : \t'] state = {'epoch': epoch+1, 'model_state': model.state_dict(), 'optimizer_state' : optimizer.state_dict(),} torch.save(state, "{}_{}_best_model.pkl".format(args.arch, args.dataset))
def test(args, cfg): # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # device=torch.device("cuda:0") # device_1=torch.device("cpu") model_file_name = os.path.split(args.model_path)[1] model_name = model_file_name[:model_file_name.find("_")] IMG_Path = Path(args.img_path) IMG_File = natsort.natsorted(list(IMG_Path.glob("*.png")), alg=natsort.PATH) IMG_Str = [] for i in IMG_File: IMG_Str.append(str(i)) # Setup image print("Read Input Image from : {}".format(args.img_path)) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset, config_file=cfg) loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm) n_classes = loader.n_classes # Setup Model model = get_model(cfg['model'], n_classes) state = convert_state_dict(torch.load(args.model_path)["model_state"]) # state=torch.load(args.model_path)["model_state"] model.load_state_dict(state) model.eval() model.to(device) for j in tqdm(range(len(IMG_Str))): img_path = IMG_Str[j] img_input = misc.imread(img_path) sp = list(img_input.shape) #shape height*width*channel sp = sp[0:2] ori_size = tuple(sp) # img = img[:, :, ::-1] # multiscale # img_125=cv.resize(img,dsize=(0,0),fx=1.25,fy=1.25,interpolation=cv.INTER_LINEAR) # img_075=cv.resize(img,dsize=(0,0),fx=0.75,fy=0.75,interpolation=cv.INTER_LINEAR) # scale_list=[2.0,1.75,1.5,1.25,1,0.75,0.5] scale_list = [1.5, 1.25, 0.75, 0.5] # scale_list=[1.4,1.2,0.8,0.6] # scale_list=[2.0] multi_avg = torch.zeros((1, 6, 512, 512), dtype=torch.float32).to(device) # torch.zeros(batch-size,num-classes,height,width) for scale in scale_list: if scale != 1: img = cv.resize(img_input, dsize=(0, 0), fx=scale, fy=scale, interpolation=cv.INTER_LINEAR) else: img = img_input img = img.astype(np.float64) # img -= loader.mean if args.img_norm: img = img.astype(float) / 255.0 # NHWC -> NCHW img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() images = img.to(device) outputs = model(images) # del images # bilinear is ok for both upsample and downsample if scale != 1: outputs = F.upsample(outputs, ori_size, mode='bilinear', align_corners=False) # outputs=outputs.to(device) multi_avg = multi_avg + outputs # del outputs # outputs=multi_avg/len(scale_list) outputs = multi_avg out_path = "test_out/mv3_1_true_2_res50_data10_MS/mv3_1_true_2_res50_data10_MS_7/" + Path( img_path).stem + "_S4_not_1.pt" torch.save(outputs, out_path)
def train(args): # Setup Augmentations data_aug = Compose([RandomRotate(10), RandomHorizontallyFlip()]) # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) t_loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols), augmentations=data_aug, img_norm=args.img_norm) v_loader = data_loader(data_path, is_transform=True, split='val', img_size=(args.img_rows, args.img_cols), img_norm=args.img_norm) n_classes = t_loader.n_classes trainloader = data.DataLoader(t_loader, batch_size=args.batch_size, num_workers=8, shuffle=True) valloader = data.DataLoader(v_loader, batch_size=args.batch_size, num_workers=8) # Setup Metrics running_metrics = runningScore(n_classes) # Setup visdom for visualization if args.visdom: vis = visdom.Visdom() loss_window = vis.line(X=torch.zeros((1, )).cpu(), Y=torch.zeros((1)).cpu(), opts=dict(xlabel='minibatches', ylabel='Loss', title='Training Loss', legend=['Loss'])) # Setup Model model = get_model(args.arch, n_classes) model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count())) model.cuda() # Check if model has custom optimizer / loss if hasattr(model.module, 'optimizer'): optimizer = model.module.optimizer else: optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) if hasattr(model.module, 'loss'): print('Using custom loss') loss_fn = model.module.loss else: loss_fn = cross_entropy2d if args.resume is not None: if os.path.isfile(args.resume): print("Loading model and optimizer from checkpoint '{}'".format( args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['optimizer_state']) print("Loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("No checkpoint found at '{}'".format(args.resume)) best_iou = -100.0 for epoch in range(args.n_epoch): model.train() for i, (images, labels) in enumerate(trainloader): images = images.cuda() labels = labels.cuda() optimizer.zero_grad() outputs = model(images) loss = loss_fn(input=outputs, target=labels) loss.backward() optimizer.step() if args.visdom: vis.line(X=torch.ones((1, )).cpu() * i, Y=torch.Tensor([loss.item()]).cpu(), win=loss_window, update='append') if (i + 1) % 20 == 0: print("Epoch [%d/%d] Loss: %.4f" % (epoch + 1, args.n_epoch, loss.item())) model.eval() with torch.no_grad(): for i_val, (images_val, labels_val) in tqdm(enumerate(valloader)): images_val = images_val.cuda() labels_val = labels_val.cuda() outputs = model(images_val) pred = outputs.detach().max(1)[1].cpu().numpy() gt = labels_val.detach().cpu().numpy() running_metrics.update(gt, pred) score, class_iou = running_metrics.get_scores() for k, v in score.items(): print(k, v) running_metrics.reset() if score['Mean IoU : \t'] >= best_iou: best_iou = score['Mean IoU : \t'] state = { 'epoch': epoch + 1, 'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict(), } torch.save(state, "{}_{}_best_model.pkl".format(args.arch, args.dataset))
def validate(cfg, args): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_file_name = os.path.split(args.model_path)[1] model_name = model_file_name[:model_file_name.find("_")] # Setup Dataloader data_loader = get_loader(cfg['data']['dataset']) data_path = get_data_path(cfg['data']['dataset']) loader = data_loader( data_path, split=cfg['data']['val_split'], is_transform=True, img_size=(cfg['data']['img_rows'], cfg['data']['img_rows']), ) n_classes = loader.n_classes valloader = data.DataLoader(loader, batch_size=cfg['training']['batch_size'], num_workers=8) running_metrics = runningScore(n_classes) # Setup Model model = get_model(model_name, n_classes, version=cfg['data']['dataset']) state = convert_state_dict(torch.load(args.model_path)["model_state"]) model.load_state_dict(state) model.eval() model.to(device) for i, (images, labels) in enumerate(valloader): start_time = timeit.default_timer() images = images.to(device) if args.eval_flip: outputs = model(images) # Flip images in numpy (not support in tensor) outputs = outputs.data.cpu().numpy() flipped_images = np.copy(images.data.cpu().numpy()[:, :, :, ::-1]) flipped_images = torch.from_numpy(flipped_images).float().to( device) outputs_flipped = model(flipped_images) outputs_flipped = outputs_flipped.data.cpu().numpy() outputs = (outputs + outputs_flipped[:, :, :, ::-1]) / 2.0 pred = np.argmax(outputs, axis=1) else: outputs = model(images) pred = outputs.data.max(1)[1].cpu().numpy() gt = labels.numpy() if args.measure_time: elapsed_time = timeit.default_timer() - start_time print("Inference time \ (iter {0:5d}): {1:3.5f} fps".format( i + 1, pred.shape[0] / elapsed_time)) running_metrics.update(gt, pred) score, class_iou = running_metrics.get_scores() for k, v in score.items(): print(k, v) for i in range(n_classes): print(i, class_iou[i])
def test(args): model_file_name = os.path.split(args.model_path)[1] model_name = model_file_name[:model_file_name.find('_')] # Setup image print("Read Input Image from : {}".format(args.img_path)) img = misc.imread(args.img_path) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm) n_classes = loader.n_classes resized_img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]), interp='bicubic') orig_size = img.shape[:-1] if model_name in ['pspnet', 'icnet', 'icnetBN']: img = misc.imresize(img, (orig_size[0]//2*2+1, orig_size[1]//2*2+1)) # uint8 with RGB mode, resize width and height which are odd numbers else: img = misc.imresize(img, (loader.img_size[0], loader.img_size[1])) img = img[:, :, ::-1] img = img.astype(np.float64) img -= loader.mean if args.img_norm: img = img.astype(float) / 255.0 # NHWC -> NCHW img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() # Setup Model model = get_model(model_name, n_classes, version=args.dataset) state = convert_state_dict(torch.load(args.model_path)['model_state']) model.load_state_dict(state) model.eval() if torch.cuda.is_available(): model.cuda(0) images = Variable(img.cuda(0), volatile=True) else: images = Variable(img, volatile=True) outputs = model(images) #outputs = F.softmax(outputs, dim=1) if args.dcrf: unary = outputs.data.cpu().numpy() unary = np.squeeze(unary, 0) unary = -np.log(unary) unary = unary.transpose(2, 1, 0) w, h, c = unary.shape unary = unary.transpose(2, 0, 1).reshape(loader.n_classes, -1) unary = np.ascontiguousarray(unary) resized_img = np.ascontiguousarray(resized_img) d = dcrf.DenseCRF2D(w, h, loader.n_classes) d.setUnaryEnergy(unary) d.addPairwiseBilateral(sxy=5, srgb=3, rgbim=resized_img, compat=1) q = d.inference(50) mask = np.argmax(q, axis=0).reshape(w, h).transpose(1, 0) decoded_crf = loader.decode_segmap(np.array(mask, dtype=np.uint8)) dcrf_path = args.out_path[:-4] + '_drf.png' misc.imsave(dcrf_path, decoded_crf) print("Dense CRF Processed Mask Saved at: {}".format(dcrf_path)) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0) if model_name in ['pspnet', 'icnet', 'icnetBN']: pred = pred.astype(np.float32) pred = misc.imresize(pred, orig_size, 'nearest', mode='F') # float32 with F mode, resize back to orig_size decoded = loader.decode_segmap(pred) print('Classes found: ', np.unique(pred)) misc.imsave(args.out_path, decoded) print("Segmentation Mask Saved at: {}".format(args.out_path))
def encode_segmap(self, mask): #Put all void classes to zero for _voidc in self.void_classes: mask[mask==_voidc] = 0 for _validc in self.valid_classes: mask[mask==_validc] = self.class_map[_validc] return mask # Run it as a standalone file from project root directory if __name__ == '__main__': import torchvision import matplotlib.pyplot as plt from ptsemseg.loader import get_loader, get_data_path dst = cityscapesLoader(get_data_path('cityscapes'), is_transform=True) trainloader = data.DataLoader(dst, batch_size=4, num_workers=0) for i, data in enumerate(trainloader): imgs, labels = data img = imgs.numpy()[0, ::-1, :, :] img = np.transpose(img, [1,2,0]) f, axarr = plt.subplots(2,1) axarr[0].imshow(img) axarr[1].imshow(dst.decode_segmap(labels.numpy()[0])) plt.show() a = input() if a == 'ex': break else: plt.close()
def test(args): # Setup image print("Read Input Image from : {}".format(args.img_path)) img = misc.imread(args.img_path) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, is_transform=True) n_classes = loader.n_classes resized_img = misc.imresize(img, (loader.img_size[0], loader.img_size[1]), interp='bicubic') img = img[:, :, ::-1] img = img.astype(np.float64) img -= loader.mean img = misc.imresize(img, (loader.img_size[0], loader.img_size[1])) img = img.astype(float) / 255.0 # NHWC -> NCWH img = img.transpose(2, 0, 1) img = np.expand_dims(img, 0) img = torch.from_numpy(img).float() # Setup Model model = get_model(args.model_path[:args.model_path.find('_')], n_classes) state = convert_state_dict(torch.load(args.model_path)['model_state']) model.load_state_dict(state) model.eval() model.cuda(0) images = Variable(img.cuda(0), volatile=True) outputs = F.softmax(model(images), dim=1) if args.dcrf == "True": unary = outputs.data.cpu().numpy() unary = np.squeeze(unary, 0) unary = -np.log(unary) unary = unary.transpose(2, 1, 0) w, h, c = unary.shape unary = unary.transpose(2, 0, 1).reshape(loader.n_classes, -1) unary = np.ascontiguousarray(unary) resized_img = np.ascontiguousarray(resized_img) d = dcrf.DenseCRF2D(w, h, loader.n_classes) d.setUnaryEnergy(unary) d.addPairwiseBilateral(sxy=5, srgb=3, rgbim=resized_img, compat=1) q = d.inference(50) mask = np.argmax(q, axis=0).reshape(w, h).transpose(1, 0) decoded_crf = loader.decode_segmap(np.array(mask, dtype=np.uint8)) dcrf_path = args.out_path[:-4] + '_drf.png' misc.imsave(dcrf_path, decoded_crf) print("Dense CRF Processed Mask Saved at: {}".format(dcrf_path)) if torch.cuda.is_available(): model.cuda(0) images = Variable(img.cuda(0), volatile=True) else: images = Variable(img, volatile=True) pred = np.squeeze(outputs.data.max(1)[1].cpu().numpy(), axis=0) decoded = loader.decode_segmap(pred) print('Classes found: ', np.unique(pred)) misc.imsave(args.out_path, decoded) print("Segmentation Mask Saved at: {}".format(args.out_path))
def validate_bayesian(args, model, split, labeled_index=None, verbose=False): # Setup Data data_loader = get_loader('camvid') data_path = get_data_path('camvid') dataset = data_loader(data_path, split, is_transform=True, labeled_index=labeled_index) valloader = data.DataLoader(dataset, batch_size=1, num_workers=0, shuffle=False) # Setup Model model.eval() model.apply(set_dropout) # Uncertainty Hyperparameter T = args.sample_num inference_time = 0 gts, preds, uncts = [], [], [] uncts_r, uncts_e, uncts_v, uncts_b = [], [], [], [] image_names_all = [] for i, (images, labels, image_names) in enumerate(valloader): torch.cuda.synchronize() t1 = time.time() if torch.cuda.is_available(): model.cuda() images = Variable(images.cuda(), volatile=True) labels_var = Variable(labels.cuda(async=True), volatile=True) else: images = Variable(images, volatile=True) labels_var = Variable(labels, volatile=True) output_list = [] # MC dropout for t in range(T): output = F.softmax(model(images)) if t == 0: output_mean = output * 0 output_square = output * 0 entropy_mean = output.mean(1) * 0 output_mean += output output_square += output.pow(2) entropy_mean += acquisition_func('e', output) output_mean = output_mean / T output_square = output_square / T entropy_mean = entropy_mean / T # Uncertainty estimation if args.acqu_func != 'all': unc_map = acquisition_func(args.acqu_func, output_mean,\ square_mean=output_square, entropy_mean=entropy_mean) else: unc_map_r = acquisition_func('r', output_mean,\ square_mean=output_square, entropy_mean=entropy_mean) unc_map_e = acquisition_func('e', output_mean,\ square_mean=output_square, entropy_mean=entropy_mean) unc_map_b = acquisition_func('b', output_mean,\ square_mean=output_square, entropy_mean=entropy_mean) unc_map_v = acquisition_func('v', output_mean,\ square_mean=output_square, entropy_mean=entropy_mean) pred = torch.max(output_mean, 1)[1] torch.cuda.synchronize() t2 = time.time() gts += list(labels.numpy()) preds += list(pred.data.cpu().numpy()) if args.acqu_func != 'all': uncts += list(unc_map.data.cpu().numpy()) else: uncts_r += list(unc_map_r.data.cpu().numpy()) uncts_e += list(unc_map_e.data.cpu().numpy()) uncts_b += list(unc_map_b.data.cpu().numpy()) uncts_v += list(unc_map_v.data.cpu().numpy()) image_names_all += list(image_names) inference_time += t2 - t1 if verbose: print '[Info] evaluate ', image_names print '[Time] Average Inference Time = ', inference_time / (i + 1) # Save unct_map and pred_map if args.save_output: for index in range(len(preds)): out_name = os.path.basename(image_names_all[index]).replace( '.png', '') np.save(os.path.join(args.out_pred_dir, out_name), preds[index]) if args.acqu_func != 'all': np.save(os.path.join(args.out_unct_dir, out_name), uncts[index]) else: np.save(os.path.join(args.out_unct_dir_r, out_name), uncts_r[index]) np.save(os.path.join(args.out_unct_dir_e, out_name), uncts_e[index]) np.save(os.path.join(args.out_unct_dir_b, out_name), uncts_b[index]) np.save(os.path.join(args.out_unct_dir_v, out_name), uncts_v[index]) return gts, preds, uncts
def merge(args): result_root_path = make_result_dir(args.dataset, args.split) # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) loader = data_loader(data_path, split=args.split, is_transform=True, img_size=(args.img_rows, args.img_cols), img_norm=args.img_norm, no_gt=args.no_gt) n_classes = loader.n_classes testloader = data.DataLoader(loader, batch_size=args.batch_size, num_workers=4, pin_memory=True) running_metrics = runningScore(n_classes) rm = 0 pred_dict = {} map = AverageMeter() with open('list_test_18000') as f: id_list = f.read().splitlines() # Average all the probability maps from all folds all_prob = np.zeros((len(testloader), args.img_rows, args.img_cols), dtype=np.float32) for i in range(1, args.max_k_split + 1): prob = np.load('prob-{}_{}_{}.npy'.format(args.split, i, args.max_k_split)) all_prob = all_prob + prob all_prob = all_prob / args.max_k_split np.save('prob-{}_avg_{}.npy'.format(args.split, args.max_k_split), all_prob) for i, id in tqdm(enumerate(id_list)): lbl = id + '.png' pred = all_prob[i, :, :] if not args.no_gt: gt_path = os.path.join(data_path, args.split, 'masks', lbl) gt = loader.encode_segmap(cv2.imread(gt_path, cv2.IMREAD_GRAYSCALE)) pred = np.where(pred < args.pred_thr, 0, 1) if pred.sum( ) <= loader.lbl_thr: # Remove salt masks for sum of salts <= a threshold lbl_thr pred = np.zeros((args.img_rows, args.img_cols), dtype=np.uint8) rm = rm + 1 decoded = loader.decode_segmap(pred) rle_mask = loader.RLenc(decoded) pred_dict[id] = rle_mask save_result_path = os.path.join(result_root_path, id + '.png') cv2.imwrite(save_result_path, decoded) if not args.no_gt: map_val = running_metrics.comput_map(gt, pred) map.update(map_val.mean(), n=map_val.size) if not args.no_gt: print('Mean Average Precision: {:.5f}'.format(map.avg)) # Create final submission sub = pd.DataFrame.from_dict(pred_dict, orient='index') sub.index.names = ['id'] sub.columns = ['rle_mask'] sub.to_csv(args.split + '.csv') print('To black: ', rm)
def train(args): logger.auto_set_dir() os.environ['CUDA_VISIBLE_DEVICES'] = '3' # Setup Augmentations data_aug = Compose([RandomRotate(10), RandomHorizontallyFlip()]) # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) t_loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols), epoch_scale=4, augmentations=data_aug, img_norm=args.img_norm) v_loader = data_loader(data_path, is_transform=True, split='val', img_size=(args.img_rows, args.img_cols), img_norm=args.img_norm) n_classes = t_loader.n_classes trainloader = data.DataLoader(t_loader, batch_size=args.batch_size, num_workers=8, shuffle=True) valloader = data.DataLoader(v_loader, batch_size=args.batch_size, num_workers=8) # Setup Metrics running_metrics = runningScore(n_classes) # Setup Model from model_zoo.deeplabv1 import VGG16_LargeFoV model = VGG16_LargeFoV(class_num=n_classes, image_size=[args.img_cols, args.img_rows], pretrained=True) #model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count())) model.cuda() # Check if model has custom optimizer / loss if hasattr(model, 'optimizer'): logger.warn("don't have customzed optimizer, use default setting!") optimizer = model.module.optimizer else: optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) optimizer_summary(optimizer) if args.resume is not None: if os.path.isfile(args.resume): logger.info( "Loading model and optimizer from checkpoint '{}'".format( args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['optimizer_state']) logger.info("Loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: logger.info("No checkpoint found at '{}'".format(args.resume)) best_iou = -100.0 for epoch in tqdm(range(args.n_epoch), total=args.n_epoch): model.train() for i, (images, labels) in tqdm(enumerate(trainloader), total=len(trainloader), desc="training epoch {}/{}".format( epoch, args.n_epoch)): cur_iter = i + epoch * len(trainloader) cur_lr = adjust_learning_rate(optimizer, args.l_rate, cur_iter, args.n_epoch * len(trainloader), power=0.9) #if i > 10:break images = Variable(images.cuda()) labels = Variable(labels.cuda()) optimizer.zero_grad() outputs = model(images) #print(np.unique(outputs.data[0].cpu().numpy())) loss = CrossEntropyLoss2d_Seg(input=outputs, target=labels, class_num=n_classes) loss.backward() optimizer.step() if (i + 1) % 100 == 0: logger.info("Epoch [%d/%d] Loss: %.4f, lr: %.7f" % (epoch + 1, args.n_epoch, loss.data[0], cur_lr)) model.eval() for i_val, (images_val, labels_val) in tqdm(enumerate(valloader), total=len(valloader), desc="validation"): images_val = Variable(images_val.cuda(), volatile=True) labels_val = Variable(labels_val.cuda(), volatile=True) outputs = model(images_val) pred = outputs.data.max(1)[1].cpu().numpy() gt = labels_val.data.cpu().numpy() running_metrics.update(gt, pred) score, class_iou = running_metrics.get_scores() for k, v in score.items(): logger.info("{}: {}".format(k, v)) running_metrics.reset() if score['Mean IoU : \t'] >= best_iou: best_iou = score['Mean IoU : \t'] state = { 'epoch': epoch + 1, 'mIoU': best_iou, 'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict(), } torch.save(state, os.path.join(logger.get_logger_dir(), "best_model.pkl"))
args.deep_features_size = 256 elif args.backend == 'resnet50' or args.backend == 'resnet101' or args.backend == 'resnet152': args.psp_size = 2048 args.deep_features_size = 1024 # Setup Dataloader val_transforms = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(256), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) val_dataset = data_loader(data_path, split='val', transform=val_transforms) args.n_classes = val_dataset.n_classes valloader = data.DataLoader(val_dataset, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=False, pin_memory=True) # Setup Model model = get_model(args) model.load_state_dict(torch.load(args.model_path)['state_dict']) model = torch.nn.DataParallel(model, device_ids=range( torch.cuda.device_count())).cuda()