class TiledSingleImageDataset(Dataset): def __init__(self, image_fname: str, mask_fname: str, image_loader: Callable, target_loader: Callable, tile_size, tile_step, image_margin=0, transform=None, target_shape=None, keep_in_mem=False): self.image_fname = image_fname self.mask_fname = mask_fname self.image_loader = image_loader self.mask_loader = target_loader self.image = None self.mask = None if target_shape is None or keep_in_mem: image = image_loader(image_fname) mask = target_loader(mask_fname) if image.shape[0] != mask.shape[0] or image.shape[1] != mask.shape[ 1]: raise ValueError( f"Image size {image.shape} and mask shape {image.shape} must have equal width and height" ) target_shape = image.shape self.slicer = ImageSlicer(target_shape, tile_size, tile_step, image_margin) if keep_in_mem: self.images = self.slicer.split(image) self.masks = self.slicer.split(mask) else: self.images = None self.masks = None self.transform = transform self.image_ids = [ id_from_fname(image_fname) + f' [{crop[0]};{crop[1]};{crop[2]};{crop[3]};]' for crop in self.slicer.crops ] def _get_image(self, index): if self.images is None: image = self.image_loader(self.image_fname) image = self.slicer.cut_patch(image, index) else: image = self.images[index] return image def _get_mask(self, index): if self.masks is None: mask = self.mask_loader(self.mask_fname) mask = self.slicer.cut_patch(mask, index) else: mask = self.masks[index] return mask def __len__(self): return len(self.slicer.crops) def __getitem__(self, index): image = self._get_image(index) mask = self._get_mask(index) data = self.transform(image=image, mask=mask) return { 'features': tensor_from_rgb_image(data['image']), 'targets': tensor_from_mask_image(data['mask']).float(), 'image_id': self.image_ids[index] }
class _InrialTiledImageMaskDataset(Dataset): def __init__( self, image_fname: str, mask_fname: str, image_loader: Callable, target_loader: Callable, tile_size, tile_step, image_margin=0, transform=None, target_shape=None, need_weight_mask=False, keep_in_mem=False, make_mask_target_fn: Callable = mask_to_bce_target, ): self.image_fname = image_fname self.mask_fname = mask_fname self.image_loader = image_loader self.mask_loader = target_loader self.image = None self.mask = None self.need_weight_mask = need_weight_mask if target_shape is None or keep_in_mem: image = image_loader(image_fname) mask = target_loader(mask_fname) if image.shape[0] != mask.shape[0] or image.shape[1] != mask.shape[ 1]: raise ValueError( f"Image size {image.shape} and mask shape {image.shape} must have equal width and height" ) target_shape = image.shape self.slicer = ImageSlicer(target_shape, tile_size, tile_step, image_margin) self.transform = transform self.image_ids = [fs.id_from_fname(image_fname)] * len( self.slicer.crops) self.crop_coords_str = [ f"[{crop[0]};{crop[1]};{crop[2]};{crop[3]};]" for crop in self.slicer.crops ] self.make_mask_target_fn = make_mask_target_fn def _get_image(self, index): image = self.image_loader(self.image_fname) image = self.slicer.cut_patch(image, index) return image def _get_mask(self, index): mask = self.mask_loader(self.mask_fname) mask = self.slicer.cut_patch(mask, index) return mask def __len__(self): return len(self.slicer.crops) def __getitem__(self, index): image = self._get_image(index) mask = self._get_mask(index) data = self.transform(image=image, mask=mask) image = data["image"] mask = data["mask"] data = { INPUT_IMAGE_KEY: image_to_tensor(image), INPUT_MASK_KEY: self.make_mask_target_fn(mask), INPUT_IMAGE_ID_KEY: self.image_ids[index], "crop_coords": self.crop_coords_str[index], } if self.need_weight_mask: data[INPUT_MASK_WEIGHT_KEY] = tensor_from_mask_image( compute_weight_mask(mask)).float() return data