"""
Convolutional Neural Network

Objective : To train a CNN model detect if TB is present in Lung X-ray or not.

Resources CNN Theory :
 https://en.wikipedia.org/wiki/Convolutional_neural_network
Resources Tensorflow : https://www.tensorflow.org/tutorials/images/cnn

Download dataset from :
https://lhncbc.nlm.nih.gov/LHC-publications/pubs/TuberculosisChestXrayImageDataSets.html

1. Download the dataset folder and create two folder training set and test set
in the parent dataste folder
2. Move 30-40 image from both TB positive and TB Negative folder
in the test set folder
3. The labels of the iamges will be extracted from the folder name
the image is present in.

"""

# Part 1 - Building the CNN

import numpy as np

# Importing the Keras libraries and packages
import tensorflow as tf
from tensorflow.keras import layers, models

if __name__ == "__main__":

 # Initialising the CNN
 classifier = models.Sequential()

 # Step 1 - Convolution
 classifier.add(
 layers.Conv2D(32, (3, 3), input_shape=(64, 64, 3), activation="relu")
 )

 # Step 2 - Pooling
 classifier.add(layers.MaxPooling2D(pool_size=(2, 2)))

 # Adding a second convolutional layer
 classifier.add(layers.Conv2D(32, (3, 3), activation="relu"))
 classifier.add(layers.MaxPooling2D(pool_size=(2, 2)))

 # Step 3 - Flattening
 classifier.add(layers.Flatten())

 # Step 4 - Full connection
 classifier.add(layers.Dense(units=128, activation="relu"))
 classifier.add(layers.Dense(units=1, activation="sigmoid"))

 # Compiling the CNN
 classifier.compile(
 optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"]
 )

 # Part 2 - Fitting the CNN to the images

 # Load Trained model weights

 # from keras.models import load_model
 # regressor=load_model('cnn.h5')

 train_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
 rescale=1.0 / 255, shear_range=0.2, zoom_range=0.2, horizontal_flip=True
 )

 test_datagen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1.0 / 255)

 training_set = train_datagen.flow_from_directory(
 "dataset/training_set", target_size=(64, 64), batch_size=32, class_mode="binary"
 )

 test_set = test_datagen.flow_from_directory(
 "dataset/test_set", target_size=(64, 64), batch_size=32, class_mode="binary"
 )

 classifier.fit_generator(
 training_set, steps_per_epoch=5, epochs=30, validation_data=test_set
 )

 classifier.save("cnn.h5")

 # Part 3 - Making new predictions

 test_image = tf.keras.preprocessing.image.load_img(
 "dataset/single_prediction/image.png", target_size=(64, 64)
 )
 test_image = tf.keras.preprocessing.image.img_to_array(test_image)
 test_image = np.expand_dims(test_image, axis=0)
 result = classifier.predict(test_image)
 training_set.class_indices
 if result[0][0] == 0:
 prediction = "Normal"
 if result[0][0] == 1:
 prediction = "Abnormality detected"