import sys
import os

from keras.preprocessing.image  import ImageDataGenerator
from keras import optimizers

sys.path.append('/home/mike/densenet/DenseNet-master')
import densenet


train_dir = r'/home/mike/datasets/train/train_main'
valid_dir = r'/home/mike/datasets/validate/validate_balanced'

if __name__ == '__main__':

    # Create two generators, one for Train images, one for validation images
    train_gen = ImageDataGenerator().flow_from_directory(train_dir, target_size=(240, 240), batch_size=32,
                                                            class_mode='categorical', shuffle=True)

    valid_gen = ImageDataGenerator().flow_from_directory(valid_dir, target_size=(240, 240), batch_size=32,
                                                         class_mode='categorical')

    # Initialize my GPU
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = '2'

    # Create the model
    model = densenet.DenseNet(classes=13,
                              input_shape=(240,240,3),
                              depth=40,
                              growth_rate=12,
                              bottleneck=True,
                              reduction=0.5)

    # Create the optimizer
    optimizer = optimizers.SGD(lr=0.001,decay=0.0,momentum=0.8,nesterov=True)

    # Compile the model
    model.compile(loss='categorical_crossentropy',optimizer=optimizer,metrics=['accuracy'])

    # train the model
    try:
        model.fit_generator(train_gen,
                            steps_per_epoch=8572,
                            epochs=1,
                            validation_data=valid_gen,
                            validation_steps=117)
    except Exception as e:
        print(e)