cnn_keras.py 21.8 KB
Newer Older
1 2 3 4 5 6
#!/usr/bin/python
# -*- coding: utf-8 -*-
#
"""
    Generic classifier with multiple models
    Models -> (Xception, VGG16, VGG19, ResNet50, InceptionV3, MobileNet)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
7

8 9 10 11 12 13
    Name: cnn_keras.py
    Author: Gabriel Kirsten Menezes (gabriel.kirsten@hotmail.com)

"""
import time
import os
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
14 15
import shutil
import random
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
16
import numpy as np
17
import json
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
18 19 20
import logging
import sys

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
21
from PIL import Image
22 23 24
from keras import applications
from keras.preprocessing.image import ImageDataGenerator
from keras import optimizers
25
from keras.models import Model, load_model
26
from keras.layers import Dropout, Flatten, Dense
27 28
from keras.callbacks import ModelCheckpoint, TensorBoard
from keras import backend as K
29

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
30
from interface.interface import InterfaceException as IException
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
31
from classification.classifier import Classifier
32 33 34 35 36 37 38

from collections import OrderedDict

from util.config import Config
from util.file_utils import File
from util.utils import TimeUtils

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
39 40 41 42

logger = logging.getLogger('PIL')
logger.setLevel(logging.WARNING)

43 44 45 46 47 48 49
START_TIME = time.time()

# =========================================================
# Constants
# =========================================================

IMG_WIDTH, IMG_HEIGHT = 256, 256
50
weight_path = None
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
51

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
52

53
class CNNKeras(Classifier):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
54 55
    """ Class for CNN classifiers based on Keras applications """

56
    def __init__(self, architecture="ResNet50", learning_rate=0.001, momentum=0.9, batch_size=32, epochs=50, fine_tuning_rate=100, transfer_learning=False, save_weights=True, perc_train=80, perc_validation=20, recreate_dataset=False):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
57 58 59 60
        """
            Constructor of CNNKeras
        """

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76
        self.architecture = Config(
            "Architecture", architecture, str)
        self.learning_rate = Config(
            "Learning rate", learning_rate, float)
        self.momentum = Config(
            "Momentum", momentum, float)
        self.batch_size = Config(
            "Batch size", batch_size, int)
        self.epochs = Config(
            "Epochs", epochs, int)
        self.fine_tuning_rate = Config(
            "Fine Tuning Rate", fine_tuning_rate, int)
        self.transfer_learning = Config(
            "Transfer Learning", transfer_learning, bool)
        self.save_weights = Config(
            "Save weights", save_weights, bool)
77 78 79 80 81 82
        self.perc_train = Config(
            "Perc Train", perc_train, float)
        self.perc_validation = Config(
            "Perc Validation", perc_validation, float)
        self.recreate_dataset = Config(
            "Recreate Dataset", recreate_dataset, bool)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
83
        self.file_name = "kerasCNN"
84

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
85
        self.model = None
86

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
87
        self.trained = False
88 89

    def get_config(self):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
90
        """Return configuration of classifier.
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
91

92 93 94 95 96
        Returns
        -------
        config : OrderedDict
            Current configs of classifier.
        """
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
97 98
        keras_config = OrderedDict()

99 100 101 102 103 104 105 106
        keras_config["Architecture"] = self.architecture
        keras_config["Learning rate"] = self.learning_rate
        keras_config["Momentum"] = self.momentum
        keras_config["Batch size"] = self.batch_size
        keras_config["Epochs"] = self.epochs
        keras_config["Fine Tuning rate"] = self.fine_tuning_rate
        keras_config["Transfer Learning"] = self.transfer_learning
        keras_config["Save weights"] = self.save_weights
107 108 109
        keras_config["Perc Train"] = self.perc_train
        keras_config["Perc Validation"] = self.perc_validation
        keras_config["Recreate Dataset"] = self.recreate_dataset
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
110
        return keras_config
111 112

    def set_config(self, configs):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
113
        """Update configuration of classifier.
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
114

115 116 117 118 119
        Parameters
        ----------
        configs : OrderedDict
            New configs of classifier.
        """
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
120 121 122 123
        self.architecture = Config.nvl_config(
            configs["Architecture"], self.architecture)
        self.learning_rate = Config.nvl_config(
            configs["Learning rate"], self.learning_rate)
124
        self.momentum = Config.nvl_config(configs["Momentum"], self.momentum)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
125 126
        self.batch_size = Config.nvl_config(
            configs["Batch size"], self.batch_size)
127
        self.epochs = Config.nvl_config(configs["Epochs"], self.epochs)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
128 129 130 131 132 133 134 135 136 137 138 139
        self.fine_tuning_rate = Config.nvl_config(
            configs["Fine Tuning rate"], self.fine_tuning_rate)
        self.transfer_learning = Config.nvl_config(
            configs["Transfer Learning"], self.transfer_learning)
        self.save_weights = Config.nvl_config(
            configs["Save weights"], self.save_weights)
        self.perc_train = Config.nvl_config(
            configs["Perc Train"], self.perc_train)
        self.perc_validation = Config.nvl_config(
            configs["Perc Validation"], self.perc_validation)
        self.recreate_dataset = Config.nvl_config(
            configs["Recreate Dataset"], self.recreate_dataset)
140 141

    def get_summary_config(self):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
142
        """Return fomatted summary of configuration.
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
143

144 145 146 147 148
        Returns
        -------
        summary : string
            Formatted string with summary of configuration.
        """
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
149
        keras_config = OrderedDict()
150

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
151
        keras_config[self.architecture.label] = self.architecture.value
152 153 154 155
        keras_config[self.learning_rate.label] = self.learning_rate.value
        keras_config[self.momentum.label] = self.momentum.value
        keras_config[self.batch_size.label] = self.batch_size.value
        keras_config[self.epochs.label] = self.epochs.value
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
156
        keras_config[self.fine_tuning_rate.label] = self.fine_tuning_rate.value
157 158
        keras_config[self.transfer_learning.label] = self.transfer_learning.value
        keras_config[self.save_weights.label] = self.save_weights.value
159 160 161
        keras_config[self.perc_train.label] = self.perc_train.value
        keras_config[self.perc_validation.label] = self.perc_validation.value
        keras_config[self.recreate_dataset.label] = self.recreate_dataset.value
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
162 163 164
        summary = ''
        for config in keras_config:
            summary += "%s: %s\n" % (config, str(keras_config[config]))
165

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
166
        return summary
167

168
    def classify(self, dataset, test_dir, test_data, image):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
169
        """"Perform the classification.
170

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
171 172 173 174 175 176 177 178
        Parameters
        ----------
        dataset : string
            Path to image dataset.
        test_dir : string
            Not used.
        test_data : string
            Name of test data file.
179

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
180 181 182 183 184 185 186 187
        Returns
        -------
        summary : list of string
            List of predicted classes for each instance in test data in ordered way.
        """

        predict_directory = File.make_path(dataset, test_dir)

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
188
        # Create a Keras class
189 190 191 192 193 194
        if not os.path.exists(File.make_path(predict_directory, "png")):
            os.makedirs(File.make_path(predict_directory, "png"))

        for file in os.listdir(predict_directory):
            print(File.make_path(predict_directory, file))
            if os.path.splitext(file)[-1] == ".tif":
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
195 196 197 198 199 200 201 202
                try:
                    img = Image.open(File.make_path(predict_directory, file))
                    # img.thumbnail(img.size)
                    new_file = os.path.splitext(file)[0] + ".png"
                    img.save(File.make_path(predict_directory,
                                            'png', new_file), "PNG", quality=100)
                except Exception, e:
                    print e
203
            else:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
204 205
                os.symlink(File.make_path(predict_directory, file),
                           File.make_path(predict_directory, 'png', file))
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
206

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
207 208 209
        classify_datagen = ImageDataGenerator()

        classify_generator = classify_datagen.flow_from_directory(
210
            File.make_path(predict_directory, 'png'),
211
            target_size=(IMG_HEIGHT, IMG_WIDTH),
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
212 213 214 215 216
            batch_size=1,
            shuffle=False,
            class_mode=None)

        try:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
217 218
            # self.model.load_weights(
            #"../models_checkpoints/" + self.file_name + ".h5")
219 220 221
            K.clear_session()
            if self.weight_path is not None:
                self.model = load_model(self.weight_path)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
222 223
                path_classes = self.weight_path.replace(
                    "_model.h5", "_classes.npy")
224
                CLASS_NAMES = np.load(path_classes).item().keys()
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
225 226
        except Exception, e:
            raise IException("Can't load the model in " +
227
                             self.weight_path + str(e))
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
228

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
229
        output_classification = self.model.predict_generator(
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
230
            classify_generator, classify_generator.samples, verbose=2)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
231 232 233

        one_hot_output = np.argmax(output_classification, axis=1)

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
234 235
        one_hot_output = one_hot_output.tolist()

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252
        for index in range(0, len(one_hot_output)):
            one_hot_output[index] = CLASS_NAMES[one_hot_output[index]]

        return one_hot_output

    def train(self, dataset, training_data, force=False):
        """Perform the training of classifier.

        Parameters
        ----------
        dataset : string
            Path to image dataset.
        training_data : string
            Name of ARFF training file.
        force : boolean, optional, default = False
            If False don't perform new training if there is trained data.
        """
253

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
254
        # select .h5 filename
255 256 257
        if self.fine_tuning_rate.value == 100:
            self.file_name = str(self.architecture.value) + \
                '_learning_rate' + str(self.learning_rate.value) + \
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
258
                '_transfer_learning'
259 260 261
        elif self.fine_tuning_rate.value == -1:
            self.file_name = str(self.architecture.value) + \
                '_learning_rate' + str(self.learning_rate.value) + \
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
262 263
                '_without_transfer_learning'
        else:
264 265
            self.file_name = str(self.architecture.value) + \
                '_learning_rate' + str(self.learning_rate.value) + \
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
266
                '_fine_tunning_' + str(self.fine_tuning_rate.value)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
267

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
268 269
        File.remove_dir(File.make_path(dataset, ".tmp"))

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
270 271
        train_generator, validation_generator, test_generator = self.make_dataset(
            dataset)
272 273 274 275 276

        # Save the model according to the conditions
        if self.save_weights:
            if not os.path.exists("../models_checkpoints/"):
                os.makedirs("../models_checkpoints/")
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
277

278
            checkpoint = ModelCheckpoint("../models_checkpoints/" + self.file_name + ".h5", monitor='val_acc',
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
279 280
                                         verbose=1, save_best_only=True, save_weights_only=False,
                                         mode='auto', period=1)
281 282 283
        else:
            checkpoint = None

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
284 285
        self.model = self.select_model_params(train_generator.num_classes)

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
286 287 288
        tensorboard = TensorBoard(
            log_dir="../models_checkpoints/logs_" + self.file_name, write_images=False)
        # tensorboard.set_model(self.model)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
289 290 291 292 293
        # compile the model
        self.model.compile(loss="categorical_crossentropy",
                           optimizer=optimizers.SGD(
                               lr=self.learning_rate.value, momentum=self.momentum.value),
                           metrics=["accuracy"])
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
294

295
        # Train the model
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
296
        self.model.fit_generator(
297
            train_generator,
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
298 299
            steps_per_epoch=train_generator.samples // self.batch_size.value,
            epochs=self.epochs.value,
300 301 302
            callbacks=[checkpoint, tensorboard],
            validation_data=validation_generator,
            validation_steps=validation_generator.samples // self.batch_size.value)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
303 304

        if self.save_weights:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
305
            # self.model.save_weights(
306 307 308 309 310 311
            #    "../models_checkpoints/" + self.file_name + ".h5")
            self.model.save(
                "../models_checkpoints/" + self.file_name + "_model.h5")
            self.weight_path = "../models_checkpoints/" + self.file_name + "_model.h5"

            dict_classes = validation_generator.class_indices
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
312 313
            np.save("../models_checkpoints/" + self.file_name +
                    "_classes.npy", dict_classes)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
314 315

    def must_train(self):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
316
        """Return if classifier must be trained.
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
317 318 319 320 321

        Returns
        -------
        True
        """
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
322
        return not self.trained
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
323 324

    def must_extract_features(self):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
325
        """Return if classifier must be extracted features.
326

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
327 328 329 330 331
        Returns
        -------
        False
        """
        return False
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353

    def select_model_params(self, num_classes):
        if self.fine_tuning_rate.value != -1:
            if self.architecture.value == "Xception":
                model = applications.Xception(
                    weights="imagenet", include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "VGG16":
                model = applications.VGG16(
                    weights="imagenet", include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "VGG19":
                model = applications.VGG19(
                    weights="imagenet", include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "ResNet50":
                model = applications.ResNet50(
                    weights="imagenet", include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "InceptionV3":
                model = applications.InceptionV3(
                    weights="imagenet", include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "MobileNet":
                model = applications.MobileNet(
                    weights="imagenet", include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))

354
            for layer in model.layers[:int(len(model.layers) * (self.fine_tuning_rate.value / 100.0))]:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395
                layer.trainable = False

        else:  # without transfer learning
            if self.architecture.value == "Xception":
                model = applications.Xception(
                    weights=None, include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "VGG16":
                model = applications.VGG16(
                    weights=None, include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "VGG19":
                model = applications.VGG19(
                    weights=None, include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "ResNet50":
                model = applications.ResNet50(
                    weights=None, include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "InceptionV3":
                model = applications.InceptionV3(
                    weights=None, include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            elif self.architecture.value == "MobileNet":
                model = applications.MobileNet(
                    weights=None, include_top=False, input_shape=(IMG_WIDTH, IMG_HEIGHT, 3))
            for layer in model.layers:
                layer.trainable = True

        # Adding custom Layers
        new_custom_layers = model.output
        new_custom_layers = Flatten()(new_custom_layers)
        new_custom_layers = Dense(1024, activation="relu")(new_custom_layers)
        new_custom_layers = Dropout(0.5)(new_custom_layers)
        new_custom_layers = Dense(1024, activation="relu")(new_custom_layers)
        predictions = Dense(num_classes,
                            activation="softmax")(new_custom_layers)

        # creating the final model
        model = Model(inputs=model.input, outputs=predictions)

        return model

    def make_dataset(self, dataset):

        # create symbolic links to the dataset
396 397
        KERAS_DATASET_DIR_NAME = ".keras_dataset"
        #KERAS_DATASET_DIR_NAME = File.make_path("..", os.path.split(dataset)[-1] + "_keras_dataset")
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
398 399
        KERAS_DIR_TRAIN_NAME = "train"
        KERAS_DIR_VALIDATION_NAME = "validation"
400 401 402 403
        KERAS_DIR_TEST_NAME = "test"
        PERC_TRAIN = self.perc_train.value
        PERC_VALIDATION = self.perc_validation.value

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
404
        # create keras dir dataset
405 406 407 408
        if not os.path.exists(File.make_path(dataset, KERAS_DATASET_DIR_NAME)) or self.recreate_dataset.value:
            if os.path.exists(File.make_path(dataset, KERAS_DATASET_DIR_NAME)):
                shutil.rmtree(File.make_path(dataset, KERAS_DATASET_DIR_NAME))

Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
409 410
            os.makedirs(File.make_path(dataset, KERAS_DATASET_DIR_NAME))

411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439
            # create keras dir train
            if not os.path.exists(File.make_path(dataset, KERAS_DATASET_DIR_NAME, KERAS_DIR_TRAIN_NAME)):
                os.makedirs(File.make_path(
                    dataset, KERAS_DATASET_DIR_NAME, KERAS_DIR_TRAIN_NAME))

            # create keras dir validation
            if not os.path.exists(File.make_path(dataset, KERAS_DATASET_DIR_NAME, KERAS_DIR_VALIDATION_NAME)):
                os.makedirs(File.make_path(
                    dataset, KERAS_DATASET_DIR_NAME, KERAS_DIR_VALIDATION_NAME))

            # create keras dir test
            if not os.path.exists(File.make_path(dataset, KERAS_DATASET_DIR_NAME, KERAS_DIR_TEST_NAME)):
                os.makedirs(File.make_path(
                    dataset, KERAS_DATASET_DIR_NAME, KERAS_DIR_TEST_NAME))

            dir_classes = sorted(File.list_dirs(dataset))

            if KERAS_DATASET_DIR_NAME in dir_classes:
                dir_classes.remove(KERAS_DATASET_DIR_NAME)

            for dir_class in dir_classes:
                root = File.make_path(dataset, dir_class)
                files = os.listdir(root)
                random.shuffle(files)
                quant_files = len(files)
                quant_train = int((quant_files / 100.0) * PERC_TRAIN)
                quant_validation = int((quant_files / 100.0) * PERC_VALIDATION)

                files_train = files[0:quant_train]
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
440 441 442 443 444
                files_validation = files[quant_train:quant_train +
                                         quant_validation]
                files_test = files[quant_train + quant_validation:quant_files]
                print("Processing class %s - %d itens - %d train items - %d validation items" %
                      (dir_class, quant_files, quant_train, quant_validation))
445 446

                for file in files_train:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
447 448
                    dir_class_train = File.make_path(
                        dataset, KERAS_DATASET_DIR_NAME, KERAS_DIR_TRAIN_NAME, dir_class)
449 450
                    if not os.path.exists(dir_class_train):
                        os.makedirs(dir_class_train)
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
451

452 453
                    if os.path.splitext(file)[-1] == ".tif":
                        img = Image.open(File.make_path(root, file))
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
454 455 456 457
                        # img.thumbnail(img.size)
                        new_file = os.path.splitext(file)[0] + ".png"
                        img.save(File.make_path(dir_class_train,
                                                new_file), "PNG", quality=100)
458
                    else:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
459 460
                        os.symlink(File.make_path(root, file),
                                   File.make_path(dir_class_train, file))
461 462

                for file in files_validation:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
463 464
                    dir_class_validation = File.make_path(
                        dataset, KERAS_DATASET_DIR_NAME, KERAS_DIR_VALIDATION_NAME, dir_class)
465
                    if not os.path.exists(dir_class_validation):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
466 467
                        os.makedirs(dir_class_validation)

468 469
                    if os.path.splitext(file)[-1] == ".tif":
                        img = Image.open(File.make_path(root, file))
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
470 471 472 473
                        # img.thumbnail(img.size)
                        new_file = os.path.splitext(file)[0] + ".png"
                        img.save(File.make_path(dir_class_validation,
                                                new_file), "PNG", quality=100)
474
                    else:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
475 476
                        os.symlink(File.make_path(root, file),
                                   File.make_path(dir_class_validation, file))
477 478

                for file in files_test:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
479 480
                    dir_class_test = File.make_path(
                        dataset, KERAS_DATASET_DIR_NAME, KERAS_DIR_TEST_NAME, dir_class)
481
                    if not os.path.exists(dir_class_test):
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
482 483
                        os.makedirs(dir_class_test)

484 485
                    if os.path.splitext(file)[-1] == ".tif":
                        img = Image.open(File.make_path(root, file))
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
486 487 488 489
                        # img.thumbnail(img.size)
                        new_file = os.path.splitext(file)[0] + ".png"
                        img.save(File.make_path(dir_class_test,
                                                new_file), "PNG", quality=100)
490
                    else:
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
491 492
                        os.symlink(File.make_path(root, file),
                                   File.make_path(dir_class_test, file))
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
493 494 495 496

        train_datagen = ImageDataGenerator()

        train_generator = train_datagen.flow_from_directory(
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
497 498
            File.make_path(dataset, KERAS_DATASET_DIR_NAME,
                           KERAS_DIR_TRAIN_NAME),
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
499 500 501 502 503 504 505 506
            target_size=(IMG_HEIGHT, IMG_WIDTH),
            batch_size=self.batch_size.value,
            shuffle=True,
            class_mode="categorical")

        validation_datagen = ImageDataGenerator()

        validation_generator = validation_datagen.flow_from_directory(
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
507 508
            File.make_path(dataset, KERAS_DATASET_DIR_NAME,
                           KERAS_DIR_VALIDATION_NAME),
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
509 510 511 512 513
            target_size=(IMG_HEIGHT, IMG_WIDTH),
            batch_size=self.batch_size.value,
            shuffle=True,
            class_mode="categorical")

514 515 516
        test_datagen = ImageDataGenerator()

        test_generator = test_datagen.flow_from_directory(
Gabriel Kirsten's avatar
 
Gabriel Kirsten committed
517 518
            File.make_path(dataset, KERAS_DATASET_DIR_NAME,
                           KERAS_DIR_TEST_NAME),
519 520 521 522 523 524
            target_size=(IMG_HEIGHT, IMG_WIDTH),
            batch_size=self.batch_size.value,
            shuffle=True,
            class_mode="categorical")

        return train_generator, validation_generator, test_generator
525 526

    def single_classify(self, image_path, directory, extractors, dict_classes):
Geazy Menezes's avatar
Geazy Menezes committed
527
        from matplotlib.pylab import imread
528 529 530
        image = imread(image_path)
        predict = self.model.predict([image])[0]
        return dict_classes[predict]