Brain_Tumor_Detection / utils /model_train.py
the
Upload 55 files
eb6c18c verified
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()