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)