CNN classification added to computer vision

computer_vision/cnn_classification.py

@@ -0,0 +1,98 @@
+"""
+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"