import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.applications import VGG16
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten, Dropout
import os


DATA_DIR = 'dataset/Training' 
IMAGE_SIZE = (224, 224) # VGG16 ke liye standard input size
BATCH_SIZE = 32
NUM_CLASSES = 4 # gliomat_tumor, meningioma_tumor, no_tumor, pituitary_tumor

train_datagen = ImageDataGenerator(
    rescale=1./255,
    validation_split=0.2, # 20% data validation ke liye use hoga

)


train_generator = train_datagen.flow_from_directory(
    DATA_DIR,
    target_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE,
    class_mode='categorical', # 4 classes ke liye 'categorical' use karte hain
    subset='training'
)


validation_generator = train_datagen.flow_from_directory(
    DATA_DIR,
    target_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE,
    class_mode='categorical',
    subset='validation'
)


base_model = VGG16(weights='imagenet', 
                   include_top=False, 
                   input_shape=(IMAGE_SIZE[0], IMAGE_SIZE[1], 3))


for layer in base_model.layers:
    layer.trainable = False

model = Sequential([
    base_model,
    Flatten(), # 2D output ko 1D vector mein badalna
    Dense(256, activation='relu'), # Pehli hidden layer
    Dropout(0.5), # Overfitting se bachne ke liye
    Dense(NUM_CLASSES, activation='softmax') # Output layer (4 classes)
])


model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

model.summary()


print("\n--- Starting Model Training ---")

history = model.fit(
    train_generator,
    epochs=10, 
    validation_data=validation_generator
)

MODEL_FILENAME = 'brain_tumor_model.keras'
model.save(MODEL_FILENAME) 
print(f"\nModel successfully saved as: {MODEL_FILENAME}")