CodingMaster24/StreamlitWasteSorting

1# Save this as `waste_sorting_app.py` and run it with `streamlit run waste_sorting_app.py`2 3import streamlit as st4import tensorflow as tf5from tensorflow.keras.preprocessing.image import ImageDataGenerator6from tensorflow.keras.models import Sequential7from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout8import os9 10# Title and Description11st.title("Waste Sorting Model Training App")12st.write("This app trains a CNN model on the waste sorting dataset.")13 14# Dataset Directory Input15data_dir = st.text_input("Enter Dataset Directory Path", 'https://github.com/vikram3192/Waste-sorting-using-computer-vision/tree/e17c3e5bf2b1f2c88c07f14e91d658184fbc7602/dataset-resized')16img_size, batch_size = (180, 180), 3217 18# Prepare data generators19@st.cache_data  # Caches data so it doesn't reload every time20def prepare_data():21    train_datagen = ImageDataGenerator(22        rescale=1./255, rotation_range=15, width_shift_range=0.1,23        height_shift_range=0.1, shear_range=0.1, zoom_range=0.2,24        horizontal_flip=True, validation_split=0.225    )26    train_data = train_datagen.flow_from_directory(27        data_dir, target_size=img_size, batch_size=batch_size,28        class_mode='categorical', subset='training'29    )30    validation_data = train_datagen.flow_from_directory(31        data_dir, target_size=img_size, batch_size=batch_size,32        class_mode='categorical', subset='validation'33    )34    return train_data, validation_data35 36train_data, validation_data = prepare_data()37 38# Display Class Labels39class_labels = {v: k for k, v in train_data.class_indices.items()}40st.subheader("Class Labels:")41st.write(class_labels)42 43# Define model44def create_model():45    model = Sequential([46        Conv2D(32, (3, 3), activation='relu', input_shape=(img_size[0], img_size[1], 3)),47        MaxPooling2D((2, 2)),48        Conv2D(64, (3, 3), activation='relu'),49        MaxPooling2D((2, 2)),50        Conv2D(128, (3, 3), activation='relu'),51        MaxPooling2D((2, 2)),52        Flatten(),53        Dense(512, activation='relu'),54        Dropout(0.5),55        Dense(len(train_data.class_indices), activation='softmax')56    ])57    return model58 59model = create_model()60 61# Display Model Summary62st.subheader("Model Summary:")63model_summary = []64model.summary(print_fn=lambda x: model_summary.append(x))65st.text("\n".join(model_summary))66 67# Compile and Train the Model68if st.button("Start Training"):69    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])70    early_stop = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)71    72    history = model.fit(73        train_data, epochs=30, validation_data=validation_data, callbacks=[early_stop]74    )75    76    # Display training accuracy and loss77    st.subheader("Training Progress:")78    st.line_chart({"Training Accuracy": history.history["accuracy"], "Validation Accuracy": history.history["val_accuracy"]})79    st.line_chart({"Training Loss": history.history["loss"], "Validation Loss": history.history["val_loss"]})80 81    st.success("Training complete!")82