def get_input_pair(self, data_info_row): if len(self.channels) == 0: raise Exception('You have to specify at least one channel.') instance_name = '_'.join( [data_info_row['name'], str(data_info_row['position'])]) image_path = get_filepath(self.dataset_path, data_info_row['name'], self.images_folder, instance_name, file_type=self.image_type) mask_path = get_filepath(self.dataset_path, data_info_row['name'], self.masks_folder, instance_name, file_type=self.mask_type) images_array = filter_by_channels(read_tensor(image_path), self.channels, self.neighbours) if images_array.ndim == 2: images_array = np.expand_dims(images_array, -1) masks_array = read_tensor(mask_path) aug = Compose([ RandomRotate90(), Flip(), OneOf( [ RandomSizedCrop(min_max_height=(int( self.image_size * 0.7), self.image_size), height=self.image_size, width=self.image_size), RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15), #MedianBlur(blur_limit=3, p=0.2), MaskDropout(p=0.6), ElasticTransform(alpha=15, sigma=5, alpha_affine=5), GridDistortion(p=0.6) ], p=0.8), ToTensor() ]) augmented = aug(image=images_array, mask=masks_array) augmented_images = augmented['image'] augmented_masks = augmented['mask'] if self.classification_head: masks_class = ((augmented_masks.sum() > 0) * 1).unsqueeze(-1).float() #.type(torch.FloatTensor) return augmented_images, [augmented_masks, masks_class] else: return {'features': augmented_images, 'targets': augmented_masks}
def predict(data_path, model_weights_path, network, test_df_path, save_path, size, channels, neighbours, classification_head): model = get_model(network, classification_head) model.encoder.conv1 = nn.Conv2d(count_channels(channels) * neighbours, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) model, device = UtilsFactory.prepare_model(model) if classification_head: model.load_state_dict(torch.load(model_weights_path)) else: checkpoint = torch.load(model_weights_path, map_location='cpu') model.load_state_dict(checkpoint['model_state_dict']) test_df = pd.read_csv(test_df_path) predictions_path = os.path.join(save_path, "predictions") if not os.path.exists(predictions_path): os.makedirs(predictions_path, exist_ok=True) print("Prediction directory created.") for _, image_info in tqdm(test_df.iterrows()): filename = '_'.join([image_info['name'], image_info['position']]) image_path = get_filepath(data_path, image_info['dataset_folder'], 'images', filename, file_type='tiff') image_tensor = filter_by_channels(read_tensor(image_path), channels, neighbours) if image_tensor.ndim == 2: image_tensor = np.expand_dims(image_tensor, -1) image = transforms.ToTensor()(image_tensor) if classification_head: prediction, label = model.predict( image.view(1, count_channels(channels) * neighbours, size, size).to(device, dtype=torch.float)) else: prediction = model.predict( image.view(1, count_channels(channels) * neighbours, size, size).to(device, dtype=torch.float)) result = prediction.view(size, size).detach().cpu().numpy() cv.imwrite(get_filepath(predictions_path, filename, file_type='png'), result * 255)
def __getitem__(self, idx): if len(self.channels) < 2: raise Exception('You have to specify at least two channels.') data_info_row = self.df.iloc[idx] instance_name = '_'.join( [data_info_row['name'], data_info_row['position']]) images_array, masks_array = [], [] #for k in range(1,self.num_images+1): for k in range(self.num_images, 0, -1): image_path = get_filepath(self.dataset_path, data_info_row['dataset_folder'], self.images_folder, instance_name + f'_{k}', file_type=self.image_type) img = filter_by_channels(read_tensor(image_path), self.channels, 1) images_array.append(img) mask_path = get_filepath(self.dataset_path, data_info_row['dataset_folder'], self.masks_folder, instance_name + f'_{k}', file_type=self.mask_type) msk = read_tensor(mask_path) masks_array.append(np.expand_dims(msk, axis=-1)) aug = Compose([ RandomRotate90(), Flip(), OneOf([ RandomSizedCrop(min_max_height=(int( self.image_size * 0.7), self.image_size), height=self.image_size, width=self.image_size), RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15), ElasticTransform(alpha=15, sigma=5, alpha_affine=5), GridDistortion(p=0.6) ], p=0.8), ToTensor() ]) augmented = aug(image=np.concatenate(images_array, axis=-1), mask=np.concatenate(masks_array, axis=-1)) augmented_images = torch.stack([ augmented['image'][num_img * count_channels(self.channels):(num_img + 1) * count_channels(self.channels), :, :] for num_img in range(self.num_images) ]) if self.all_masks: augmented_masks = torch.stack([ augmented['mask'][:, :, :, i] for i in range(augmented['mask'].shape[-1]) ]).squeeze() else: augmented_masks = torch.stack([augmented['mask'][:, :, :, -1]]) return { 'features': augmented_images, 'targets': augmented_masks, 'name': data_info_row['name'], 'position': data_info_row['position'] }
def __getitem__(self, idx): if len(self.channels) < 2: raise Exception('You have to specify at least two channels.') data_info_row = self.df.iloc[idx] instance_name = '_'.join( [data_info_row['name'], data_info_row['position']]) images_array, masks_array = [], [] for k in range(1, self.num_images + 1): image_path = get_filepath(self.dataset_path, data_info_row['dataset_folder'], self.images_folder, instance_name + f'_{k}', file_type=self.image_type) img = filter_by_channels(read_tensor(image_path), self.channels, 1) images_array.append(img) mask_path = get_filepath(self.dataset_path, data_info_row['dataset_folder'], self.masks_folder, instance_name, file_type=self.mask_type) masks_array = read_tensor(mask_path) if self.phase == 'train': aug = Compose([ RandomRotate90(), Flip(), OneOf( [ RandomSizedCrop(min_max_height=(int( self.image_size * 0.7), self.image_size), height=self.image_size, width=self.image_size), RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15), #MedianBlur(blur_limit=3, p=0.2), MaskDropout(p=0.6), ElasticTransform(alpha=15, sigma=5, alpha_affine=5), GridDistortion(p=0.6) ], p=0.8), ToTensor() ]) else: aug = ToTensor() ''' keys = ['image'] values = [images_array[0]] for k in range(self.num_images-1): keys.append(f'image{k}') values.append(images_array[k+1]) keys.append('mask') values.append(masks_array) #{"image" : images_array[0], "image2" : images_array[1], ..., "mask": masks_array, ...} aug_input = { keys[i] : values[i] for i in range(len(keys)) } augmented = aug(**aug_input) augmented_images = [augmented['image']] for k in range(self.num_images-1): augmented_images.append(np.transpose(augmented[f'image{k}'], ( 2, 0, 1))/255) augmented_masks = [augmented['mask']] return {'features': augmented_images, 'targets': augmented_masks, 'name': data_info_row['name'], 'position': data_info_row['position']} ''' augmented = aug(image=np.concatenate( (images_array[0], images_array[1]), axis=-1), mask=masks_array) augmented_images = [ augmented['image'][:count_channels(self.channels), :, :], augmented['image'][count_channels(self.channels):, :, :] ] augmented_masks = [augmented['mask']] return { 'features': augmented_images, 'targets': augmented_masks, 'name': data_info_row['name'], 'position': data_info_row['position'] }