text: import os as _os
  import random as _random
  import warnings as _warn
  from abc import ABC as _ABC
  
  import numpy as _np
  import torchvision.transforms as _tmf
  from easytorch.data import ETDataset as _ETDataset
  import easytorch.vision.imageutils as _imgutils
  
  sep = _os.sep
  
  
  class BaseImageDataset(_ETDataset, _ABC):
      def __init__(self, **kw):
          super().__init__(**kw)
          self.transforms = self.get_transforms()
          self.pil_to_tensor = _tmf.Compose([_tmf.ToPILImage(), _tmf.ToTensor()])
  
      def get_transforms(self):
          return _tmf.Compose([_tmf.ToPILImage(), _tmf.ToTensor()])
  
      def _validate_image_data(self, dspec, img_obj):
          return img_obj
  
      def load_img(self, dspec, file):
          img_obj = _imgutils.Image()
  
          """Load Image"""
          img_obj.load(dspec['data_dir'], file)
          img_obj.apply_clahe()
          if len(img_obj.array.shape) > 2 and img_obj.array.shape[-1] != 3:
              _warn.warn(f"Suspicious Image shape: {img_obj.array.shape}, clipping to RGB: {dspec['name']}")
              img_obj.array = img_obj.array[:, :, :3]
  
          """Load ground truth"""
          dspec['has_gt'] = any(['label_dir' in dspec.keys()])
          if dspec['has_gt']:
              img_obj.load_ground_truth(dspec["label_dir"], dspec["label_getter"])
          else:
              _warn.warn(f"Random label initialized: {dspec['name']}")
              img_obj.ground_truth = _np.random.randint(0, 2, img_obj.array.shape[:2]).astype(_np.uint8) * 255
  
          """Load mask"""
          dspec['has_mask'] = any(['mask_dir' in dspec.keys()])
          if dspec['has_mask']:
              img_obj.load_mask(dspec['mask_dir'], dspec['mask_getter'])
  
          return self._validate_image_data(dspec, img_obj)
  
      def __getitem__(self, index):
          raise NotImplementedError('Must implement')
  
  
  class PatchedImgDataset(BaseImageDataset, _ABC):
      """dataspec must have patch_shape, patch_offset, """
  
      def load_index(self, dataset_name, file):
          r"""
          :param dataset_name: name of teh dataset as provided in train_dataspecs
          :param file: Name of an image
          :return:
          Logic split an image to patches and feed to U-Net. Meanwhile we need to store the four-corners
              of each patch so that we can rejoin the full image from the patches' corresponding predictions.
          """
          dt = self.dataspecs[dataset_name]
          obj = self.load_img(dt, file)
  
          cache_key = self.diskcache.add(f"{dataset_name}_{file}", obj)
          for corners in _imgutils.get_chunk_indexes(
                  obj.array.shape[:2],
                  dt['patch_shape'],
                  dt['patch_offset']
          ):
              """
              get_chunk_indexes will return the list of four corners of all patches of the images  
              by using window size of self.patch_shape, and offset  of elf.patch_offset
              """
              self.indices.append([dataset_name, file] + corners + [cache_key])
  
  
  class BinaryPatchDataset(PatchedImgDataset):
      def __init__(self, **kw):
          super().__init__(**kw)
  
      def _validate_image_data(self, dspec, img_obj):
          thr_manual = dspec.setdefault('thr_manual', 50)
          if dspec.get('has_gt'):
              gt_unique = _np.unique(img_obj.ground_truth)
              if len(img_obj.ground_truth.shape) > 2:
                  _warn.warn(
                      f"Ground truth shape suspicious: {img_obj.ground_truth.shape}, "
                      f"using 1st channel onl:  {dspec['name']}"
                  )
                  img_obj.ground_truth = img_obj.ground_truth[:, :, 0]
  
              if sum(gt_unique) == 1:
                  _warn.warn(f"Ground truth 1 converted to 255")
                  img_obj.ground_truth[img_obj.ground_truth == 1] = 255
  
              if len(gt_unique) != 2:
                  _warn.warn(
                      f"Number of unique ground truth items != {2} in  {dspec['name']}. "
                      f"\nBinarizing ... {gt_unique}: "
                  )
                  _imgutils.binarize(img_obj.ground_truth, thr_manual)
  
          if dspec.get('has_mask'):
              mask_unique = _np.unique(img_obj.mask)
              if len(img_obj.mask.shape) > 2:
                  _warn.warn(f"Mask shape suspicious: {img_obj.mask.shape}, using 1st channel only: {dspec['name']}")
                  img_obj.mask = img_obj.mask[:, :, 0]
  
              if sum(mask_unique) == 1:
                  _warn.warn(f"Mask truth 1 converted to 255")
                  img_obj.mask[img_obj.mask == 1] = 255
  
              if len(mask_unique) != 2:
                  _warn.warn(
                      f"Number of unique items in mask: {mask_unique} in {dspec['name']}"
                  )
  
          if dspec.get('bbox_crop'):
              copy = img_obj.copy()
              img_obj.array, img_obj.ground_truth, img_obj.mask = _imgutils.masked_bboxcrop(
                  img_obj.array,
                  img_obj.ground_truth,
                  offset=dspec.setdefault('bbox_crop_offset', 51),
                  threshold=dspec.setdefault('bbox_crop_threshold', 5)
              )
  
              if img_obj.array.shape[0] < dspec['patch_shape'][0] or img_obj.array.shape[1] < dspec['patch_shape'][1]:
                  _warn.warn(
                      f"BBOX crop reversing for "
                      f"{dspec['name']}:{img_obj.file}, shape: {img_obj.array.shape}, {copy.array.shape}"
  
                  )
                  img_obj = copy.copy()
  
          if dspec.get('resize'):
              img_obj.array = _imgutils.resize(img_obj.array, dspec['resize'])
              img_obj.ground_truth = _imgutils.resize(img_obj.ground_truth, dspec['resize'])
  
              if img_obj.mask is not None:
                  img_obj.mask = _imgutils.resize(img_obj.mask, dspec['resize'])
  
          """Must binarize after resize"""
          if dspec.get('resize'):
              _imgutils.binarize(img_obj.ground_truth, thr_manual)
              if dspec.get('has_mask'):
                  _imgutils.binarize(img_obj.mask, thr_manual)
          return img_obj
  
      def __getitem__(self, index):
          """
          :param index:
          :return: dict with keys - indices, input, label
              We need indices to get the file name to save the respective predictions.
          """
          dname, file, row_from, row_to, col_from, col_to, cache_key = self.indices[index]
  
          obj = self.diskcache.get(cache_key)
          img = obj.array
          gt = obj.ground_truth[row_from:row_to, col_from:col_to]
  
          p, q, r, s, pad = _imgutils.expand_and_mirror_patch(
              img.shape,
              [row_from, row_to, col_from, col_to],
              self.dataspecs[dname]['expand_by']
          )
          if len(img.shape) == 3:
              pad = [*pad, (0, 0)]
  
          img = _np.pad(img[p:q, r:s], pad, 'reflect')
          if self.mode == 'train' and _random.uniform(0, 1) <= 0.5:
              img = _np.flip(img, 0)
              gt = _np.flip(gt, 0)
  
          if self.mode == 'train' and _random.uniform(0, 1) <= 0.5:
              img = _np.flip(img, 1)
              gt = _np.flip(gt, 1)
  
          img = self.transforms(img)
          gt = self.pil_to_tensor(gt)
          return {'indices': self.indices[index], 'input': img, 'label': gt.squeeze()}
  
  
  class FullImgDataset(BaseImageDataset):
      def __init__(self, **kw):
          r"""
          Initialize necessary shapes for unet.
          """
          super().__init__(**kw)
          self.labels = None
  
      def get_transforms(self):
          if self.mode == "train":
              return _tmf.Compose(
                  [_tmf.ToPILImage(),
                   _tmf.RandomHorizontalFlip(),
                   _tmf.RandomVerticalFlip(),
                   _tmf.ToTensor()]
              )
          return self.pil_to_tensor
  
      def _load_labels(self, dspec):
          return None
  
      def _get_label(self, file):
          pass
  
      def _validate_image_data(self, dspec, img_obj):
  
          if dspec.get('bbox_crop'):
              copy = img_obj.copy()
              img_obj.array, img_obj.mask = _imgutils.masked_bboxcrop(
                  img_obj.array,
                  offset=dspec.setdefault('bbox_crop_offset', 31),
                  threshold=dspec.setdefault('bbox_crop_threshold', 5)
              )
  
              if dspec.get('resize') and (
                      img_obj.array.shape[0] < 0.55 * copy.array.shape[0]
                      or
                      img_obj.array.shape[1] < 0.55 * copy.array.shape[1]
              ):
                  _warn.warn(f"BBOX crop reversing for {dspec['name']}:{img_obj.file}, shape: {img_obj.array.shape}")
                  img_obj = copy.copy()
  
          if dspec.get('resize'):
              img_obj.array = _imgutils.resize(img_obj.array, dspec['resize'])
  
              if img_obj.mask is not None:
                  img_obj.mask = _imgutils.resize(img_obj.mask, dspec['resize'])
  
          return img_obj
  
      def load_img(self, dspec, file):
          img_obj = _imgutils.Image()
  
          """Load Image"""
          img_obj.load(dspec['data_dir'], file)
          img_obj.apply_clahe()
          if len(img_obj.array.shape) > 2 and img_obj.array.shape[-1] != 3:
              _warn.warn(f"Suspicious Image shape: {img_obj.array.shape}, clipping to RGB: {dspec['name']}")
              img_obj.array = img_obj.array[:, :, :3]
  
          """Load mask"""
          dspec['has_mask'] = any(['mask_dir' in dspec.keys()])
          if dspec['has_mask']:
              img_obj.load_mask(dspec['mask_dir'], dspec['mask_getter'])
  
          return self._validate_image_data(dspec, img_obj)
  
      def load_index(self, dataset_name, file):
          dspec = self.dataspecs[dataset_name]
          if self.labels is None:
              self.labels = self._load_labels(dspec)
  
          img_obj = self.load_img(dspec, file)
          cache_key = self.diskcache.add(f"{dataset_name}_{file}", img_obj)
          self.indices.append([dataset_name, file, self._get_label(file), cache_key])
  
      def __getitem__(self, index):
          _dname, file, label, cache_key = self.indices[index]
          obj = self.diskcache.get(cache_key)
          img = obj.array
  
          img = self.transforms(img)
          return {'indices': self.indices[index], 'input': img, 'label': _np.array(label)}
  