Spaces:
Sleeping
Sleeping
| from tensorflow.keras.models import Sequential | |
| from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense | |
| from tensorflow.keras.preprocessing.image import ImageDataGenerator | |
| import os | |
| def train_model(): | |
| data_path = 'data/brain_tumor_dataset' | |
| img_size = (150, 150) | |
| datagen = ImageDataGenerator(rescale=1./255, validation_split=0.2) | |
| train_data = datagen.flow_from_directory( | |
| data_path, | |
| target_size=img_size, | |
| batch_size=32, | |
| class_mode='binary', | |
| subset='training' | |
| ) | |
| val_data = datagen.flow_from_directory( | |
| data_path, | |
| target_size=img_size, | |
| batch_size=32, | |
| class_mode='binary', | |
| subset='validation' | |
| ) | |
| model = Sequential([ | |
| Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)), | |
| MaxPooling2D(2, 2), | |
| Flatten(), | |
| Dense(64, activation='relu'), | |
| Dense(1, activation='sigmoid') # Binary classifier: tumor vs no tumor | |
| ]) | |
| model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy']) | |
| model.fit(train_data, epochs=10, validation_data=val_data) | |
| os.makedirs("model", exist_ok=True) | |
| model.save("model/tumor_classifier.h5") | |
| print("✅ Model trained and saved to model/tumor_classifier.h5") | |
| if __name__ == "__main__": | |
| train_model() | |