File size: 2,389 Bytes
b620cf3 b79a585 b620cf3 b79a585 b620cf3 b79a585 b620cf3 b79a585 b620cf3 b79a585 |
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 64 65 66 67 68 69 70 71 72 |
import os
import numpy as np
from PIL import Image
from re import split, compile
from tensorflow.keras.utils import Sequence
def list_filenames(data_path, img_extension='png', filename_prefix=None):
if filename_prefix is None:
files_list = [file for file in os.listdir(data_path) if file.endswith(img_extension)]
else:
files_list = [file for file in os.listdir(data_path) if file.endswith(img_extension) and file.startswith(filename_prefix)]
files_list.sort(key=lambda l: [int(s) if s.isdigit() else s.lower() for s in split(compile(r'(\d+)'), l)])
files_list = [os.path.join(data_path, file) for file in files_list]
print('Found {} files in {}'.format(len(files_list), data_path))
return files_list
class Dataset(Sequence):
def __init__(self, file_list, batch_size=32, crop_dim=None, resize_dim=None, shuffle=True, mode='RGB'):
self.files_list = file_list
self.batch_size = batch_size
self.crop_dim = crop_dim
self.resize_dim = resize_dim
self.shuffle = shuffle
self.on_epoch_end()
self.mode=mode
def __len__(self):
return int(np.ceil(len(self.files_list) / self.batch_size))
def __getitem__(self, idx):
batch_files = self.files_list[idx * self.batch_size : (idx + 1) * self.batch_size]
images = [self.load_images(f) for f in batch_files]
return np.stack(images)
def on_epoch_end(self):
if self.shuffle:
np.random.shuffle(self.files_list)
@staticmethod
def center_crop(image, crop_dim):
h, w = image.size
crop_h, crop_w = crop_dim
top = max(0, (w - crop_w) // 2)
left = max(0, (h - crop_h) // 2)
right = min(h - 0, (h + crop_h) // 2)
bottom = min(w - 0, (w + crop_w) // 2)
return image.crop((left, top, right, bottom))
def load_images(self, filepath):
if self.mode=='RGB':
image = Image.open(filepath).convert('RGB')
else:
image = Image.open(filepath)
if self.crop_dim:
image = self.center_crop(image, crop_dim=self.crop_dim)
if self.resize_dim:
image = image.resize(self.resize_dim)
image = np.array(image).astype(np.float32)
image = image / 255.0
if image.ndim == 2:
image = np.expand_dims(image, -1)
return image
|