Using External Datasets with HuggingFace Data Loader

mariosasko
Thu, 21 Apr 2022 16:26:51 +0000

Hi! Note that

ds = load_dataset('imagefolder', data_dir='./tiny-imagenet-200')

puts all the images in a singe train split. Also, this fails to infer the lables since tiny imagenet is not a valid image folder.

This is the corrected code:

import os
from datasets import DownloadConfig, load_dataset, Image, DatasetDict

PATH_TO_TINY_IMAGENET_DIR = "..."

# avoids issues with filelock + multiprocessing on some platforms
dc = DownloadConfig(num_proc=1)

# prepare train split
train_dset = load_dataset("imagefolder", download_config=dc, data_dir=os.path.join(PATH_TO_TINY_IMAGENET_DIR, "train"), split="train", ignore_verifications=True)

train_dset = train_dset.remove_columns("label")
train_dset = train_dset.cast_column("image", Image(decode=False))
train_dset = train_dset.map(lambda ex: {"label": ex["image"]["path"].replace("\\", "/").split("/")[-3]})
train_dset = train_dset.cast_column("image", Image(decode=True))

train_dset = train_dset.class_encode_column("label")

# prepare validation split
val_dset = load_dataset("imagefolder", download_config=dc, data_dir=os.path.join(PATH_TO_TINY_IMAGENET_DIR, "val", "images"), split="train", ignore_verifications=True)

annotations = {}
with open(os.path.join(PATH_TO_TINY_IMAGENET_DIR, "val", "val_annotations.txt")) as f:
    for line in f:
        if line:
            img_file, label, *_ = line.split()
            annotations[img_file.lower()] = label

val_dset = val_dset.remove_columns("label")
val_dset = val_dset.cast_column("image", Image(decode=False))
val_dset = val_dset.map(lambda ex: {"label": train_dset.features["label"].str2int(annotations[os.path.basename(ex["image"]["path"]).lower()])}, features=train_dset.features)
val_dset = val_dset.cast_column("image", Image(decode=True))

# prepare test split
test_dset = load_dataset("imagefolder", download_config=dc, data_dir=os.path.join(PATH_TO_TINY_IMAGENET_DIR, "test", "images"), split="train", ignore_verifications=True)
test_dset = test_dset.remove_columns("label")
test_dset = test_dset.map(lambda ex: {"label": None}, features=train_dset.features)

# final dataset with all the splits
ds = DatasetDict({"train": train_dset, "validation": val_dset, "test": test_dset})