Python/computer_vision/cnn_classification.py
Harshit Agarwal 2899cdac20
feat: CNN classification added to computer vision (#4350)
* cnn classification file
* black formatted
* flake8 corrected
* added cnn classification
* Delete requirements.txt
* Update cnn_classification.py
* Create cnn_classification.py
* using keras from tensorflow only
* update tensorflow
* Update cnn_classification.py
* Delete computer_vision/cnn_classification directory
2021-06-24 11:58:23 +05:30

99 lines
2.9 KiB
Python

"""
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"