|
- import cv2
- import mindspore
- import mindspore.dataset as ds
- import mindspore.dataset.transforms.c_transforms as C2
- import mindspore.dataset.vision.c_transforms as C
- import pandas as pd
-
- from make_imagenet_label import make_imgpath_and_label, get_batch_csv
-
-
- class IN1K:
- def __init__(self, info_path):
- self.img_path_list = pd.read_csv(info_path)["path"]
- self.label_list = pd.read_csv(info_path)["label"]
-
- def __getitem__(self, idx):
- img = cv2.imread(self.img_path_list[idx])
- label = self.label_list[idx]
-
- return img, label
-
- def __len__(self):
- return len(self.img_path_list)
-
-
- def get_dataloader(imagenet_dir, P, K):
-
- data = make_imgpath_and_label(imagenet_dir)
- info_path = get_batch_csv(data, P, K, imagenet_dir, 3)
- # info_path = get_batch_csv(data, P, K, imagenet_dir, 3)
-
- # transform
- mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
- std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
- trans = [
- C.RandomResizedCrop(224),
- C.RandomHorizontalFlip(prob=0.5),
- C.RandomColorAdjust(brightness=0.4, contrast=0.4, saturation=0.4)
- ]
- trans += [
- C.Normalize(mean=mean, std=std),
- C.HWC2CHW(),
- ]
- type_cast_op = C2.TypeCast(mindspore.dtype.int32)
-
- # bulid dataloader
- dataset = IN1K(info_path)
- dataloader = ds.GeneratorDataset(dataset, ["image", "label"], shuffle=False)
- dataloader = dataloader.map(input_columns="image", operations=trans)
- dataloader = dataloader.map(input_columns="label", operations=type_cast_op)
-
- return dataloader
-
- if __name__ == '__main__':
- imagenet_dir = "/docker_output/ILSVRC2012_img_train/"
- A = get_dataloader(imagenet_dir, 2, 5)
- A = A.batch(30)
-
- for a in A.create_dict_iterator():
- print(a["image"].shape)
- break
|