<jupyter_start><jupyter_text>A custom dataset should inherit `Dataset` and override `__len__`,`__getitem__` methods. We use `__len__` to get the size of the data and `__getitem__` to support the indexing such that `dataset[i]` can be used to get $i_{th}$ example.<jupyter_code>class FaceLandmarksDataset(Dataset):
  """ Face Landmarks Dataset """

  def __init__(self, csv_file, root_dir, transform = None):
    """
    Args:
      csv_file (string): Path to the csv file with annotations.
      root_dir (string): Directory with all the images.
      transform (callable, optional): Optional transform to be applied on
        a sample.
    """
    self.landmarks_frame = pd.read_csv(csv_file)
    self.root_dir = root_dir
    self.transform = transform
  
  def __len__(self):
    return len(self.landmarks_frame)
  
  def __getitem__(self, idx):
      if torch.is_tensor(idx):
        idx = idx.tolist()
      
      img_name = os.path.join(self.root_dir,
                              self.landmarks_frame.iloc[idx, 0])
      image = io.imread(img_name)
      landmarks = self.landmarks_frame.iloc[idx, 1:]
      landmarks = np.array([landmarks])
      landmarks = landmarks.astype(np.float32).reshape(-1,2)
      sample = {'image': image, 'landmarks': landmarks}

      if self.transform:
        sample = self.transform(sample)
      
      return sample
face_dataset = FaceLandmarksDataset(csv_file="faces/face_landmarks.csv",
                                    root_dir='faces/')

fig = plt.figure()

for i in range(len(face_dataset)):
  sample = face_dataset[i]

  print(i, sample['image'].shape, sample['landmarks'].shape)

  ax = plt.subplot(1, 4, i + 1)
  plt.tight_layout()
  ax.set_title('Sample #{}'.format(i))
  ax.axis('off')
  show_landmarks(**sample)

  if i == 3:
    plt.show()
    break

# create custom transformers
class Rescale(object):
  """
  Rescale the image in a sample to a given size
  Args:
    output_size (tuple or int): Desired output size. If tuple,
      output is matched to output_size. If int, smaller of image edges
      is matched to output_size keeping aspect ratio the same.
  
  """
  def __init__(self, output_size):
    assert isinstance(output_size, (int,tuple))
    self.output_size = output_size
  
  def __call__(self,sample):
    image, landmarks = sample['image'], sample['landmarks']

    h, w = image.shape[:2]
    if isinstance(self.output_size,int):
      if h>w:
        new_h, new_w = self.output_size * h / w, self.output_size
      else:
        new_h, new_w = self.output_size, self.output_size * w / h
    else:
      new_h, new_w = self.output_size
    
    new_h, new_w = int(new_h), int(new_w)

    img = transform.resize(image, (new_h, new_w))

    # h and w are swapper for landmarks because of images.
    # x and y axes are axis 1 and 0 respectively
    landmarks = landmarks * [new_w / w, new_h / h]
    
    return {'image':img, 'landmarks': landmarks}

class RandomCrop(object):
  """
  Crop randomly the image in a sample
  Doesn't crop the parts with face landmarks.
  """

  def __call__(self, sample):
    image, landmarks = sample['image'], sample['landmarks']
    
    lm_min_x = sample['landmarks'][:,0].min()
    lm_max_x = sample['landmarks'][:,0].max()
    lm_min_y = sample['landmarks'][:,1].min()
    lm_max_y = sample['landmarks'][:,1].max()

    h, w = image.shape[:2]

    top = np.random.randint(lm_max_y+5, h)
    down = np.random.randint(0, lm_min_y-5)
    left = np.random.randint(0, lm_min_x-5)
    right = np.random.randint(lm_max_x+5, w)


    image = image[down: top,
                  left: right]

    landmarks = landmarks - [left,down]

    return {'image': image, 'landmarks': landmarks}

class ToTensor(object):
  """ Convert ndarrays in sample to Tensors """

  def __call__(self,sample):
    image, landmarks = sample['image'], sample['landmarks']

    # swap color axis because
    # numpy image: H x W x C
    # torch image: C x H x W
    image = image.transpose((2,0,1))

    return {'image': torch.from_numpy(image),
            'landmarks': torch.from_numpy(landmarks)}
scale = Rescale(256)
crop = RandomCrop()
composed = transforms.Compose([RandomCrop(),
                               Rescale((128,128))])

# Apply each of the above transforms on sample
fig = plt.figure()
sample = face_dataset[r]
for i, tsfrm in enumerate([scale, crop, composed]):
  transformed_sample = tsfrm(sample)

  ax = plt.subplot(1, 3, i + 1)
  plt.tight_layout()
  ax.set_title(type(tsfrm).__name__)
  show_landmarks(**transformed_sample)

plt.show()
transformed_dataset = FaceLandmarksDataset(csv_file = 'faces/face_landmarks.csv',
                                           root_dir='faces/',
                                           transform=transforms.Compose([
                                            RandomCrop(),
                                            Rescale((124,124)),
                                            ToTensor()
                                           ]))
for i in range(len(transformed_dataset)):
  sample = transformed_dataset[i]

  print(i, sample['image'].size(), sample['landmarks'].size())

  if i == 3:
    break
data_loader = DataLoader(transformed_dataset, batch_size = 4,
                         shuffle = True, num_workers = 4)

# Helper function to show a batch
def show_landmarks_batch(sample_batched):
  """ Show image with landmarks for a batch of samples. """
  images_batch, landmarks_batch = sample_batched['image'], sample_batched['landmarks']
  batch_size = len(images_batch)
  im_size = images_batch.size(2)
  grid_border_size = 2

  grid = utils.make_grid(images_batch)
  plt.imshow(grid.numpy().transpose((1,2,0)))

  for i in range(batch_size):
    plt.scatter(landmarks_batch[i, :, 0].numpy() + i * im_size + (i + 1) * grid_border_size,
                landmarks_batch[i, :, 1].numpy() + grid_border_size,
                s = 10, marker = '.', c = 'r')
    plt.title('Batch from Data Loader')
  
for i_batch, sample_batched in enumerate(data_loader):
  print(i_batch, sample_batched['image'].size(), sample_batched['landmarks'].size())

  # observe a batch and stop
  if i_batch == 1:
    plt.figure()
    show_landmarks_batch(sample_batched)
    plt.axis('off')
    plt.ioff()
    plt.show()
    break<jupyter_output>0 torch.Size([4, 3, 124, 124]) torch.Size([4, 68, 2])
1 torch.Size([4, 3, 124, 124]) torch.Size([4, 68, 2])
{'fraction_non_alphanumeric': 0.12711598746081504, 'fraction_numerical': 0.013322884012539185, 'mean_word_length': 2.7869436201780418, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 13, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '305119', 'license': 'MIT', 'n_tokens_mistral': 2156, 'n_tokens_neox': 2042, 'n_words': 519, 'path': 'PyTorch/Writing Custom Datasets, DataLoaders and Transforms.ipynb', 'repo_name': 'Raqhea/DL-DS-ML', 'size': '6380'}