import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from google.colab import drive
drive.mount('/content/drive')
Define constants
image_size = (150, 150)
batch_size = 32
Data augmentation for the training set
train_datagen = ImageDataGenerator(
rescale=1./255,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True
)
Rescaling for the testing set
test_datagen = ImageDataGenerator(rescale=1./255)
Load the training set
train_set = train_datagen.flow_from_directory(
'/content/drive/MyDrive/chest_xray/train',
target_size=image_size,
batch_size=batch_size,
class_mode='binary'
)
Load the testing set
test_set = test_datagen.flow_from_directory(
'/content/drive/MyDrive/chest_xray/test',
target_size=image_size,
batch_size=batch_size,
class_mode='binary'
)
Build the CNN model
model = Sequential()
model.add(Conv2D(32, (3, 3), input_shape=(image_size[0], image_size[1], 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(units=128, activation='relu'))
model.add(Dense(units=1, activation='sigmoid'))
Compile the model
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
Train the model
model.fit(train_set, epochs=10, validation_data=test_set)
Save the model
model.save('pneumonia_model.h5')
Evaluate the model on the testing set
accuracy = model.evaluate(test_set)[1]
print(f'Test Accuracy: {accuracy}')
Make predictions on new images
def predict_image(file_path):
img = tf.keras.preprocessing.image.load_img(file_path, target_size=image_size)
img_array = tf.keras.preprocessing.image.img_to_array(img)
img_array = tf.expand_dims(img_array, 0) # Create a batch
predictions = model.predict(img_array)
if predictions[0] > 0.5:
print("Prediction: Pneumonia")
else:
print("Prediction: Normal")
Example usage:
image_path = "/content/drive/MyDrive/chest_xray/train/PNEUMONIA/BACTERIA-1033441-0001.jpeg"
predict_image(image_path)