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() | |