def create_dataloader(samples_folder, batch_size, num_devices, params):
"""
Function to create dataloader objects for training, validation and test datasets.
:param samples_folder: path to the folder containting .hdf5 files if task is segmentation
:param batch_size: (int) batch size
:param num_devices: (int) number of GPUs used
:param params: (dict) Parameters found in the yaml config file.
:return: trn_dataloader, val_dataloader, tst_dataloader
"""
debug = get_key_def('debug_mode', params['global'], False)
dontcare_val = get_key_def("ignore_index", params["training"], -1)
assert samples_folder.is_dir(), f'Could not locate: {samples_folder}'
assert len([f for f in samples_folder.glob('**/*.hdf5')]) >= 1, f"Couldn't locate .hdf5 files in {samples_folder}"
num_samples = get_num_samples(samples_path=samples_folder, params=params)
assert num_samples['trn'] >= batch_size and num_samples['val'] >= batch_size, f"Number of samples in .hdf5 files is less than batch size"
print(f"Number of samples : {num_samples}\n")
meta_map = get_key_def("meta_map", params["global"], {})
num_bands = get_key_def("number_of_bands", params["global"], {})
if not meta_map:
dataset_constr = CreateDataset.SegmentationDataset
else:
dataset_constr = functools.partial(CreateDataset.MetaSegmentationDataset, meta_map=meta_map)
datasets = []
for subset in ["trn", "val", "tst"]:
# TODO: transform the aug.compose_transforms in a class with radiom, geom, totensor in def
datasets.append(dataset_constr(samples_folder, subset, num_bands,
max_sample_count=num_samples[subset],
dontcare=dontcare_val,
radiom_transform=aug.compose_transforms(params, subset, type='radiometric'),
geom_transform=aug.compose_transforms(params, subset, type='geometric',
ignore_index=dontcare_val),
totensor_transform=aug.compose_transforms(params, subset, type='totensor'),
debug=debug))
trn_dataset, val_dataset, tst_dataset = datasets
# https://discuss.pytorch.org/t/guidelines-for-assigning-num-workers-to-dataloader/813/5
num_workers = num_devices * 4 if num_devices > 1 else 4
# Shuffle must be set to True.
trn_dataloader = DataLoader(trn_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True,
drop_last=True)
# Using batch_metrics with shuffle=False on val dataset will always mesure metrics on the same portion of the val samples.
# Shuffle should be set to True.
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True,
drop_last=True)
tst_dataloader = DataLoader(tst_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False,
drop_last=True) if num_samples['tst'] > 0 else None
return trn_dataloader, val_dataloader, tst_dataloader
def create_dataloader(data_path,
batch_size,
task,
num_devices,
params,
samples_folder=None):
"""
Function to create dataloader objects for training, validation and test datasets.
:param data_path: (str) path to the samples folder
:param batch_size: (int) batch size
:param task: (str) classification or segmentation
:param num_devices: (int) number of GPUs used
:param samples_folder: path to folder containting .hdf5 files if task is segmentation
:param params: (dict) Parameters found in the yaml config file.
:return: trn_dataloader, val_dataloader, tst_dataloader
"""
assert Path(samples_folder).is_dir(), f'Could not locate: {samples_folder}'
assert len([f for f in Path(samples_folder).glob('**/*.hdf5')
]) >= 1, f"Couldn't locate .hdf5 files in {samples_folder}"
num_samples = get_num_samples(samples_path=samples_folder, params=params)
print(f"Number of samples : {num_samples}\n")
meta_map = get_key_def("meta_map", params["global"], {})
if not meta_map:
dataset_constr = CreateDataset.SegmentationDataset
else:
dataset_constr = functools.partial(
CreateDataset.MetaSegmentationDataset, meta_map=meta_map)
dontcare = get_key_def("ignore_index", params["training"], None)
if dontcare == 0:
warnings.warn(
"The 'dontcare' value (or 'ignore_index') used in the loss function cannot be zero;"
" all valid class indices should be consecutive, and start at 0. The 'dontcare' value"
" will be remapped to -1 while loading the dataset, and inside the config from now on."
)
params["training"]["ignore_index"] = -1
datasets = []
for subset in ["trn", "val", "tst"]:
datasets.append(
dataset_constr(samples_folder,
subset,
max_sample_count=num_samples[subset],
dontcare=dontcare,
transform=aug.compose_transforms(params, subset)))
trn_dataset, val_dataset, tst_dataset = datasets
# https://discuss.pytorch.org/t/guidelines-for-assigning-num-workers-to-dataloader/813/5
num_workers = num_devices * 4 if num_devices > 1 else 4
# Shuffle must be set to True.
trn_dataloader = DataLoader(trn_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True,
drop_last=True)
val_dataloader = DataLoader(val_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False)
tst_dataloader = DataLoader(
tst_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False) if num_samples['tst'] > 0 else None
return trn_dataloader, val_dataloader, tst_dataloader
def sem_seg_inference(model,
nd_array,
overlay,
chunk_size,
num_classes,
device,
meta_map=None,
metadata=None,
output_path=Path(os.getcwd()),
index=0,
debug=False):
"""Inference on images using semantic segmentation
Args:
model: model to use for inference
nd_array: input raster as array
overlay: amount of overlay to apply
chunk_size: size of individual chunks to be processed during inference
num_classes: number of different classes that may be predicted by the model
device: device used by pytorch (cpu ou cuda)
meta_map:
metadata:
output_path: path to save debug files
index: (int) index of array from list of images on which inference is performed
returns a numpy array of the same size (h,w) as the input image, where each value is the predicted output.
"""
# switch to evaluate mode
model.eval()
if len(nd_array.shape) == 3:
h, w, nb = nd_array.shape
# Pad with overlay on left and top and pad with chunk_size on right and bottom
padded_array = np.pad(nd_array, ((overlay, chunk_size), (overlay, chunk_size), (0, 0)), mode='constant')
elif len(nd_array.shape) == 2:
h, w = nd_array.shape
padded_array = np.expand_dims(np.pad(nd_array, ((overlay, chunk_size), (overlay, chunk_size)),
mode='constant'), axis=0)
else:
h = 0
w = 0
padded_array = None
h_padded, w_padded = padded_array.shape[:2]
# Create an empty array of dimensions (c x h x w): num_classes x height of padded array x width of padded array
output_probs = np.empty([num_classes, h_padded, w_padded], dtype=np.float32)
# Create identical 0-filled array without channels dimension to receive counts for number of outputs generated in specific area.
output_counts = np.zeros([output_probs.shape[1], output_probs.shape[2]], dtype=np.int32)
if padded_array.any():
with torch.no_grad():
for row in tqdm(range(overlay, h + chunk_size, chunk_size - overlay), position=1, leave=False,
desc=f'Inferring rows with "{device}"'):
row_start = row - overlay
row_end = row_start + chunk_size
with tqdm(range(overlay, w + chunk_size, chunk_size - overlay), position=2, leave=False, desc='Inferring columns') as _tqdm:
for col in _tqdm:
sample = {'sat_img': None, 'metadata': {'dtype': None}}
sample['metadata'] = metadata
totensor_transform = augmentation.compose_transforms(params, dataset="tst", type='totensor')
col_start = col - overlay
col_end = col_start + chunk_size
sample['sat_img'] = padded_array[row_start:row_end, col_start:col_end, :]
if meta_map:
sample['sat_img'] = MetaSegmentationDataset.append_meta_layers(sample['sat_img'], meta_map, metadata)
sample = totensor_transform(sample)
inputs = sample['sat_img'].unsqueeze_(0) # Add dummy batch dimension
inputs = inputs.to(device)
# forward
outputs = model(inputs)
# torchvision models give output in 'out' key. May cause problems in future versions of torchvision.
if isinstance(outputs, OrderedDict) and 'out' in outputs.keys():
outputs = outputs['out']
if debug:
if index == 0:
tqdm.write(f'(debug mode) Visualizing inferred tiles...')
vis_from_batch(params, inputs, outputs, batch_index=0, vis_path=output_path,
dataset=f'{row_start}_{col_start}_inf', ep_num=index, debug=True)
outputs = F.softmax(outputs, dim=1)
output_counts[row_start:row_end, col_start:col_end] += 1
# Add inference on sub-image to all completed inferences on previous sub-images.
# TODO: This operation need to be optimized. Using a lot of RAM on large images.
output_probs[:, row_start:row_end, col_start:col_end] += np.squeeze(outputs.cpu().numpy(),
axis=0)
if debug and device.type == 'cuda':
res, mem = gpu_stats(device=device.index)
_tqdm.set_postfix(OrderedDict(gpu_perc=f'{res.gpu} %',
gpu_RAM=f'{mem.used / (1024 ** 2):.0f}/{mem.total / (1024 ** 2):.0f} MiB',
inp_size=inputs.cpu().numpy().shape,
out_size=outputs.cpu().numpy().shape,
overlay=overlay))
# Divide array according to output counts. Manages overlap and returns a softmax array as if only one forward pass had been done.
output_mask_raw = np.divide(output_probs, np.maximum(output_counts, 1)) # , 1 is added to overwrite 0 values.
# Resize the output array to the size of the input image and write it
output_mask_raw_cropped = np.moveaxis(output_mask_raw, 0, -1)
output_mask_raw_cropped = output_mask_raw_cropped[overlay:(h + overlay), overlay:(w + overlay), :]
return output_mask_raw_cropped
else:
raise IOError(f"Error classifying image : Image shape of {len(nd_array.shape)} is not recognized")
def segmentation(
param,
input_image,
label_arr,
num_classes: int,
gpkg_name,
model,
chunk_size: int,
device,
scale: List,
BGR_to_RGB: bool,
tp_mem,
debug=False,
):
"""
Args:
param: parameter dict
input_image: opened image (rasterio object)
label_arr: numpy array of label if available
num_classes: number of classes
gpkg_name: geo-package name if available
model: model weights
chunk_size: image tile size
device: cuda/cpu device
scale: scale range
BGR_to_RGB: True/False
tp_mem: memory temp file for saving numpy array to disk
debug: True/False
Returns:
"""
xmin, ymin, xmax, ymax = (input_image.bounds.left,
input_image.bounds.bottom,
input_image.bounds.right, input_image.bounds.top)
xres, yres = (abs(input_image.transform.a), abs(input_image.transform.e))
mx = chunk_size * xres
my = chunk_size * yres
padded = chunk_size * 2
h = input_image.height
w = input_image.width
h_ = h + padded
w_ = w + padded
dist_samples = int(round(chunk_size * (1 - 1.0 / 2.0)))
# switch to evaluate mode
model.eval()
# initialize test time augmentation
transforms = tta.Compose([
tta.HorizontalFlip(),
])
# construct window for smoothing
WINDOW_SPLINE_2D = _window_2D(window_size=padded, power=2.0)
WINDOW_SPLINE_2D = torch.as_tensor(np.moveaxis(WINDOW_SPLINE_2D, 2,
0), ).type(torch.float)
WINDOW_SPLINE_2D = WINDOW_SPLINE_2D.to(device)
fp = np.memmap(tp_mem,
dtype='float16',
mode='w+',
shape=(h_, w_, num_classes))
sample = {'sat_img': None, 'map_img': None, 'metadata': None}
cnt = 0
img_gen = gen_img_samples(input_image, chunk_size)
start_seg = time.time()
for img in tqdm(img_gen,
position=1,
leave=False,
desc='inferring on window slices'):
row = img[1]
col = img[2]
sub_image = img[0]
image_metadata = add_metadata_from_raster_to_sample(
sat_img_arr=sub_image,
raster_handle=input_image,
meta_map={},
raster_info={})
sample['metadata'] = image_metadata
totensor_transform = augmentation.compose_transforms(
param,
dataset="tst",
input_space=BGR_to_RGB,
scale=scale,
aug_type='totensor')
sample['sat_img'] = sub_image
sample = totensor_transform(sample)
inputs = sample['sat_img'].unsqueeze_(0)
inputs = inputs.to(device)
if inputs.shape[1] == 4 and any("module.modelNIR" in s
for s in model.state_dict().keys()):
############################
# Test Implementation of the NIR
############################
# Init NIR TODO: make a proper way to read the NIR channel
# and put an option to be able to give the idex of the NIR channel
# Extract the NIR channel -> [batch size, H, W] since it's only one channel
inputs_NIR = inputs[:, -1, ...]
# add a channel to get the good size -> [:, 1, :, :]
inputs_NIR.unsqueeze_(1)
# take out the NIR channel and take only the RGB for the inputs
inputs = inputs[:, :-1, ...]
# Suggestion of implementation
# inputs_NIR = data['NIR'].to(device)
inputs = [inputs, inputs_NIR]
# outputs = model(inputs, inputs_NIR)
############################
# End of the test implementation module
############################
output_lst = []
for transformer in transforms:
# augment inputs
augmented_input = transformer.augment_image(inputs)
augmented_output = model(augmented_input)
if isinstance(augmented_output,
OrderedDict) and 'out' in augmented_output.keys():
augmented_output = augmented_output['out']
logging.debug(
f'Shape of augmented output: {augmented_output.shape}')
# reverse augmentation for outputs
deaugmented_output = transformer.deaugment_mask(augmented_output)
deaugmented_output = F.softmax(deaugmented_output,
dim=1).squeeze(dim=0)
output_lst.append(deaugmented_output)
outputs = torch.stack(output_lst)
outputs = torch.mul(outputs, WINDOW_SPLINE_2D)
outputs, _ = torch.max(outputs, dim=0)
outputs = outputs.permute(1, 2, 0)
outputs = outputs.reshape(padded, padded,
num_classes).cpu().numpy().astype('float16')
outputs = outputs[dist_samples:-dist_samples,
dist_samples:-dist_samples, :]
fp[row:row + chunk_size, col:col + chunk_size, :] = \
fp[row:row + chunk_size, col:col + chunk_size, :] + outputs
cnt += 1
fp.flush()
del fp
fp = np.memmap(tp_mem,
dtype='float16',
mode='r',
shape=(h_, w_, num_classes))
subdiv = 2.0
step = int(chunk_size / subdiv)
pred_img = np.zeros((h_, w_), dtype=np.uint8)
for row in tqdm(range(0, input_image.height, step),
position=2,
leave=False):
for col in tqdm(range(0, input_image.width, step),
position=3,
leave=False):
arr1 = fp[row:row + chunk_size, col:col + chunk_size, :] / (2**2)
arr1 = arr1.argmax(axis=-1).astype('uint8')
pred_img[row:row + chunk_size, col:col + chunk_size] = arr1
pred_img = pred_img[:h, :w]
end_seg = time.time() - start_seg
logging.info('Segmentation operation completed in {:.0f}m {:.0f}s'.format(
end_seg // 60, end_seg % 60))
if debug:
logging.debug(
f'Bin count of final output: {np.unique(pred_img, return_counts=True)}'
)
gdf = None
if label_arr is not None:
start_seg_ = time.time()
feature = defaultdict(list)
cnt = 0
for row in tqdm(range(0, h, chunk_size), position=2, leave=False):
for col in tqdm(range(0, w, chunk_size), position=3, leave=False):
label = label_arr[row:row + chunk_size, col:col + chunk_size]
pred = pred_img[row:row + chunk_size, col:col + chunk_size]
pixelMetrics = ComputePixelMetrics(label.flatten(),
pred.flatten(), num_classes)
eval = pixelMetrics.update(pixelMetrics.iou)
feature['id_image'].append(gpkg_name)
for c_num in range(num_classes):
feature['L_count_' + str(c_num)].append(
int(np.count_nonzero(label == c_num)))
feature['P_count_' + str(c_num)].append(
int(np.count_nonzero(pred == c_num)))
feature['IoU_' + str(c_num)].append(eval['iou_' +
str(c_num)])
feature['mIoU'].append(eval['macro_avg_iou'])
x_1, y_1 = (xmin + (col * xres)), (ymax - (row * yres))
x_2, y_2 = (xmin + ((col * xres) + mx)), y_1
x_3, y_3 = x_2, (ymax - ((row * yres) + my))
x_4, y_4 = x_1, y_3
geom = Polygon([(x_1, y_1), (x_2, y_2), (x_3, y_3),
(x_4, y_4)])
feature['geometry'].append(geom)
feature['length'].append(geom.length)
feature['pointx'].append(geom.centroid.x)
feature['pointy'].append(geom.centroid.y)
feature['area'].append(geom.area)
cnt += 1
gdf = gpd.GeoDataFrame(feature, crs=input_image.crs)
gdf.to_crs(crs="EPSG:4326", inplace=True)
end_seg_ = time.time() - start_seg_
logging.info('Benchmark operation completed in {:.0f}m {:.0f}s'.format(
end_seg_ // 60, end_seg_ % 60))
input_image.close()
return pred_img, gdf
def create_dataloader(samples_folder: Path,
batch_size: int,
eval_batch_size: int,
gpu_devices_dict: dict,
sample_size: int,
dontcare_val: int,
crop_size: int,
meta_map,
num_bands: int,
BGR_to_RGB: bool,
scale: Sequence,
params: dict,
dontcare2backgr: bool = False,
calc_eval_bs: bool = False,
debug: bool = False):
"""
Function to create dataloader objects for training, validation and test datasets.
:param samples_folder: path to folder containting .hdf5 files if task is segmentation
:param batch_size: (int) batch size
:param gpu_devices_dict: (dict) dictionary where each key contains an available GPU with its ram info stored as value
:param sample_size: (int) size of hdf5 samples (used to evaluate eval batch-size)
:param dontcare_val: (int) value in label to be ignored during loss calculation
:param meta_map: metadata mapping object
:param num_bands: (int) number of bands in imagery
:param BGR_to_RGB: (bool) if True, BGR channels will be flipped to RGB
:param scale: (List) imagery data will be scaled to this min and max value (ex.: 0 to 1)
:param params: (dict) Parameters found in the yaml config file.
:param dontcare2backgr: (bool) if True, all dontcare values in label will be replaced with 0 (background value)
before training
:return: trn_dataloader, val_dataloader, tst_dataloader
"""
if not samples_folder.is_dir():
raise FileNotFoundError(f'Could not locate: {samples_folder}')
if not len([f for f in samples_folder.glob('**/*.hdf5')]) >= 1:
raise FileNotFoundError(f"Couldn't locate .hdf5 files in {samples_folder}")
num_samples, samples_weight = get_num_samples(samples_path=samples_folder, params=params, dontcare=dontcare_val)
if not num_samples['trn'] >= batch_size and num_samples['val'] >= batch_size:
raise ValueError(f"Number of samples in .hdf5 files is less than batch size")
logging.info(f"Number of samples : {num_samples}\n")
if not meta_map:
dataset_constr = create_dataset.SegmentationDataset
else:
dataset_constr = functools.partial(create_dataset.MetaSegmentationDataset, meta_map=meta_map)
datasets = []
for subset in ["trn", "val", "tst"]:
datasets.append(dataset_constr(samples_folder, subset, num_bands,
max_sample_count=num_samples[subset],
dontcare=dontcare_val,
radiom_transform=aug.compose_transforms(params=params,
dataset=subset,
aug_type='radiometric'),
geom_transform=aug.compose_transforms(params=params,
dataset=subset,
aug_type='geometric',
dontcare=dontcare_val,
crop_size=crop_size),
totensor_transform=aug.compose_transforms(params=params,
dataset=subset,
input_space=BGR_to_RGB,
scale=scale,
dontcare2backgr=dontcare2backgr,
dontcare=dontcare_val,
aug_type='totensor'),
params=params,
debug=debug))
trn_dataset, val_dataset, tst_dataset = datasets
# https://discuss.pytorch.org/t/guidelines-for-assigning-num-workers-to-dataloader/813/5
num_workers = len(gpu_devices_dict.keys()) * 4 if len(gpu_devices_dict.keys()) > 1 else 4
samples_weight = torch.from_numpy(samples_weight)
sampler = torch.utils.data.sampler.WeightedRandomSampler(samples_weight.type('torch.DoubleTensor'),
len(samples_weight))
if gpu_devices_dict and calc_eval_bs:
max_pix_per_mb_gpu = 280 # TODO: this value may need to be finetuned
eval_batch_size = calc_eval_batchsize(gpu_devices_dict, batch_size, sample_size, max_pix_per_mb_gpu)
trn_dataloader = DataLoader(trn_dataset, batch_size=batch_size, num_workers=num_workers, sampler=sampler,
drop_last=True)
val_dataloader = DataLoader(val_dataset, batch_size=eval_batch_size, num_workers=num_workers, shuffle=False,
drop_last=True)
tst_dataloader = DataLoader(tst_dataset, batch_size=eval_batch_size, num_workers=num_workers, shuffle=False,
drop_last=True) if num_samples['tst'] > 0 else None
return trn_dataloader, val_dataloader, tst_dataloader
def segmentation_with_smoothing(raster, clip_gpkg, model, sample_size, overlap,
num_bands, device):
# switch to evaluate mode
model.eval()
img_array, input_image, dataset_nodata = image_reader_as_array(
input_image=raster, clip_gpkg=clip_gpkg)
metadata = add_metadata_from_raster_to_sample(img_array,
input_image,
meta_map=None,
raster_info=None)
h, w, bands = img_array.shape
assert num_bands <= bands, f"Num of specified bands is not compatible with image shape {img_array.shape}"
if num_bands < bands:
img_array = img_array[:, :, :num_bands]
padding = int(round(sample_size * (1 - 1.0 / overlap)))
padded_img = pad(img_array, padding=padding, fill=0)
WINDOW_SPLINE_2D = _window_2D(window_size=sample_size, power=1)
WINDOW_SPLINE_2D = np.moveaxis(WINDOW_SPLINE_2D, 2, 0)
step = int(sample_size / overlap)
h_, w_ = padded_img.shape[:2]
pred_img = np.empty((h_, w_), dtype=np.uint8)
for row in tqdm(range(0, h_ - sample_size + 1, step),
position=1,
leave=False,
desc='Inferring rows'):
with tqdm(range(0, w_ - sample_size + 1, step),
position=2,
leave=False,
desc='Inferring columns') as _tqdm:
for col in _tqdm:
sample = {'sat_img': None, 'metadata': None}
sample['metadata'] = metadata
totensor_transform = augmentation.compose_transforms(
params, dataset="tst", type='totensor')
sub_images = padded_img[row:row + sample_size,
col:col + sample_size, :]
sample['sat_img'] = sub_images
sample = totensor_transform(sample)
inputs = sample['sat_img'].unsqueeze_(0)
inputs = inputs.to(device)
if inputs.shape[1] == 4 and any(
"module.modelNIR" in s
for s in model.state_dict().keys()):
############################
# Test Implementation of the NIR
############################
# Init NIR TODO: make a proper way to read the NIR channel
# and put an option to be able to give the idex of the NIR channel
# Extract the NIR channel -> [batch size, H, W] since it's only one channel
inputs_NIR = inputs[:, -1, ...]
# add a channel to get the good size -> [:, 1, :, :]
inputs_NIR.unsqueeze_(1)
# take out the NIR channel and take only the RGB for the inputs
inputs = inputs[:, :-1, ...]
# Suggestion of implementation
#inputs_NIR = data['NIR'].to(device)
inputs = [inputs, inputs_NIR]
#outputs = model(inputs, inputs_NIR)
############################
# End of the test implementation module
############################
outputs = model(inputs)
# torchvision models give output in 'out' key.
# May cause problems in future versions of torchvision.
if isinstance(outputs,
OrderedDict) and 'out' in outputs.keys():
outputs = outputs['out']
outputs = F.softmax(outputs,
dim=1).squeeze(dim=0).cpu().numpy()
outputs = WINDOW_SPLINE_2D * outputs
outputs = outputs.argmax(axis=0)
pred_img[row:row + sample_size,
col:col + sample_size] = outputs
pred_img = pred_img[padding:-padding, padding:-padding]
return pred_img[:h, :w]
def segmentation(raster, clip_gpkg, model, sample_size, num_bands, device):
# switch to evaluate mode
model.eval()
img_array, input_image, dataset_nodata = image_reader_as_array(
input_image=raster, clip_gpkg=clip_gpkg)
metadata = add_metadata_from_raster_to_sample(img_array,
input_image,
meta_map=None,
raster_info=None)
h, w, bands = img_array.shape
assert num_bands <= bands, f"Num of specified bands is not compatible with image shape {img_array.shape}"
if num_bands < bands:
img_array = img_array[:, :, :num_bands]
h_ = sample_size * math.ceil(h / sample_size)
w_ = sample_size * math.ceil(w / sample_size)
pred_img = np.empty((h_, w_), dtype=np.uint8)
for row in tqdm(range(0, h, sample_size),
position=1,
leave=False,
desc='Inferring rows'):
with tqdm(range(0, w, sample_size),
position=2,
leave=False,
desc='Inferring columns') as _tqdm:
for column in _tqdm:
sample = {'sat_img': None, 'metadata': None}
sample['metadata'] = metadata
totensor_transform = augmentation.compose_transforms(
params, dataset="tst", type='totensor')
sub_images = img_array[row:row + sample_size,
column:column + sample_size, :]
sub_images_row = sub_images.shape[0]
sub_images_col = sub_images.shape[1]
if sub_images_row < sample_size or sub_images_col < sample_size:
padding = pad_diff(actual_height=sub_images_row,
actual_width=sub_images_col,
desired_shape=sample_size)
sub_images = pad(
sub_images, padding, fill=0
) # FIXME combine pad and pad_diff into one function
sample['sat_img'] = sub_images
sample = totensor_transform(sample)
inputs = sample['sat_img'].unsqueeze_(0)
inputs = inputs.to(device)
if inputs.shape[1] == 4 and any(
"module.modelNIR" in s
for s in model.state_dict().keys()):
############################
# Test Implementation of the NIR
############################
# Init NIR TODO: make a proper way to read the NIR channel
# and put an option to be able to give the idex of the NIR channel
# Extract the NIR channel -> [batch size, H, W] since it's only one channel
inputs_NIR = inputs[:, -1, ...]
# add a channel to get the good size -> [:, 1, :, :]
inputs_NIR.unsqueeze_(1)
# take out the NIR channel and take only the RGB for the inputs
inputs = inputs[:, :-1, ...]
# Suggestion of implementation
#inputs_NIR = data['NIR'].to(device)
inputs = [inputs, inputs_NIR]
#outputs = model(inputs, inputs_NIR)
############################
# End of the test implementation module
############################
outputs = model(inputs)
# torchvision models give output in 'out' key. May cause problems in future versions of torchvision.
if isinstance(outputs,
OrderedDict) and 'out' in outputs.keys():
outputs = outputs['out']
outputs = F.softmax(
outputs, dim=1).argmax(dim=1).squeeze(dim=0).cpu().numpy()
pred_img[row:row + sample_size,
column:column + sample_size] = outputs
return pred_img[:h, :w]
def create_dataloader(data_path, num_samples, batch_size, task):
"""
Function to create dataloader objects for training, validation and test datasets.
:param data_path: (str) path to the samples folder
:param num_samples: (dict) number of samples for training, validation and test
:param batch_size: (int) batch size
:param task: (str) classification or segmentation
:return: trn_dataloader, val_dataloader, tst_dataloader
"""
if task == 'classification':
trn_dataset = torchvision.datasets.ImageFolder(
os.path.join(data_path, "trn"),
transform=transforms.Compose([
transforms.RandomRotation((0, 275)),
transforms.RandomHorizontalFlip(),
transforms.Resize(299),
transforms.ToTensor()
]),
loader=loader)
val_dataset = torchvision.datasets.ImageFolder(
os.path.join(data_path, "val"),
transform=transforms.Compose(
[transforms.Resize(299),
transforms.ToTensor()]),
loader=loader)
tst_dataset = torchvision.datasets.ImageFolder(
os.path.join(data_path, "tst"),
transform=transforms.Compose(
[transforms.Resize(299),
transforms.ToTensor()]),
loader=loader)
elif task == 'segmentation':
trn_dataset = CreateDataset.SegmentationDataset(
os.path.join(data_path, "samples"),
num_samples['trn'],
"trn",
transform=aug.compose_transforms(params, 'trn'))
val_dataset = CreateDataset.SegmentationDataset(
os.path.join(data_path, "samples"),
num_samples['val'],
"val",
transform=aug.compose_transforms(params, 'tst'))
tst_dataset = CreateDataset.SegmentationDataset(
os.path.join(data_path, "samples"),
num_samples['tst'],
"tst",
transform=aug.compose_transforms(params, 'tst'))
else:
raise ValueError(
f"The task should be either classification or segmentation. The provided value is {task}"
)
# Shuffle must be set to True.
# https://discuss.pytorch.org/t/guidelines-for-assigning-num-workers-to-dataloader/813/5
if torch.cuda.device_count() > 1:
num_workers = torch.cuda.device_count() * 4
else:
num_workers = 4
trn_dataloader = DataLoader(trn_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
val_dataloader = DataLoader(val_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
tst_dataloader = DataLoader(tst_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
return trn_dataloader, val_dataloader, tst_dataloader
def create_dataloader(data_path, batch_size, task, num_devices, params):
"""
Function to create dataloader objects for training, validation and test datasets.
:param data_path: (str) path to the samples folder
:param batch_size: (int) batch size
:param task: (str) classification or segmentation
:param num_devices: (int) number of GPUs used
:param params: (dict) Parameters found in the yaml config file.
:return: trn_dataloader, val_dataloader, tst_dataloader
"""
if task == 'classification':
trn_dataset = torchvision.datasets.ImageFolder(
os.path.join(data_path, "trn"),
transform=transforms.Compose([
transforms.RandomRotation((0, 275)),
transforms.RandomHorizontalFlip(),
transforms.Resize(299),
transforms.ToTensor()
]),
loader=loader)
val_dataset = torchvision.datasets.ImageFolder(
os.path.join(data_path, "val"),
transform=transforms.Compose(
[transforms.Resize(299),
transforms.ToTensor()]),
loader=loader)
tst_dataset = torchvision.datasets.ImageFolder(
os.path.join(data_path, "tst"),
transform=transforms.Compose(
[transforms.Resize(299),
transforms.ToTensor()]),
loader=loader)
elif task == 'segmentation':
num_samples = get_num_samples(data_path=data_path, params=params)
print(f"Number of samples : {num_samples}")
meta_map = get_key_def("meta_map", params["global"], {})
if not meta_map:
dataset_constr = CreateDataset.SegmentationDataset
else:
dataset_constr = functools.partial(
CreateDataset.MetaSegmentationDataset, meta_map=meta_map)
dontcare = get_key_def("ignore_index", params["training"], None)
if dontcare == 0:
warnings.warn(
"The 'dontcare' value (or 'ignore_index') used in the loss function cannot be zero;"
" all valid class indices should be consecutive, and start at 0. The 'dontcare' value"
" will be remapped to -1 while loading the dataset, and inside the config from now on."
)
params["training"]["ignore_index"] = -1
datasets = []
for subset in ["trn", "val", "tst"]:
datasets.append(
dataset_constr(os.path.join(data_path, "samples"),
subset,
max_sample_count=num_samples[subset],
dontcare=dontcare,
transform=aug.compose_transforms(
params, subset)))
trn_dataset, val_dataset, tst_dataset = datasets
else:
raise ValueError(
f"The task should be either classification or segmentation. The provided value is {task}"
)
# Shuffle must be set to True.
# https://discuss.pytorch.org/t/guidelines-for-assigning-num-workers-to-dataloader/813/5
if num_devices > 1:
num_workers = num_devices * 4
else:
num_workers = 4
trn_dataloader = DataLoader(trn_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
val_dataloader = DataLoader(val_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
tst_dataloader = DataLoader(tst_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
return trn_dataloader, val_dataloader, tst_dataloader
