Fix AFHQ Dataset

[omar-abdelgawad]

Describe the bug
The current AFHQ dataset implementation that has been merged needs correction as the code inside is for colorizatoin datasets and doesn't work with the afhq dataset format.

Expected behavior

• Please refer to the following link for the dataset structure https://www.kaggle.com/datasets/andrewmvd/animal-faces.
• Also please refer to the original implementation of TUNIT to see how they implemented the dataset class https://github.com/clovaai/tunit.

Screenshots
Just for reference the following code in src/tunit/dataset.py is the code to be changed. Please note that the method to fully change is the getitem only.

class AFHQCatDataset(Dataset):
    """Dataset class for AFHQ Cat Colorization dataset.

    Args:
        root_dir (root_dir): Path for dataset dir.
    """

    def __init__(self, root_dir: Path | str) -> None:
        self.root_dir = root_dir
        self.list_files = os.listdir(self.root_dir)[: cfg.NUM_IMAGES_DATASET]
        print(f"The length of the dataset is: {len(self.list_files)}")

    def __len__(self):
        return len(self.list_files)

    def __getitem__(self, index):
        img_file_name = self.list_files[index]
        img_path = os.path.join(self.root_dir, img_file_name)
        image = np.array(Image.open(img_path))
        target_image = image[:, : image.shape[1] // 2, :]
        input_image = image[:, image.shape[1] // 2 :, :]

        augmentations = cfg.both_transform(image=input_image, image0=target_image)
        input_image, target_image = augmentations["image"], augmentations["image0"]

        input_image = cfg.transform_only_input(image=input_image)["image"]
        target_image = cfg.transform_only_mask(image=target_image)["image"]

        return input_image, target_image

Additional context
For additional information please ask or contact me directly whichever is easier for you.