Spaces:
Sleeping
Sleeping
File size: 1,385 Bytes
eb6c18c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
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()
|