def eval_model(net, save_dir, batch_size=10): # Setting parameters relax_crop = 50 # Enlarge the bounding box by relax_crop pixels zero_pad_crop = True # Insert zero padding when cropping the image net.eval() composed_transforms_ts = transforms.Compose([ tr.CropFromMask(crop_elems=('image', 'gt'), relax=relax_crop, zero_pad=zero_pad_crop), tr.FixedResize(resolutions={ 'gt': None, 'crop_image': (512, 512), 'crop_gt': (512, 512) }), tr.ExtremePoints(sigma=10, pert=0, elem='crop_gt'), tr.ToImage(norm_elem='extreme_points'), tr.ConcatInputs(elems=('crop_image', 'extreme_points')), tr.ToTensor() ]) db_test = pascal.VOCSegmentation(split='val', transform=composed_transforms_ts, retname=True) testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=2) save_dir.mkdir(exist_ok=True) with torch.no_grad(): test(net, testloader, save_dir)
def create_transforms(relax_crop, zero_crop): # Preparation of the data loaders first = [ tr.CropFromMask(crop_elems=('image', 'gt'), relax=relax_crop, zero_pad=zero_crop), tr.FixedResize(resolutions={ 'crop_image': (512, 512), 'crop_gt': (512, 512) }) ] second = [ tr.ToImage(norm_elem='extreme_points'), tr.ConcatInputs(elems=('crop_image', 'extreme_points')), tr.ToTensor() ] train_tf = transforms.Compose([ tr.RandomHorizontalFlip(), tr.ScaleNRotate(rots=(-20, 20), scales=(.75, 1.25)), *first, tr.ExtremePoints(sigma=10, pert=5, elem='crop_gt'), *second ]) test_tf = transforms.Compose( [*first, tr.ExtremePoints(sigma=10, pert=0, elem='crop_gt'), *second]) return train_tf, test_tf
log_dir = os.path.join(save_dir, 'models', datetime.now().strftime('%b%d_%H-%M-%S') + '_' + socket.gethostname()) # writer = SummaryWriter(log_dir=log_dir) # Use the following optimizer optimizer = optim.SGD(train_params, lr=p['lr'], momentum=p['momentum'], weight_decay=p['wd']) p['optimizer'] = str(optimizer) # Preparation of the data loaders composed_transforms_tr = transforms.Compose([ tr.RandomHorizontalFlip(), tr.ScaleNRotate(rots=(-20, 20), scales=(.75, 1.25)), tr.CropFromMask(crop_elems=('image', 'gt'), relax=relax_crop, zero_pad=zero_pad_crop), tr.FixedResize(resolutions={'crop_image': (512, 512), 'crop_gt': (512, 512)}), tr.ExtremePoints(sigma=10, pert=5, elem='crop_gt'), tr.ToImage(norm_elem='extreme_points'), tr.ConcatInputs(elems=('crop_image', 'extreme_points')), tr.ToTensor()]) composed_transforms_ts = transforms.Compose([ tr.CropFromMask(crop_elems=('image', 'gt'), relax=relax_crop, zero_pad=zero_pad_crop), tr.FixedResize(resolutions={'crop_image': (512, 512), 'crop_gt': (512, 512)}), tr.ExtremePoints(sigma=10, pert=0, elem='crop_gt'), tr.ToImage(norm_elem='extreme_points'), tr.ConcatInputs(elems=('crop_image', 'extreme_points')), tr.ToTensor()]) voc_train = pascal.VOCSegmentation(split='train', transform=composed_transforms_tr) voc_val = pascal.VOCSegmentation(split='val', transform=composed_transforms_ts) if use_sbd: sbd = sbd.SBDSegmentation(split=['train', 'val'], transform=composed_transforms_tr, retname=True) db_train = combine_dbs([voc_train, sbd], excluded=[voc_val])
# load pretrain_dict pretrain_dict = torch.load(os.path.join(save_dir, 'models', modelName + '_epoch-' + str(resume_epoch - 1) + '.pth')) print("Initializing weights from: {}".format( os.path.join(save_dir, 'models', modelName + '_epoch-' + str(resume_epoch - 1) + '.pth'))) net.load_state_dict(pretrain_dict) net.to(device) # Generate result of the validation images net.eval() composed_transforms_ts = transforms.Compose([ tr.CropFromMask(crop_elems=('image', 'gt','void_pixels'), relax=30, zero_pad=True), tr.FixedResize(resolutions={'gt': None, 'crop_image': (512, 512), 'crop_gt': (512, 512), 'crop_void_pixels': (512, 512)},flagvals={'gt':cv2.INTER_LINEAR,'crop_image':cv2.INTER_LINEAR,'crop_gt':cv2.INTER_LINEAR,'crop_void_pixels': cv2.INTER_LINEAR}), tr.IOGPoints(sigma=10, elem='crop_gt',pad_pixel=10), tr.ToImage(norm_elem='IOG_points'), tr.ConcatInputs(elems=('crop_image', 'IOG_points')), tr.ToTensor()]) db_test = pascal.VOCSegmentation(split='val', transform=composed_transforms_ts, retname=True) testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=1) save_dir_res = os.path.join(save_dir, 'Results') if not os.path.exists(save_dir_res): os.makedirs(save_dir_res) save_dir_res_list=[save_dir_res] print('Testing Network') with torch.no_grad(): for ii, sample_batched in enumerate(testloader): inputs, gts, metas = sample_batched['concat'], sample_batched['gt'], sample_batched['meta'] inputs = inputs.to(device) coarse_outs1,coarse_outs2,coarse_outs3,coarse_outs4,fine_out = net.forward(inputs) outputs = fine_out.to(torch.device('cpu'))
if __name__ == '__main__': from dataloaders import custom_transforms as tr from torch.utils.data import DataLoader from torchvision import transforms import matplotlib.pyplot as plt composed_transforms_tr = transforms.Compose([ tr.RandomHorizontalFlip(), tr.ScaleNRotate(rots=(-15, 15), scales=(.75, 1.25)), tr.FixedResize(resolutions={ 'image': (450, 450), 'gt': (450, 450) }), tr.DistanceMap(v=0.15, elem='gt'), tr.ConcatInputs(elems=('image', 'distance_map')), tr.ToTensor() ]) voc_train = VOCSegmentation(split='train', retname=False, transform=composed_transforms_tr) dataloader = DataLoader(voc_train, batch_size=2, shuffle=True, num_workers=2) for ii, sample in enumerate(dataloader): for jj in range(sample["image"].size()[0]): dismap = sample['distance_map'][jj].numpy()
tr.FixedResizePoints( resolutions={'extreme_points': (cfg['lr_size'], cfg['lr_size'])}, mask_elem='gt', prefix='lr_'), tr.FixedResize(resolutions={ 'image': (cfg['lr_size'], cfg['lr_size']), 'gt': (cfg['lr_size'], cfg['lr_size']), 'void_pixels': (cfg['lr_size'], cfg['lr_size']) }, prefix='lr_'), tr.GaussianTransform(tr_elems=['lr_extreme_points'], mask_elem='lr_gt', sigma=10, tr_name='lr_points'), tr.ToImage(norm_elem=['points', 'image_grad', 'lr_points']), tr.ConcatInputs(cat_elems=['lr_image', 'lr_points'], cat_name='concat_lr'), tr.ConcatInputs(cat_elems=['image', 'points'], cat_name='concat'), tr.ConcatInputs(cat_elems=['image', 'image_grad'], cat_name='grad'), tr.ToTensor() ] composed_transforms_ts = transforms.Compose(composed_transforms) # Setup dataset if args.test_set == 'thinobject5k_test': db = thinobject5k.ThinObject5K(split='test', transform=composed_transforms_ts) elif args.test_set == 'coift': db = coift.COIFT(split='test', transform=composed_transforms_ts) elif args.test_set == 'hrsod': db = hrsod.HRSOD(split='test', transform=composed_transforms_ts) else:
optimizer = optim.SGD(net.parameters(), lr=p['lr'], momentum=p['momentum'], weight_decay=p['wd']) criterion = SegmentationMultiLosses(1) p['optimizer'] = str(optimizer) # Preparation of the data loaders composed_transforms_tr = transforms.Compose([ tr.RandomHorizontalFlip(), tr.ScaleNRotate(rots=(-20, 20), scales=(.75, 1.25)), tr.CropFromMaskStatic(crop_elems=('image', 'gt'), relax = 50, zero_pad=zero_pad_crop), tr.FixedResize(resolutions={'crop_image': (512, 512), 'crop_gt': (512, 512)}), # tr.ExtremePoints(sigma=10, pert=0, elem='crop_gt', is_val = False), # tr.ToImage(norm_elem='extreme_points'), # tr.NEllipse(is_val = False), tr.NEllipseWithGaussians(alpha = 0.6, is_val = False), # tr.AddConfidenceMap(elem = ('crop_image'), hm_type = 'l1l2', tau = 7), tr.ConcatInputs(elems=('crop_image', 'nellipseWithGaussians')), tr.ToTensor()]) composed_transforms_ts = transforms.Compose([ # tr.CreateBBMask(), tr.CropFromMaskStatic(crop_elems=('image', 'gt'), relax = 50, zero_pad=zero_pad_crop), tr.FixedResize(resolutions={'void_pixels': None, 'gt': None, 'crop_image': (512, 512), 'crop_gt': (512, 512)}), # tr.ExtremePoints(sigma=10, pert=0, elem='crop_gt', is_val = True), # tr.ToImage(norm_elem='extreme_points'), # tr.NEllipse(is_val = True), tr.NEllipseWithGaussians(alpha = 0.6, is_val = True), # tr.AddConfidenceMap(elem = ('crop_image'), hm_type = 'l1l2', tau = 7), tr.ConcatInputs(elems=('crop_image', 'nellipseWithGaussians')), tr.ToTensor()]) voc_train = pascal.VOCSegmentation(split = 'train', transform=composed_transforms_tr) voc_val = pascal.VOCSegmentation(split = 'val', transform=composed_transforms_ts)
def process(image_name): # Set gpu_id to -1 to run in CPU mode, otherwise set the id of the corresponding gpu gpu_id = 0 device = torch.device("cuda:"+str(gpu_id) if torch.cuda.is_available() else "cpu") if torch.cuda.is_available(): print('Using GPU: {} '.format(gpu_id)) # Setting parameters resume_epoch = 100 # test epoch nInputChannels = 5 # Number of input channels (RGB + heatmap of IOG points) # Network definition modelName = 'IOG_pascal' net = Network(nInputChannels=nInputChannels, num_classes=1, backbone='resnet101', output_stride=16, sync_bn=None, freeze_bn=False) # load pretrain_dict pretrain_dict = torch.load('IOG_PASCAL_SBD.pth') net.load_state_dict(pretrain_dict) # net.to(device) # Generate result of the validation images net.eval() image = np.array(Image.open(image_name).convert('RGB')) im_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) roi = cv2.selectROI(im_rgb) image = image.astype(np.float32) bbox = np.zeros_like(image[..., 0]) # bbox[0: 130, 220: 320] = 1 # for ponny # bbox[220: 390, 370: 570] = 1 bbox[int(roi[1]):int(roi[1]+roi[3]), int(roi[0]):int(roi[0]+roi[2])] = 1 void_pixels = 1 - bbox sample = {'image': image, 'gt': bbox, 'void_pixels': void_pixels} trns = transforms.Compose([ tr.CropFromMask(crop_elems=('image', 'gt','void_pixels'), relax=30, zero_pad=True), tr.FixedResize(resolutions={'gt': None, 'crop_image': (512, 512), 'crop_gt': (512, 512), 'crop_void_pixels': (512, 512)},flagvals={'gt':cv2.INTER_LINEAR,'crop_image':cv2.INTER_LINEAR,'crop_gt':cv2.INTER_LINEAR,'crop_void_pixels': cv2.INTER_LINEAR}), tr.IOGPoints(sigma=10, elem='crop_gt',pad_pixel=10), tr.ToImage(norm_elem='IOG_points'), tr.ConcatInputs(elems=('crop_image', 'IOG_points')), tr.ToTensor()]) tr_sample = trns(sample) inputs = tr_sample['concat'][None] # inputs = inputs.to(device) outputs = net.forward(inputs)[-1] # outputs = fine_out.to(torch.device('cpu')) pred = np.transpose(outputs.data.numpy()[0, :, :, :], (1, 2, 0)) pred = 1 / (1 + np.exp(-pred)) pred = np.squeeze(pred) gt = tens2image(tr_sample['gt']) bbox = get_bbox(gt, pad=30, zero_pad=True) result = crop2fullmask(pred, bbox, gt, zero_pad=True, relax=0,mask_relax=False) light = np.zeros_like(image) light[:, :, 2] = 255. alpha = 0.5 blending = (alpha * light + (1 - alpha) * image) * result[..., None] + (1 - result[..., None]) * image blending[blending > 255.] = 255 cv2.imshow('resulting segmentation', cv2.cvtColor(blending.astype(np.uint8), cv2.COLOR_RGB2BGR)) cv2.waitKey(0) cv2.destroyAllWindows()