From 46be83d396bc65071b40c6f1e3cd6bd601803815 Mon Sep 17 00:00:00 2001 From: Dmitry Barsukoff Date: Sat, 19 Sep 2020 00:34:12 +0300 Subject: [PATCH] [ADD] Generator class --- .gitignore | 4 +-- .idea/Facades.iml | 2 +- .idea/misc.xml | 2 +- facades.py | 24 ++++++++----- jupyter_config_itmo.py | 2 +- ml/__init__.py | 1 + ml/generator.py | 52 ++++++++++++++++++++++++++++ utils.py | 77 ++++++++++++++++++++++++++++++++++++++++++ 8 files changed, 150 insertions(+), 14 deletions(-) create mode 100644 ml/generator.py create mode 100644 utils.py diff --git a/.gitignore b/.gitignore index 03a2f44..5821b85 100644 --- a/.gitignore +++ b/.gitignore @@ -256,6 +256,6 @@ dmypy.json # Pyre type checker .pyre/ +# User ignores dataset/sources/~$Labels.docx -ml/generator.py -utils.py +jupyter_config_itmo.py \ No newline at end of file diff --git a/.idea/Facades.iml b/.idea/Facades.iml index 1f1ea62..7352b64 100644 --- a/.idea/Facades.iml +++ b/.idea/Facades.iml @@ -2,7 +2,7 @@ - + diff --git a/.idea/misc.xml b/.idea/misc.xml index 8656114..66da066 100644 --- a/.idea/misc.xml +++ b/.idea/misc.xml @@ -3,5 +3,5 @@ - + \ No newline at end of file diff --git a/facades.py b/facades.py index a9f9cd3..94daba7 100644 --- a/facades.py +++ b/facades.py @@ -20,9 +20,10 @@ keras.backend.set_image_data_format('channels_last') -class Facades: # TODO: Predict # TODO: Comments +class Facades: # TODO: Predict + """Class for segmentation buildings' facades using Deep Learning""" def __init__(self, dim=(720, 1080), n_channels=3, use_datagenerator=True): - """Initialize main parameters of Facades class + """Initialize main parameters of Facades class s :param dim: Images resolution in format (height, width) :param n_channels: Number of channels in image (1 or 3) """ @@ -78,11 +79,12 @@ def training_init(self, dataset_dir="dataset", batch_size=8): # Initialize data generators seed = 909 # to transform image and masks with same augmentation parameter. - image_generator = data_datagen.flow_from_directory(dataset_dir + "train/data/", class_mode=None, seed=seed, - batch_size=batch_size, target_size=self.dim) - masks_generator = masks_datagen.flow_from_directory(dataset_dir + "train/masks/", class_mode=None, seed=seed, - batch_size=batch_size, target_size=self.dim, - color_mode="grayscale") + image_generator = data_datagen.flow_from_directory(dataset_dir + "train/data/", class_mode=None, + seed=seed, batch_size=batch_size, + target_size=self.dim) + masks_generator = masks_datagen.flow_from_directory(dataset_dir + "train/masks/", class_mode=None, + seed=seed, batch_size=batch_size, + target_size=self.dim, color_mode="grayscale") self.train_generator = zip(image_generator, masks_generator) @@ -117,11 +119,10 @@ def training_init(self, dataset_dir="dataset", batch_size=8): self.model = Model(base_model.input, output_layer) # keras.utils.plot_model(self.model, 'model.png', show_shapes=True) - def train(self, epochs=1, number_of_GPUs=0, remote_monitor_ip="localhost:4959"): + def train(self, epochs=1, number_of_GPUs=0): """Train model :param epochs: Number of epochs to train model :param number_of_GPUs: Number of devices to train on multiple GPUs - :param remote_monitor_ip: IP where to run Keras RemoteMonitor (default: localhost:4959) # TODO """ # Check training init before train if not self.training_init_check: @@ -151,6 +152,11 @@ def train(self, epochs=1, number_of_GPUs=0, remote_monitor_ip="localhost:4959"): if __name__ == "__main__": + # Initialize main class facades = Facades((512, 512)) + + # Initialize training dataset facades.training_init("dataset/", batch_size=5) + + # Run model training facades.train(10) diff --git a/jupyter_config_itmo.py b/jupyter_config_itmo.py index 11b6927..ddd7614 100644 --- a/jupyter_config_itmo.py +++ b/jupyter_config_itmo.py @@ -2,4 +2,4 @@ c.NotebookApp.port = 4958 c.NotebookApp.password = u'sha1:1236633a1c55:9f2f5745eb1da1304369e01b65e5ad6d64f1cb13' -c.NotebookApp.open_browser = False \ No newline at end of file +c.NotebookApp.open_browser = False diff --git a/ml/__init__.py b/ml/__init__.py index 267d4b0..e9d578f 100644 --- a/ml/__init__.py +++ b/ml/__init__.py @@ -1 +1,2 @@ +from .generator import Generator from .residual_model import residual_model diff --git a/ml/generator.py b/ml/generator.py new file mode 100644 index 0000000..5a19814 --- /dev/null +++ b/ml/generator.py @@ -0,0 +1,52 @@ +import numpy as np +from cv2 import imread +import os +from keras.utils import Sequence + + +class Generator(Sequence): + """Generates data for Keras""" + + def __init__(self, dataset_dir, batch_size=16, dim=(1920, 1080), n_channels=3, shuffle=True): + """Initialization""" + data_folder = dataset_dir + "/data" + self.data_paths = os.listdir(data_folder) + + labels_path = dataset_dir + "/masks" + self.labels_paths = os.listdir(labels_path) + + self.batch_size = batch_size + self.dim = dim + self.n_channels = n_channels + + self.shuffle = shuffle + + self.len = len(self.data_paths) + + def __len__(self): + """Denotes the number of batches per epoch + :return: number of batches per epoch + """ + return len(np.floor(self.len / self.batch_size)) + + def __getitem__(self, index): + """Generate one batch of data + :param index: index of the batch + :return: X and y when fitting. X only when predicting + """ + # Generate indexes of the batch + indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size] + + # Download data by indexes + X, y = [], [] + for i in indexes: + X.append(imread(self.data_paths[i])) + y.append(get_label(self.labels_paths[i])) + + return X, y + + def on_epoch_end(self): + """Updates indexes after each epoch""" + self.indexes = np.arange(self.len) + if self.shuffle: + np.random.shuffle(self.indexes) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..f03f4df --- /dev/null +++ b/utils.py @@ -0,0 +1,77 @@ +import os + +import numpy as np +from PIL import Image +from tqdm import tqdm + + +# def fix_dataset_labels(labels_path, colors=None): +# if colors is None: +# colors = { # RGB +# (0, 0, 255): 0, # Wall +# (0, 0, 170): 1, # Background +# (255, 255, 0): 2, # Window (closed) +# (0, 85, 255): 2, # Window +# (0, 170, 255): 7, # Door +# (170, 0, 0): 8, # Shop +# (170, 255, 85): 3, # Balcony +# (255, 85, 0): 6, # Molding +# (255, 3, 0): 11, # Pillar +# (0, 255, 255): 12, # Cornice +# (85, 255, 170): 4, # Sill +# (255, 170, 0): 9, # Deco +# # (): 11, # Blind +# } +# +# try: +# os.mkdir(labels_path + "_FIXED/") +# except: +# pass +# +# files = [f for f in os.listdir(labels_path) if f[-4:] in [".jpg", ".png"]] +# for id, image_name in enumerate(tqdm(files)): +# image_path = labels_path + "/" + image_name +# image = cv.imread(image_path) +# +# for h in range(len(image)): +# for w in range(len(image[h])): +# c = image[h][w] +# for key, value in colors.items(): +# if c[2] in range(key[0] - 3, key[0] + 3) and c[1] in range(key[1] - 2, key[1] + 3) and \ +# c[0] in range(key[2] - 2, key[2] + 3): +# image[h][w] = value +# break +# +# image = cv.cvtColor(image, cv.COLOR_BGR2GRAY) +# cv.imwrite(labels_path + "_FIXED/" + image_name, image) + + +# def show_batch_photos(facades_class): +# x_batch, y_batch = next(facades_class.train_generator) +# for i in range(0, facades_class.batch_size): +# x = x_batch[i] +# y = y_batch[i] +# +# plt.imshow(x) +# plt.show() # TODO: Check needing two +# plt.imshow(y) +# plt.show() + + +def load_images_from_folder(dir, shape): + files = [file for file in os.listdir(dir) if file.endswith((".png", ".jpg"))] + result = [] + + for image_name in tqdm(files, dir): + image_path = dir + "/" + image_name + image = Image.open(image_path) + image_resized = image.resize(shape[:-1], Image.ANTIALIAS) + image_numpy = np.array(image_resized) + if shape[-1] == 1: + image_numpy = np.expand_dims(image_numpy, axis=-1) + + result.append(image_numpy) + + return np.array(result) + +