import os import streamlit as st import tensorflow as tf from tensorflow.keras import layers, models from tensorflow.keras.preprocessing.image import ImageDataGenerator import numpy as np import matplotlib.pyplot as plt import json from datetime import datetime def create_model_directories(): """Create necessary directories for model storage""" os.makedirs('models', exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") run_dir = os.path.join('models', f'run_{timestamp}') os.makedirs(run_dir, exist_ok=True) return run_dir def train_model(epochs, batch_size, learning_rate): # Create directories run_dir = create_model_directories() # Set the paths to your image folders BASE_DIR = os.path.join(os.getcwd(), 'dataset') train_dir = os.path.join(BASE_DIR, 'train') validation_dir = os.path.join(BASE_DIR, 'validation') test_dir = os.path.join(BASE_DIR, 'test') # Verify dataset directories exist for dir_path in [train_dir, validation_dir, test_dir]: if not os.path.exists(dir_path): st.error(f"Directory not found: {dir_path}") return None, None # Set the parameters for the data generators img_height, img_width = 256, 256 # Create data generators with data augmentation for training train_datagen = ImageDataGenerator( rescale=1.0/255.0, rotation_range=20, width_shift_range=0.2, height_shift_range=0.2, shear_range=0.2, zoom_range=0.2, horizontal_flip=True, vertical_flip=True, fill_mode='nearest' ) train_generator = train_datagen.flow_from_directory( train_dir, target_size=(img_height, img_width), batch_size=batch_size, class_mode='categorical' ) validation_datagen = ImageDataGenerator(rescale=1.0/255.0) validation_generator = validation_datagen.flow_from_directory( validation_dir, target_size=(img_height, img_width), batch_size=batch_size, class_mode='categorical' ) # Create a CNN model cnn_model = models.Sequential([ # First Convolutional Block layers.Conv2D(32, (3, 3), activation='relu', input_shape=(img_height, img_width, 3)), layers.BatchNormalization(), layers.Conv2D(32, (3, 3), activation='relu'), layers.BatchNormalization(), layers.MaxPooling2D((2, 2)), layers.Dropout(0.25), # Second Convolutional Block layers.Conv2D(64, (3, 3), activation='relu'), layers.BatchNormalization(), layers.Conv2D(64, (3, 3), activation='relu'), layers.BatchNormalization(), layers.MaxPooling2D((2, 2)), layers.Dropout(0.25), # Third Convolutional Block layers.Conv2D(128, (3, 3), activation='relu'), layers.BatchNormalization(), layers.Conv2D(128, (3, 3), activation='relu'), layers.BatchNormalization(), layers.MaxPooling2D((2, 2)), layers.Dropout(0.25), # Dense Layers layers.Flatten(), layers.Dense(256, activation='relu'), layers.BatchNormalization(), layers.Dropout(0.5), layers.Dense(train_generator.num_classes, activation='softmax') ]) # Compile the model optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate) cnn_model.compile( optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy'] ) # Training callbacks callbacks = [ tf.keras.callbacks.EarlyStopping( monitor='val_loss', patience=5, restore_best_weights=True ), tf.keras.callbacks.ModelCheckpoint( filepath=os.path.join(run_dir, 'best_model.h5'), monitor='val_accuracy', save_best_only=True ) ] # Train the model st.write("Starting model training...") st.write(f"Number of classes: {train_generator.num_classes}") st.write(f"Training samples: {train_generator.samples}") st.write(f"Validation samples: {validation_generator.samples}") history = cnn_model.fit( train_generator, epochs=epochs, validation_data=validation_generator, callbacks=callbacks ) # Save the model model_path = os.path.join(run_dir, 'rose_model.h5') cnn_model.save(model_path) st.success(f"Model saved to: {model_path}") # Save class names class_names = list(train_generator.class_indices.keys()) class_names_path = os.path.join(run_dir, 'class_names.json') with open(class_names_path, 'w') as f: json.dump(class_names, f) st.success(f"Class names saved to: {class_names_path}") # Plot training history fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5)) # Plot accuracy ax1.plot(history.history['accuracy'], label='Training Accuracy') ax1.plot(history.history['val_accuracy'], label='Validation Accuracy') ax1.set_title('Model Accuracy') ax1.set_ylabel('Accuracy') ax1.set_xlabel('Epoch') ax1.legend(loc='lower right') ax1.grid(True) # Plot loss ax2.plot(history.history['loss'], label='Training Loss') ax2.plot(history.history['val_loss'], label='Validation Loss') ax2.set_title('Model Loss') ax2.set_ylabel('Loss') ax2.set_xlabel('Epoch') ax2.legend(loc='upper right') ax2.grid(True) plt.tight_layout() history_path = os.path.join(run_dir, 'training_history.png') plt.savefig(history_path) plt.close() return history_path, run_dir def main(): st.title("Rose Classification Model Training") st.write("Train your rose classification model with custom parameters") # Training parameters col1, col2 = st.columns(2) with col1: epochs = st.slider("Number of Epochs", min_value=1, max_value=100, value=50, step=1) batch_size = st.slider("Batch Size", min_value=8, max_value=64, value=32, step=8) with col2: learning_rate = st.slider("Learning Rate", min_value=0.0001, max_value=0.01, value=0.001, step=0.0001) if st.button("Start Training"): with st.spinner("Training in progress..."): history_path, run_dir = train_model(epochs, batch_size, learning_rate) if history_path and run_dir: st.image(history_path, caption="Training History") st.success(f"Training completed! Files saved in: {run_dir}") if __name__ == "__main__": main()