File size: 3,630 Bytes
612c03b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import os
import numpy as np
import cv2
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.callbacks import ModelCheckpoint
# Function to load images from specified folders
def load_images_from_folders(folders):
images = []
for folder in folders:
for filename in os.listdir(folder):
if filename.lower().endswith(('.png', '.jpg', '.jpeg')): # Check for valid image extensions
img = cv2.imread(os.path.join(folder, filename), cv2.IMREAD_GRAYSCALE) # Read image as grayscale
if img is not None:
img = cv2.resize(img, (224, 224)) # Resize to 224x224 pixels
img = img.astype(np.float32) # Ensure the image is in float32 format
img /= 255.0 # Normalize to [0, 1]
images.append(img)
else:
print(f"Failed to load image: {filename}")
return np.array(images)
# Load normal and pneumonia images
normal_folders = [
os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'NORMAL'),
os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'NORMAL'),
os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'NORMAL'),
]
pneumonia_folders = [
os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'PNEUMONIA'),
os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'PNEUMONIA'),
os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'PNEUMONIA'),
]
normal_images = load_images_from_folders(normal_folders)
pneumonia_images = load_images_from_folders(pneumonia_folders)
# Print the number of images loaded for each category
print(f"Loaded {len(normal_images)} normal images.")
print(f"Loaded {len(pneumonia_images)} pneumonia images.")
# Prepare dataset for machine learning
images = np.concatenate([normal_images, pneumonia_images]) # Combine normal and pneumonia images
labels = np.array([0] * len(normal_images) + [1] * len(pneumonia_images)) # Create labels
# Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)
# Convert to RGB format
X_train_rgb = np.stack([cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_GRAY2RGB) for img in X_train])
X_test_rgb = np.stack([cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_GRAY2RGB) for img in X_test])
print(f'Training data shape: {X_train_rgb.shape}, Training labels shape: {y_train.shape}')
print(f'Testing data shape: {X_test_rgb.shape}, Testing labels shape: {y_test.shape}')
# Build a simple CNN model
model = Sequential([
Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)),
MaxPooling2D(pool_size=(2, 2)),
Conv2D(64, (3, 3), activation='relu'),
MaxPooling2D(pool_size=(2, 2)),
Flatten(),
Dense(128, activation='relu'),
Dense(1, activation='sigmoid') # Binary classification
])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
# Set up model checkpointing with the new file extension
checkpoint = ModelCheckpoint('pneumonia_model.keras', save_best_only=True, monitor='val_loss', mode='min')
# Train the model
model.fit(X_train_rgb, y_train, validation_data=(X_test_rgb, y_test), epochs=10, callbacks=[checkpoint])
# Save the model in the new format
model.save('pneumonia_model_final.keras')
print("Model saved as 'pneumonia_model_final.keras'")
|