|
- # -*- coding: utf-8 -*-
- """
- @author: xufan
- @License: (C)Copyright 2021, xufan
- @CreateTime: 2021/11/20 17:09:05
- @Filename: dataset
- service api views
- """
-
- import numpy as np
- import random
- import glob
- import h5py
- from PIL import Image
- import mindspore.dataset as ds
- import mindspore.dataset.vision.py_transforms as py_vision
- import os
- from mytransformers import get_transforms
- from mindspore import context
-
- data_transforms = get_transforms({
- 'scale': 256,
- 'horizontal_flip': True,
- 'random_affine': 0.3,
- 'random_rotation': (0, 360),
- 'random_crop': 256,
- })
-
-
- class MoNuSegGenerator:
- def __init__(self, image_dir, target_dir, shuffle=True):
- self.image_list = []
- self.target_list = []
- self.data_transform = data_transforms
- self.toTensor = py_vision.ToTensor()
- _image_path_list = glob.glob(os.path.join(image_dir, '*.png'))
- if shuffle:
- random.shuffle(_image_path_list)
- for _image_path in _image_path_list:
- base_name = os.path.basename(_image_path).replace('.png', '')
- _target_path = os.path.join(target_dir, '%s.npy' % (base_name))
- if os.path.exists(_target_path):
- self.image_list.append(_image_path)
- self.target_list.append(_target_path)
-
- def __getitem__(self, index):
- image_path = self.image_list[index]
- _image = Image.open(self.image_list[index]).convert('RGB') # (640, 640, 3)
- mask = np.zeros(shape=(_image.size[0], _image.size[1]), dtype=np.uint8)
- _target_copy = (np.load(self.target_list[index]) * 256).astype(np.uint8)
- _target = Image.fromarray(np.stack((_target_copy[0], _target_copy[1], mask), axis=2))
- sample = [_image, _target]
- if self.data_transform is not None:
- sample = self.data_transform(sample)
-
- image = self.toTensor(sample[0])
- target = self.toTensor(sample[1])[0:2]
- return [image, target, image_path]
-
- def __len__(self):
- return len(self.image_list)
-
-
- def create_dataset(image_dir, target_dir, column_names, batch_size, shuffle=True):
- monuseg = MoNuSegGenerator(image_dir=image_dir, target_dir=target_dir, shuffle=shuffle)
- dataset = ds.GeneratorDataset(monuseg, column_names=column_names)
- train_set = dataset.batch(batch_size)
- return train_set
-
-
- if __name__ == '__main__':
- train_img_dir = os.path.join("BCData", "images/train")
- train_target_dir = os.path.join("BCData", "npy_data2/train")
- MoNuSegGenerator(image_dir=train_img_dir, target_dir=train_target_dir, ).__getitem__(0)
- train_set = create_dataset(image_dir=train_img_dir, target_dir=train_target_dir, column_names=["input", "target"],
- batch_size=2)
-
- for data in train_set:
- image, target = data
- print(image.shape)
|