def INRIASliced(dataset_dir, grayscale):
x = sorted(find_in_dir(os.path.join(dataset_dir, 'images')))
y = sorted(find_in_dir(os.path.join(dataset_dir, 'gt')))
image_id = [os.path.basename(fname).split('_')[0] for fname in x]
unique_image_id = np.unique(image_id)
location = [basename[:6] for basename in unique_image_id
] # Geocode is first 6 characters
train_id, test_id = train_test_split(unique_image_id,
random_state=1234,
test_size=0.1,
stratify=location)
xy_train = [(image_fname, mask_fname)
for image_fname, mask_fname, image_id in zip(x, y, image_id)
if image_id in train_id]
xy_test = [(image_fname, mask_fname)
for image_fname, mask_fname, image_id in zip(x, y, image_id)
if image_id in test_id]
x_train, y_train = zip(*xy_train)
x_test, y_test = zip(*xy_test)
train_transform = aug.Sequential([
aug.ImageOnly(aug.RandomGrayscale(1.0 if grayscale else 0.5)),
aug.ImageOnly(aug.RandomBrightness()),
aug.ImageOnly(aug.RandomContrast()),
aug.VerticalFlip(),
aug.HorizontalFlip(),
aug.ShiftScaleRotate(rotate_limit=15),
aug.ImageOnly(aug.NormalizeImage(mean=INRIA_MEAN, std=INRIA_STD)),
aug.MaskOnly(aug.MakeBinary())
])
test_transform = aug.Sequential([
aug.ImageOnly(aug.NormalizeImage(mean=INRIA_MEAN, std=INRIA_STD)),
aug.MaskOnly(aug.MakeBinary())
])
train = ImageMaskDataset(x_train,
y_train,
image_loader=read_rgb,
target_loader=read_mask,
transform=train_transform,
load_in_ram=False)
test = ImageMaskDataset(x_test,
y_test,
image_loader=read_rgb,
target_loader=read_mask,
transform=test_transform,
load_in_ram=False)
num_classes = 1
return train, test, num_classes
def DSB2018Sliced(dataset_dir, grayscale, patch_size):
"""
Returns train & test dataset or DSB2018
:param dataset_dir:
:param grayscale:
:param patch_size:
:return:
"""
images = [read_rgb(x) for x in find_in_dir(os.path.join(dataset_dir, 'images'))]
masks = [read_mask(x) for x in find_in_dir(os.path.join(dataset_dir, 'masks'))]
image_ids = []
patch_images = []
patch_masks = []
for image_id, (image, mask) in enumerate(zip(images, masks)):
slicer = ImageSlicer(image.shape, patch_size, patch_size // 2)
patch_images.extend(slicer.split(image))
patch_masks.extend(slicer.split(mask))
image_ids.extend([image_id] * len(slicer.crops))
x_train, x_test, y_train, y_test = train_test_split(patch_images, patch_masks, random_state=1234, test_size=0.1, stratify=image_ids)
train_transform = aug.Sequential([
# aug.ImageOnly(aug.RandomGrayscale()),
# aug.ImageOnly(aug.RandomInvert()),
aug.ImageOnly(aug.NormalizeImage()),
# aug.ImageOnly(aug.RandomBrightness()),
# aug.ImageOnly(aug.RandomContrast()),
aug.RandomRotate90(),
aug.VerticalFlip(),
aug.HorizontalFlip(),
aug.ShiftScaleRotate(rotate_limit=15),
aug.MaskOnly(aug.MakeBinary())
])
test_transform = aug.Sequential([
aug.ImageOnly(aug.NormalizeImage()),
aug.MaskOnly(aug.MakeBinary())
])
train = InMemoryDataset(x_train, y_train, transform=train_transform)
test = InMemoryDataset(x_test, y_test, transform=test_transform)
num_classes = 1
return train, test, num_classes
def INRIA(dataset_dir, grayscale, patch_size):
x = find_in_dir(os.path.join(dataset_dir, 'images'))
y = find_in_dir(os.path.join(dataset_dir, 'gt'))
x_train, x_test, y_train, y_test = train_test_split(x, y, random_state=1234, test_size=0.1)
train_transform = aug.Sequential([
aug.RandomCrop(patch_size),
aug.ImageOnly(aug.RandomGrayscale(1.0 if grayscale else 0.5)),
aug.ImageOnly(aug.RandomInvert()),
aug.ImageOnly(aug.NormalizeImage()),
aug.ImageOnly(aug.RandomBrightness()),
aug.ImageOnly(aug.RandomContrast()),
aug.VerticalFlip(),
aug.HorizontalFlip(),
aug.ShiftScaleRotate(),
aug.MaskOnly(aug.MakeBinary()),
aug.ToTensors()
])
test_transform = aug.Sequential([
aug.CenterCrop(patch_size, patch_size),
aug.ImageOnly(aug.RandomGrayscale(1.0 if grayscale else 0.5)),
aug.ImageOnly(aug.NormalizeImage()),
aug.MaskOnly(aug.MakeBinary()),
aug.ToTensors()
])
train = ImageMaskDataset(x_train, y_train, image_loader=read_rgb, target_loader=read_gray, transform=train_transform, load_in_ram=False)
test = ImageMaskDataset(x_test, y_test, image_loader=read_rgb, target_loader=read_gray, transform=test_transform, load_in_ram=False)
num_classes = 1
return train, test, num_classes
def INRIA(dataset_dir, grayscale, patch_size, keep_in_mem, small=False):
x = sorted(find_in_dir(os.path.join(dataset_dir, 'images')))
y = sorted(find_in_dir(os.path.join(dataset_dir, 'gt')))
if small:
x = x[:4]
y = y[:4]
x_train, x_test, y_train, y_test = train_test_split(x,
y,
random_state=1234,
test_size=0.1)
train_transform = aug.Sequential([
aug.ImageOnly(aug.RandomGrayscale(1.0 if grayscale else 0.5)),
aug.ImageOnly(aug.RandomBrightness()),
aug.ImageOnly(aug.RandomContrast()),
aug.VerticalFlip(),
aug.HorizontalFlip(),
aug.ShiftScaleRotate(rotate_limit=15),
aug.ImageOnly(aug.NormalizeImage(mean=INRIA_MEAN, std=INRIA_STD)),
aug.MaskOnly(aug.MakeBinary())
])
test_transform = aug.Sequential([
aug.ImageOnly(aug.NormalizeImage(mean=INRIA_MEAN, std=INRIA_STD)),
aug.MaskOnly(aug.MakeBinary())
])
train = TiledImagesDataset(x_train,
y_train,
patch_size,
target_shape=(5000, 5000),
transform=train_transform,
keep_in_mem=keep_in_mem)
test = TiledImagesDataset(x_test,
y_test,
patch_size,
target_shape=(5000, 5000),
transform=test_transform,
keep_in_mem=keep_in_mem)
num_classes = 1
return train, test, num_classes
def DSB2018(dataset_dir, grayscale, patch_size):
"""
Returns train & test dataset or DSB2018
:param dataset_dir:
:param grayscale:
:param patch_size:
:return:
"""
images = find_in_dir(os.path.join(dataset_dir, 'images'))
masks = find_in_dir(os.path.join(dataset_dir, 'masks'))
x_train, x_test, y_train, y_test = train_test_split(images, masks, random_state=1234, test_size=0.1)
train_transform = aug.Sequential([
aug.RandomCrop(patch_size),
# aug.ImageOnly(aug.RandomGrayscale()),
# aug.ImageOnly(aug.RandomInvert()),
aug.ImageOnly(aug.NormalizeImage()),
# aug.ImageOnly(aug.RandomBrightness()),
# aug.ImageOnly(aug.RandomContrast()),
# aug.RandomRotate90(),
# aug.VerticalFlip(),
# aug.HorizontalFlip(),
# aug.ShiftScaleRotate(),
aug.MaskOnly(aug.MakeBinary())
])
test_transform = aug.Sequential([
aug.CenterCrop(patch_size, patch_size),
# aug.ImageOnly(aug.RandomGrayscale(1.0 if grayscale else 0.0)),
aug.ImageOnly(aug.NormalizeImage()),
aug.MaskOnly(aug.MakeBinary()),
])
train = ImageMaskDataset(x_train, y_train, image_loader=read_rgb, target_loader=read_mask, transform=train_transform, load_in_ram=False)
test = ImageMaskDataset(x_test, y_test, image_loader=read_rgb, target_loader=read_mask, transform=test_transform, load_in_ram=False)
num_classes = 1
return train, test, num_classes
def main():
cudnn.benchmark = True
parser = argparse.ArgumentParser()
parser.add_argument('-g',
'--grayscale',
action='store_true',
help='Whether to use grayscale image instead of RGB')
parser.add_argument('-m',
'--model',
required=True,
type=str,
help='Name of the model')
parser.add_argument('-c',
'--checkpoint',
required=True,
type=str,
help='Name of the model checkpoint')
parser.add_argument('-p', '--patch-size', type=int, default=224)
parser.add_argument('-b',
'--batch-size',
type=int,
default=1,
help='Batch Size during training, e.g. -b 64')
parser.add_argument('-dd',
'--data-dir',
type=str,
default='data',
help='Root directory where datasets are located.')
parser.add_argument('-x',
'--experiment',
type=str,
help='Name of the experiment')
parser.add_argument('-f', '--full', action='store_true')
args = parser.parse_args()
if args.experiment is None:
args.experiment = 'inria_%s_%d_%s' % (
args.model, args.patch_size, 'gray' if args.grayscale else 'rgb')
experiment_dir = os.path.join('submits', args.experiment)
os.makedirs(experiment_dir, exist_ok=True)
model = TT.get_model(args.model,
patch_size=args.patch_size,
num_channels=1 if args.grayscale else 3).cuda()
start_epoch, train_history, best_loss = TT.restore_snapshot(
model, None, auto_file(args.checkpoint))
print('Using weights from epoch', start_epoch - 1, best_loss)
test_transform = aug.Sequential([
aug.ImageOnly(aug.NormalizeImage(mean=INRIA_MEAN, std=INRIA_STD)),
])
x = sorted(find_in_dir(os.path.join(args.data_dir, 'images')))
# x = x[:10]
model.eval()
with torch.no_grad():
for test_fname in tqdm(x, total=len(x)):
image = read_rgb(test_fname)
basename = os.path.splitext(os.path.basename(test_fname))[0]
if args.full:
mask = predict_full(image, model, test_transform)
else:
mask = predict_tiled(image, model, test_transform,
args.patch_size, args.batch_size)
mask = ((mask > 0.5) * 255).astype(np.uint8)
cv2.imwrite(os.path.join(experiment_dir, basename + '.tif'), mask)