Shortcuts

Source code for hfai.datasets.imagenet21k

from typing import Callable, Optional

import pickle
from PIL import Image
from io import BytesIO
from ffrecord import FileReader
from .base import (
    BaseDataset,
    get_data_dir,
    register_dataset
)


"""
Expected file organization:

    [data_dir]
        train.ffr
            PART_00000.ffr
            PART_00001.ffr
            ...
        val.ffr
            PART_00000.ffr
            PART_00001.ffr
            ...
"""


[docs]@register_dataset class ImageNet21K(BaseDataset): """ 这是一个图像识别数据集 更多信息参考:https://github.com/Alibaba-MIIL/ImageNet21K Args: split (str): 数据集划分形式,包括:完整集(``full``)、训练集(``train``)或者验证集(``val``) transform (Callable): transform 函数,对图片进行 transfrom,接受一张图片作为输入,输出 transform 之后的图片 check_data (bool): 是否对每一条样本检验校验和(默认为 ``True``) miniset (bool): 是否使用 mini 集合(默认为 ``False``) Returns: pic, name, label (PIL.Image.Image, str, int): 返回的每个样本是一个元组,包括一个RGB格式的图片,图片名,以及代表这张图片类别的标签 Examples: .. code-block:: python from hfai.datasets import ImageNet21K from torchvision import transforms transform = transforms.Compose([ transforms.Resize(224), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std), ]) dataset = ImageNet21K(split, transform) loader = dataset.loader(batch_size=64, num_workers=4) for pic, name, label in loader: # training model """ def __init__( self, split: str, transform: Optional[Callable] = None, check_data: bool = True, miniset: bool = False ) -> None: super(ImageNet21K, self).__init__() assert split in ["full", "train", "val"] self.split = split self.transform = transform data_dir = get_data_dir() if miniset: data_dir = data_dir / "mini" self.data_dir = data_dir / "ImageNet21K" self.fname = self.data_dir / f"{split}.ffr" self.reader = FileReader(self.fname, check_data) def __len__(self): return self.reader.n def __getitem__(self, indices): imgs_bytes = self.reader.read(indices) samples = [] for i, bytes_ in enumerate(imgs_bytes): img_bytes, name, label = pickle.loads(bytes_) try: img = Image.open(BytesIO(img_bytes)).convert("RGB") except Exception as e: continue if self.transform: img = self.transform(img) samples.append((img, name, int(label))) return samples