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 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, 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