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