def save_predict(output, gt, img_name, dataset, save_path, output_grey=False, output_color=True, gt_color=False): if output_grey: output_grey = Image.fromarray(output) output_grey.save(os.path.join(save_path, img_name + '.png')) if output_color: if dataset == 'cityscapes': output_color = cityscapes_colorize_mask(output) elif dataset == 'camvid': output_color = camvid_colorize_mask(output) output_color.save(os.path.join(save_path, img_name + '_color.png')) if gt_color: if dataset == 'cityscapes': gt_color = cityscapes_colorize_mask(gt) elif dataset == 'camvid': gt_color = camvid_colorize_mask(gt) gt_color.save(os.path.join(save_path, img_name + '_gt.png'))
def save_predict(output, gt, img_name, dataset, save_path, output_grey=False, output_color=True, gt_color=False): if output_grey: row, col = output.shape dst = np.ones((row, col), dtype=np.uint8) * 255 for i in range(19): dst[output == i] = color_list[i] output_grey = Image.fromarray(dst) output_grey.save(os.path.join(save_path, img_name[:-1] + '.png')) if output_color: if dataset == 'cityscapes': output_color = cityscapes_colorize_mask(output) elif dataset == 'camvid': output_color = camvid_colorize_mask(output) output_color.save(os.path.join(save_path, img_name[:-1] + '_color.png')) if gt_color: if dataset == 'cityscapes': gt_color = cityscapes_colorize_mask(gt) elif dataset == 'camvid': gt_color = camvid_colorize_mask(gt) gt_color.save(os.path.join(save_path, img_name + '_gt.png'))
def save_predict(output, gt, img_name, dataset, save_path, output_grey=False, output_color=True, gt_color=False): if output_grey: if dataset == 'cityscapes': output[np.where(output == 18)] = 33 output[np.where(output == 17)] = 32 output[np.where(output == 16)] = 31 output[np.where(output == 15)] = 28 output[np.where(output == 14)] = 27 output[np.where(output == 13)] = 26 output[np.where(output == 12)] = 25 output[np.where(output == 11)] = 24 output[np.where(output == 10)] = 23 output[np.where(output == 9)] = 22 output[np.where(output == 8)] = 21 output[np.where(output == 7)] = 20 output[np.where(output == 6)] = 19 output[np.where(output == 5)] = 17 output[np.where(output == 4)] = 13 output[np.where(output == 3)] = 12 output[np.where(output == 2)] = 11 output[np.where(output == 1)] = 8 output[np.where(output == 0)] = 7 output_grey = Image.fromarray(output) output_grey.save(os.path.join(save_path, img_name + '.png')) if output_color: if dataset == 'cityscapes': output_color = cityscapes_colorize_mask(output) elif dataset == 'paris': output_color = paris_colorize_mask(output) elif dataset == 'road': output_color = road_colorize_mask(output) elif dataset == 'austin': output_color = austin_colorize_mask(output) elif dataset == 'postdam' or dataset == 'vaihingen': output_color = isprs_colorize_mask(output) output_color.save(os.path.join(save_path, img_name + '_color.png')) if gt_color: if dataset == 'cityscapes': gt_color = cityscapes_colorize_mask(gt) elif dataset == 'paris': gt_color = paris_colorize_mask(gt) elif dataset == 'road': gt_color = road_colorize_mask(gt) elif dataset == 'austin': gt_color = austin_colorize_mask(gt) elif dataset == 'postdam' or dataset == 'vaihingen': gt_color = isprs_colorize_mask(gt) gt_color.save(os.path.join(save_path, img_name + '_gt.png'))
def val(args, val_loader, model, criterion): """ args: val_loader: loaded for validation dataset model: model criterion: loss function return: IoU class, and mean IoU """ #evaluation mode model.eval() total_batches = len(val_loader) data_list = [] for i, (input, label, size, name) in enumerate(val_loader): start_time = time.time() input_var = Variable(input, volatile=True).cuda() output = model(input_var) time_taken = time.time() - start_time print('[%d/%d] time: %.2f' % (i, total_batches, time_taken)) # save seg image output = output.cpu().data[0].numpy() # 1xCxHxW ---> CxHxW gt = np.asarray(label[0].numpy(), dtype=np.uint8) output = output.transpose(1, 2, 0) # CxHxW --> HxWxC output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) data_list.append([gt.flatten(), output.flatten()]) output_color = cityscapes_colorize_mask(output) output = Image.fromarray(output) output.save('%s/%s.png' % (args.save_seg_dir, name[0])) output_color.save('%s/%s_color.png' % (args.save_seg_dir, name[0])) meanIoU, IoUs = get_iou(data_list, args.classes) print("mean IoU:", meanIoU) print(IoUs)
def test(args, test_loader, model): """ args: test_loader: loaded for test set model: model criterion: loss function return: IoU class, and mean IoU """ #evaluation mode model.eval() total_batches = len(test_loader) for i, (input, size, name) in enumerate(test_loader): start_time = time.time() input_var = Variable(input, volatile=True).cuda() output = model(input_var) time_taken = time.time() - start_time print('[%d/%d] time: %.2f' % (i, total_batches, time_taken)) # save seg image output = output.cpu().data[0].numpy() # 1xCxHxW ---> CxHxW output = output.transpose(1, 2, 0) # CxHxW --> HxWxC output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) output_color = cityscapes_colorize_mask(output) output = Image.fromarray(output) #output.save( "%s/%s.png " % (args.save_seg_dir, name[0]) ) output_color.save("%s/%s_color.png" % (args.save_seg_dir, name[0]))
def test(args, test_loader, model): ''' Args: test_loader: loaded for test dataset model: model ''' #switch to evaluation mode model.eval() total_batches = len(test_loader) data_list = [] for i, (input, size, name) in enumerate(test_loader): start_time = time.time() #input_var = Variable(input, volatile=True).cuda() torch version==0.3 with torch.no_grad(): input_var = Variable(input).cuda() # run the mdoel output = model(input_var) time_taken = time.time() - start_time print('[%d/%d] time: %.2f' % (i, total_batches, time_taken)) # save seg image output = output.cpu().data[0].numpy() # 1xCxHxW ---> CxHxW output = output.transpose(1, 2, 0) # CxHxW --> HxWxC output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) output_color = cityscapes_colorize_mask(output) output = Image.fromarray(output) output_color.save('%s/%s.png' % (args.save_seg_dir, name[0]))
def main(): args = get_arguments() print("=====> Configure dataset and model") configure_dataset_model(args) print(args) print("=====> Set GPU for training") if args.cuda: print("====> Use gpu id: '{}'".format(args.gpus)) os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus if not torch.cuda.is_available(): raise Exception( "No GPU found or Wrong gpu id, please run without --cuda") model = CoattentionNet(num_classes=args.num_classes) saved_state_dict = torch.load(args.restore_from, map_location=lambda storage, loc: storage) #print(saved_state_dict.keys()) #model.load_state_dict({k.replace('pspmodule.',''):v for k,v in torch.load(args.restore_from)['state_dict'].items()}) model.load_state_dict(convert_state_dict(saved_state_dict["model"]) ) #convert_state_dict(saved_state_dict["model"]) model.eval() model.cuda() if args.dataset == 'voc12': testloader = data.DataLoader(VOCDataTestSet(args.data_dir, args.data_list, crop_size=(505, 505), mean=args.img_mean), batch_size=1, shuffle=False, pin_memory=True) interp = nn.Upsample(size=(505, 505), mode='bilinear') voc_colorize = VOCColorize() elif args.dataset == 'cityscapes': testloader = data.DataLoader( CityscapesTestDataSet(args.data_dir, args.data_list, f_scale=args.f_scale, mean=args.img_mean), batch_size=1, shuffle=False, pin_memory=True ) # f_sale, meaning resize image at f_scale as input interp = nn.Upsample(size=(1024, 2048), mode='bilinear') #size = (h,w) voc_colorize = VOCColorize() elif args.dataset == 'davis': #for davis 2016 db_test = db.PairwiseImg( train=False, inputRes=(473, 473), db_root_dir=args.data_dir, transform=None, seq_name=None, sample_range=args.sample_range ) #db_root_dir() --> '/path/to/DAVIS-2016' train path testloader = data.DataLoader(db_test, batch_size=1, shuffle=False, num_workers=0) voc_colorize = VOCColorize() else: print("dataset error") data_list = [] if args.save_segimage: if not os.path.exists(args.seg_save_dir) and not os.path.exists( args.vis_save_dir): os.makedirs(args.seg_save_dir) os.makedirs(args.vis_save_dir) print("======> test set size:", len(testloader)) my_index = 0 old_temp = '' for index, batch in enumerate(testloader): print('%d processd' % (index)) target = batch['target'] #search = batch['search'] temp = batch['seq_name'] args.seq_name = temp[0] print(args.seq_name) if old_temp == args.seq_name: my_index = my_index + 1 else: my_index = 0 output_sum = 0 for i in range(0, args.sample_range): search = batch['search' + '_' + str(i)] search_im = search #print(search_im.size()) output = model( Variable(target, volatile=True).cuda(), Variable(search_im, volatile=True).cuda()) #print(output[0]) # output有两个 output_sum = output_sum + output[0].data[ 0, 0].cpu().numpy() #分割那个分支的结果 #np.save('infer'+str(i)+'.npy',output1) #output2 = output[1].data[0, 0].cpu().numpy() #interp' output1 = output_sum / args.sample_range first_image = np.array( Image.open(args.data_dir + '/JPEGImages/480p/blackswan/00000.jpg')) original_shape = first_image.shape output1 = cv2.resize(output1, (original_shape[1], original_shape[0])) if 0: original_image = target[0] #print('image type:',type(original_image.numpy())) original_image = original_image.numpy() original_image = original_image.transpose((2, 1, 0)) original_image = cv2.resize(original_image, (original_shape[1], original_shape[0])) unary = np.zeros((2, original_shape[0] * original_shape[1]), dtype='float32') #unary[0, :, :] = res_saliency/255 #unary[1, :, :] = 1-res_saliency/255 EPSILON = 1e-8 tau = 1.05 crf = dcrf.DenseCRF(original_shape[1] * original_shape[0], 2) anno_norm = (output1 - np.min(output1)) / ( np.max(output1) - np.min(output1)) #res_saliency/ 255. n_energy = 1.0 - anno_norm + EPSILON #-np.log((1.0 - anno_norm + EPSILON)) #/ (tau * sigmoid(1 - anno_norm)) p_energy = anno_norm + EPSILON #-np.log(anno_norm + EPSILON) #/ (tau * sigmoid(anno_norm)) #unary = unary.reshape((2, -1)) #print(unary.shape) unary[1, :] = p_energy.flatten() unary[0, :] = n_energy.flatten() crf.setUnaryEnergy(unary_from_softmax(unary)) feats = create_pairwise_gaussian(sdims=(3, 3), shape=original_shape[:2]) crf.addPairwiseEnergy(feats, compat=3, kernel=dcrf.DIAG_KERNEL, normalization=dcrf.NORMALIZE_SYMMETRIC) feats = create_pairwise_bilateral( sdims=(10, 10), schan=(1, 1, 1), # orgin is 60, 60 5, 5, 5 img=original_image, chdim=2) crf.addPairwiseEnergy(feats, compat=5, kernel=dcrf.DIAG_KERNEL, normalization=dcrf.NORMALIZE_SYMMETRIC) Q = crf.inference(5) MAP = np.argmax(Q, axis=0) output1 = MAP.reshape((original_shape[0], original_shape[1])) mask = (output1 * 255).astype(np.uint8) #print(mask.shape[0]) mask = Image.fromarray(mask) if args.dataset == 'voc12': print(output.shape) print(size) output = output[:, :size[0], :size[1]] output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) if args.save_segimage: seg_filename = os.path.join(args.seg_save_dir, '{}.png'.format(name[0])) color_file = Image.fromarray( voc_colorize(output).transpose(1, 2, 0), 'RGB') color_file.save(seg_filename) elif args.dataset == 'davis': save_dir_res = os.path.join(args.seg_save_dir, 'Results', args.seq_name) old_temp = args.seq_name if not os.path.exists(save_dir_res): os.makedirs(save_dir_res) if args.save_segimage: my_index1 = str(my_index).zfill(5) seg_filename = os.path.join(save_dir_res, '{}.png'.format(my_index1)) #color_file = Image.fromarray(voc_colorize(output).transpose(1, 2, 0), 'RGB') mask.save(seg_filename) #np.concatenate((torch.zeros(1, 473, 473), mask, torch.zeros(1, 512, 512)),axis = 0) #save_image(output1 * 0.8 + target.data, args.vis_save_dir, normalize=True) elif args.dataset == 'cityscapes': output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) if args.save_segimage: output_color = cityscapes_colorize_mask(output) output = Image.fromarray(output) output.save('%s/%s.png' % (args.seg_save_dir, name[0])) output_color.save('%s/%s_color.png' % (args.seg_save_dir, name[0])) else: print("dataset error")
def main(): args = get_arguments() print("=====> Configure dataset and model") configure_dataset_model(args) print(args) print("=====> Set GPU for training") if args.cuda: print("====> Use gpu id: '{}'".format(args.gpus)) os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus if not torch.cuda.is_available(): raise Exception( "No GPU found or Wrong gpu id, please run without --cuda") model = PSPNet(num_classes=args.num_classes) saved_state_dict = torch.load(args.restore_from) model.load_state_dict(convert_state_dict(saved_state_dict["model"])) model.eval() model.cuda() if args.dataset == 'voc12': testloader = data.DataLoader(VOCDataTestSet(args.data_dir, args.data_list, crop_size=(505, 505), mean=args.img_mean), batch_size=1, shuffle=False, pin_memory=True) interp = nn.Upsample(size=(505, 505), mode='bilinear') voc_colorize = VOCColorize() elif args.dataset == 'cityscapes': testloader = data.DataLoader( CityscapesTestDataSet(args.data_dir, args.data_list, f_scale=args.f_scale, mean=args.img_mean), batch_size=1, shuffle=False, pin_memory=True ) # f_sale, meaning resize image at f_scale as input interp = nn.Upsample(size=(1024, 2048), mode='bilinear') #size = (h,w) else: print("dataset error") data_list = [] if args.save_segimage: if not os.path.exists(args.seg_save_dir): os.makedirs(args.seg_save_dir) print("======> test set size:", len(testloader)) for index, batch in enumerate(testloader): print('%d processd' % (index)) image, size, name = batch size = size[0].numpy() output = model(Variable(image, volatile=True).cuda()) output = interp(output).cpu().data[0].numpy() if args.dataset == 'voc12': print(output.shape) print(size) output = output[:, :size[0], :size[1]] output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) if args.save_segimage: seg_filename = os.path.join(args.seg_save_dir, '{}.png'.format(name[0])) color_file = Image.fromarray( voc_colorize(output).transpose(1, 2, 0), 'RGB') color_file.save(seg_filename) elif args.dataset == 'cityscapes': output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) if args.save_segimage: output_color = cityscapes_colorize_mask(output) output = Image.fromarray(output) output.save('%s/%s.png' % (args.seg_save_dir, name[0])) output_color.save('%s/%s_color.png' % (args.seg_save_dir, name[0])) else: print("dataset error")