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
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)
if resume_epoch != nEpochs: # Logging into Tensorboard 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:
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=30, elem='crop_gt', num_pts=args.num_pts, type=args.point_type, vis=True if args.debug else False), 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,
net.to(device) net.eval() # Setup data transformations composed_transforms = [ tr.IdentityTransform(tr_elems=['gt'], prefix='ori_'), tr.CropFromMask(crop_elems=['image', 'gt'], relax=cfg['relax_crop'], zero_pad=cfg['zero_pad_crop'], adaptive_relax=cfg['adaptive_relax'], prefix=''), tr.Resize(resize_elems=['image', 'gt', 'void_pixels'], min_size=cfg['min_size'], max_size=cfg['max_size']), tr.ComputeImageGradient(elem='image'), tr.ExtremePoints(sigma=10, pert=0, elem='gt'), tr.GaussianTransform(tr_elems=['extreme_points'], mask_elem='gt', sigma=10, tr_name='points'), 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'],
p = OrderedDict() # Parameters to include in report classifier = 'psp' # Head classifier to use useTest = 1 # See evolution of the test set when training? nTestInterval = 10 # Run on test set every nTestInterval epochs snapshot = 20 # Store a model every snapshot epochs relax_crop = 50 # Enlarge the bounding box by relax_crop pixels nInputChannels = 4 # Number of input channels (RGB + heatmap of extreme points) zero_pad_crop = True # Insert zero padding when cropping the image 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=30, elem='crop_gt', num_pts=50, type='polygon', vis=False), 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=8) modelName = 'dextr_pascal' net = resnet.resnet101(1, pretrained=True, nInputChannels=nInputChannels, classifier=classifier) net = torch.nn.DataParallel(net.cuda()) # net = net.cuda() print("Initializing weights from: {}".format( args.model_path)) net.load_state_dict(