Brain_Tumor_Detection / utils /model_train.py
the11's picture
Upload 55 files
eb6c18c verified
Raw
History Blame Contribute Delete
1.39 kB
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()