-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataset.py
63 lines (51 loc) · 2.2 KB
/
dataset.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import random
import glob
import io
import numpy as np
import PIL.Image as pil_image
#import tensorflow as tf
import mindspore as mindspore
from mindspore.dataset.transforms.py_transforms import Compose
#config = tf.ConfigProto()
#config.gpu_options.allow_growth = True
#tf.enable_eager_execution(config=config)
class Dataset(object):
def __init__(self, images_dir, patch_size, jpeg_quality, use_fast_loader=False):
self.image_files = sorted(glob.glob(images_dir + '/*'))
self.patch_size = patch_size
self.jpeg_quality = jpeg_quality
self.use_fast_loader = use_fast_loader
def __getitem__(self, idx):
if self.use_fast_loader:
#label = tf.read_file(self.image_files[idx])
label = mindspore.mindrecord.FileReader(self.image_files[idx])
#label = tf.image.decode_jpeg(label, channels=3)
label = Compose([label.Decode(),
label.RandomHorizontalFlip(0.5),
label.ToTensor()])
# apply the transform to dataset through map function
image_folder_dataset = image_folder_dataset.map(operations=label,
input_columns="image")
label = pil_image.fromarray(label.numpy())
else:
label = pil_image.open(self.image_files[idx]).convert('RGB')
# randomly crop patch from training set
crop_x = random.randint(0, label.width - self.patch_size)
crop_y = random.randint(0, label.height - self.patch_size)
label = label.crop((crop_x, crop_y, crop_x + self.patch_size, crop_y + self.patch_size))
# additive jpeg noise
buffer = io.BytesIO()
label.save(buffer, format='jpeg', quality=self.jpeg_quality)
input = pil_image.open(buffer)
input = np.array(input).astype(np.float32)
label = np.array(label).astype(np.float32)
input = np.transpose(input, axes=[2, 0, 1])
label = np.transpose(label, axes=[2, 0, 1])
# normalization
input /= 255.0
label /= 255.0
return input, label
def __len__(self):
return len(self.image_files)