text: # Copyright 2021 MosaicML. All Rights Reserved.
  
  import os
  from dataclasses import dataclass
  from math import ceil
  from typing import Optional, Tuple, Union
  
  import numpy as np
  import torch
  import torchvision.transforms.functional as TF
  import yahp as hp
  from PIL import Image
  from torch.utils.data import Dataset
  from torchvision import transforms
  
  from composer.core.types import DataSpec
  from composer.datasets.hparams import DatasetHparams, SyntheticHparamsMixin
  from composer.datasets.imagenet import IMAGENET_CHANNEL_MEAN, IMAGENET_CHANNEL_STD
  from composer.datasets.synthetic import SyntheticBatchPairDataset
  from composer.datasets.utils import NormalizationFn, pil_image_collate
  from composer.utils import dist
  
  
  class RandomResizePair(torch.nn.Module):
      """Resize the image and target to `base_size` scaled by a randomly sampled value.
  
      Args:
          min_scale (float): the minimum value the samples can be rescaled.
          max_scale (float): the maximum value the samples can be rescaled.
          base_size (Tuple[int, int]): a specified base size (height x width) to scale to get the resized dimensions.
              When this is None (default), use the input image size.
      """
  
      def __init__(self, min_scale: float, max_scale: float, base_size: Optional[Tuple[int, int]] = None):
          super().__init__()
          self.min_scale = min_scale
          self.max_scale = max_scale
          self.base_size = base_size
  
      def forward(self, sample: Tuple[Image.Image, Image.Image]):
          image, target = sample
          resize_scale = np.random.random_sample() * (self.max_scale - self.min_scale) + self.min_scale
          base_height, base_width = self.base_size if self.base_size else (image.height, image.width)
          resized_dims = (int(base_height * resize_scale), int(base_width * resize_scale))
          resized_image = TF.resize(image, resized_dims, interpolation=TF.InterpolationMode.BILINEAR)  # type: ignore
          resized_target = TF.resize(target, resized_dims, interpolation=TF.InterpolationMode.NEAREST)  # type: ignore
          return resized_image, resized_target
  
  
  class RandomCropPair(torch.nn.Module):
      """Crop the image and target at a randomly sampled position.
  
      Args:
          crop_size (Tuple[int, int]): the size (height x width) of the crop.
      """
  
      def __init__(self, crop_size: Tuple[int, int]):
          super().__init__()
          self.crop_size = crop_size
  
      def forward(self, sample: Tuple[Image.Image, Image.Image]):
          image, target = sample
          if image.height > self.crop_size[0] or image.width > self.crop_size[1]:
              i, j, h, w = transforms.RandomCrop.get_params(
                  image, output_size=self.crop_size)  # type: ignore - transform typing does not include PIL.Image
              image = TF.crop(image, i, j, h, w)  # type: ignore - transform typing does not include PIL.Image
              target = TF.crop(target, i, j, h, w)  # type: ignore - transform typing does not include PIL.Image
          return image, target
  
  
  class RandomHFlipPair(torch.nn.Module):
      """Flip the image and target horizontally with a specified probability.
  
      Args:
          probability (float): the probability of flipping the image and target. Default is 0.5.
      """
  
      def __init__(self, probability: float = 0.5):
          super().__init__()
          self.probability = probability
  
      def forward(self, sample: Tuple[Image.Image, Image.Image]):
          image, target = sample
          if np.random.random_sample() > self.probability:
              image = TF.hflip(image)  # type: ignore - transform typing does not include PIL.Image
              target = TF.hflip(target)  # type: ignore - transform typing does not include PIL.Image
          return image, target
  
  
  class PadToSize(torch.nn.Module):
      """Pad an image to a specified size.
  
      Args:
          size (Tuple[int, int]): the size (height x width) of the image after padding.
          fill (Union[int, Tuple[int, int, int]]): the value to use for the padded pixels. Default is 0.
      """
  
      def __init__(self, size: Tuple[int, int], fill: Union[int, Tuple[int, int, int]] = 0):
          super().__init__()
          self.size = size
          self.fill = fill
  
      def forward(self, image: Image.Image):
          padding = max(self.size[0] - image.height, 0), max(self.size[1] - image.width, 0)
          padding = (padding[1] // 2, padding[0] // 2, ceil(padding[1] / 2), ceil(padding[0] / 2))
          image = TF.pad(image, padding, fill=self.fill)  # type: ignore - transform typing does not include PIL.Image
          return image
  
  
  class PhotometricDistoration(torch.nn.Module):
      """Applies a combination of brightness, contrast, saturation, and hue jitters with random intensity.
  
      This is a less severe form of PyTorch's ColorJitter used by the mmsegmentation library here:
      https://github.com/open-mmlab/mmsegmentation/blob/master/mmseg/datasets/pipelines/transforms.py#L837
  
      Args:
          brightness (float): max and min to jitter brightness.
          contrast (float): max and min to jitter contrast.
          saturation (float): max and min to jitter saturation.
          hue (float): max and min to jitter hue.
      """
  
      def __init__(self, brightness: float, contrast: float, saturation: float, hue: float):
          super().__init__()
          self.brightness = brightness
          self.contrast = contrast
          self.saturation = saturation
          self.hue = hue
  
      def forward(self, image: Image.Image):
          if np.random.randint(2):
              brightness_factor = np.random.uniform(1 - self.brightness, 1 + self.brightness)
              image = TF.adjust_brightness(
                  image, brightness_factor)  # type: ignore - transform typing does not include PIL.Image
  
          contrast_mode = np.random.randint(2)
          if contrast_mode == 1 and np.random.randint(2):
              contrast_factor = np.random.uniform(1 - self.contrast, 1 + self.contrast)
              image = TF.adjust_contrast(
                  image,  # type: ignore - transform typing does not include PIL.Image
                  contrast_factor)
  
          if np.random.randint(2):
              saturation_factor = np.random.uniform(1 - self.saturation, 1 + self.saturation)
              image = TF.adjust_saturation(
                  image, saturation_factor)  # type: ignore - transform typing does not include PIL.Image
  
          if np.random.randint(2):
              hue_factor = np.random.uniform(-self.hue, self.hue)
              image = TF.adjust_hue(image, hue_factor)  # type: ignore - transform typing does not include PIL.Image
  
          if contrast_mode == 0 and np.random.randint(2):
              contrast_factor = np.random.uniform(1 - self.contrast, 1 + self.contrast)
              image = TF.adjust_contrast(
                  image,  # type: ignore - transform typing does not include PIL.Image
                  contrast_factor)
  
          return image
  
  
  class ADE20k(Dataset):
      """PyTorch Dataset for ADE20k.
  
      Args:
          datadir (str): the path to the ADE20k folder.
          split (str): the dataset split to use, either 'train', 'val', or 'test'. Default is 'train'.
          both_transforms (torch.nn.Module): transformations to apply to the image and target simultaneously.
              Default is None.
          image_transforms (torch.nn.Module): transformations to apply to the image only. Default is None.
          target_transforms (torch.nn.Module): transformations to apply to the target only. Default is None.
      """
  
      def __init__(self,
                   datadir: str,
                   split: str = 'train',
                   both_transforms: Optional[torch.nn.Module] = None,
                   image_transforms: Optional[torch.nn.Module] = None,
                   target_transforms: Optional[torch.nn.Module] = None):
          super().__init__()
          self.datadir = datadir
          self.split = split
          self.both_transforms = both_transforms
          self.image_transforms = image_transforms
          self.target_transforms = target_transforms
  
          # Check datadir value
          if self.datadir is None:
              raise ValueError("datadir must be specified")
          elif not os.path.exists(self.datadir):
              raise FileNotFoundError(f"datadir path does not exist: {self.datadir}")
  
          # Check split value
          if self.split not in ["train", "val", "test"]:
              raise ValueError(f"split must be one of [`train`, `val`, `test`] but is: {self.split}")
  
          self.image_dir = os.path.join(self.datadir, 'images', self.split)
          if not os.path.exists(self.image_dir):
              raise FileNotFoundError(f"ADE20k directory structure is not as expected: {self.image_dir} does not exist")
  
          self.image_files = os.listdir(self.image_dir)
  
          # Filter for ADE files
          self.image_files = [f for f in self.image_files if f[:3] == 'ADE']
  
          # Remove grayscale samples
          if self.split == 'train':
              corrupted_samples = ['00003020', '00001701', '00013508', '00008455']
              for sample in corrupted_samples:
                  sample_file = f'ADE_train_{sample}.jpg'
                  if sample_file in self.image_files:
                      self.image_files.remove(sample_file)
  
      def __getitem__(self, index):
          # Load image
          image_file = self.image_files[index]
          image_path = os.path.join(self.image_dir, image_file)
          image = Image.open(image_path)
  
          # Load annotation target if using either train or val splits
          if self.split in ['train', 'val']:
              target_path = os.path.join(self.datadir, 'annotations', self.split, image_file.split('.')[0] + '.png')
              target = Image.open(target_path)
  
              if self.both_transforms:
                  image, target = self.both_transforms((image, target))
  
              if self.target_transforms:
                  target = self.target_transforms(target)
  
          if self.image_transforms:
              image = self.image_transforms(image)
  
          if self.split in ['train', 'val']:
              return image, target  # type: ignore
          else:
              return image
  
      def __len__(self):
          return len(self.image_files)
  
  
  @dataclass
  class ADE20kDatasetHparams(DatasetHparams, SyntheticHparamsMixin):
      """Defines an instance of the ADE20k dataset for semantic segmentation.
  
      Args:
          split (str): the dataset split to use either 'train', 'val', or 'test'. Default is `train`.
          base_size (int): initial size of the image and target before other augmentations. Default is 512.
          min_resize_scale (float): the minimum value the samples can be rescaled. Default is 0.5.
          max_resize_scale (float): the maximum value the samples can be rescaled. Default is 2.0.
          final_size (int): the final size of the image and target. Default is 512.
          ignore_background (bool): if true, ignore the background class when calculating the training loss.
              Default is true.
      """
  
      split: str = hp.optional("Which split of the dataset to use. Either ['train', 'val', 'test']", default='train')
      base_size: int = hp.optional("Initial size of the image and target before other augmentations", default=512)
      min_resize_scale: float = hp.optional("Minimum value that the image and target can be scaled", default=0.5)
      max_resize_scale: float = hp.optional("Maximum value that the image and target can be scaled", default=2.0)
      final_size: int = hp.optional("Final size of the image and target", default=512)
      ignore_background: bool = hp.optional("If true, ignore the background class in training loss", default=True)
  
      def validate(self):
          if self.datadir is None and not self.use_synthetic:
              raise ValueError("datadir must specify the path to the ADE20k dataset.")
  
          if self.split not in ['train', 'val', 'test']:
              raise ValueError(f"split value {self.split} must be one of ['train', 'val', 'test'].")
  
          if self.base_size <= 0:
              raise ValueError("base_size cannot be zero or negative.")
  
          if self.min_resize_scale <= 0:
              raise ValueError("min_resize_scale cannot be zero or negative")
  
          if self.max_resize_scale < self.min_resize_scale:
              raise ValueError("max_resize_scale cannot be less than min_resize_scale")
  
      def initialize_object(self, batch_size, dataloader_hparams) -> DataSpec:
          self.validate()
  
          if self.use_synthetic:
              if self.split == 'train':
                  total_dataset_size = 20_206
              elif self.split == 'val':
                  total_dataset_size = 2_000
              else:
                  total_dataset_size = 3_352
  
              dataset = SyntheticBatchPairDataset(
                  total_dataset_size=total_dataset_size,
                  data_shape=[3, self.final_size, self.final_size],
                  label_shape=[self.final_size, self.final_size],
                  num_classes=150,
                  num_unique_samples_to_create=self.synthetic_num_unique_samples,
                  device=self.synthetic_device,
                  memory_format=self.synthetic_memory_format,
              )
              collate_fn = None
              device_transform_fn = None
  
          else:
              # Define data transformations based on data split
              if self.split == 'train':
                  both_transforms = torch.nn.Sequential(
                      RandomResizePair(min_scale=self.min_resize_scale,
                                       max_scale=self.max_resize_scale,
                                       base_size=(self.base_size, self.base_size)),
                      RandomCropPair(crop_size=(self.final_size, self.final_size)),
                      RandomHFlipPair(),
                  )
  
                  # Photometric distoration values come from mmsegmentation: https://github.com/open-mmlab/mmsegmentation/blob/master/mmseg/datasets/pipelines/transforms.py#L837
                  r_mean, g_mean, b_mean = IMAGENET_CHANNEL_MEAN
                  image_transforms = torch.nn.Sequential(
                      PhotometricDistoration(brightness=32. / 255, contrast=0.5, saturation=0.5, hue=18. / 255),
                      PadToSize(size=(self.final_size, self.final_size), fill=(int(r_mean), int(g_mean), int(b_mean))))
  
                  target_transforms = PadToSize(size=(self.final_size, self.final_size), fill=0)
              else:
                  both_transforms = None
                  image_transforms = transforms.Resize(size=(self.final_size, self.final_size),
                                                       interpolation=TF.InterpolationMode.BILINEAR)
                  target_transforms = transforms.Resize(size=(self.final_size, self.final_size),
                                                        interpolation=TF.InterpolationMode.NEAREST)
              collate_fn = pil_image_collate
              device_transform_fn = NormalizationFn(mean=IMAGENET_CHANNEL_MEAN,
                                                    std=IMAGENET_CHANNEL_STD,
                                                    ignore_background=self.ignore_background)
  
              # Add check to avoid type ignore below
              if self.datadir is None:
                  raise ValueError("datadir must specify the path to the ADE20k dataset.")
  
              dataset = ADE20k(datadir=self.datadir,
                               split=self.split,
                               both_transforms=both_transforms,
                               image_transforms=image_transforms,
                               target_transforms=target_transforms)
          sampler = dist.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
          return DataSpec(dataloader=dataloader_hparams.initialize_object(dataset=dataset,
                                                                          batch_size=batch_size,
                                                                          sampler=sampler,
                                                                          collate_fn=collate_fn,
                                                                          drop_last=self.drop_last),
                          device_transforms=device_transform_fn)
  