Spaces:
Sleeping
Sleeping
| 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'") | |