import gradio as gr import zipfile import os import uuid import shutil import subprocess import sys import time import tensorflow as tf from tensorflow.keras.preprocessing.image import ImageDataGenerator import numpy as np # Directory setup UPLOAD_DIR = "uploads" MODEL_DIR = "models" os.makedirs(UPLOAD_DIR, exist_ok=True) os.makedirs(MODEL_DIR, exist_ok=True) def train_and_export(dataset_file, model_name, num_classes, epochs, batch_size, image_size): try: # Generate unique ID for this training session uid = str(uuid.uuid4()) zip_path = os.path.join(UPLOAD_DIR, f"{uid}.zip") # Copy uploaded file to our storage shutil.copyfile(dataset_file.name, zip_path) # Extract dataset extract_path = os.path.join(UPLOAD_DIR, uid) os.makedirs(extract_path, exist_ok=True) with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall(extract_path) # Locate train and validation directories train_dir = os.path.join(extract_path, "train") val_dir = os.path.join(extract_path, "validation") # Verify dataset structure if not os.path.exists(train_dir) or not os.path.exists(val_dir): return "Error: Dataset must contain 'train' and 'validation' folders", None, None, None # Create data generators train_datagen = ImageDataGenerator( rescale=1./255, rotation_range=20, width_shift_range=0.2, height_shift_range=0.2, horizontal_flip=True, zoom_range=0.2 ) val_datagen = ImageDataGenerator(rescale=1./255) train_generator = train_datagen.flow_from_directory( train_dir, target_size=(image_size, image_size), batch_size=batch_size, class_mode='categorical' ) val_generator = val_datagen.flow_from_directory( val_dir, target_size=(image_size, image_size), batch_size=batch_size, class_mode='categorical' ) # Update num_classes based on actual data actual_classes = train_generator.num_classes if actual_classes != num_classes: num_classes = actual_classes # Build model model = tf.keras.Sequential([ tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(image_size, image_size, 3)), tf.keras.layers.BatchNormalization(), tf.keras.layers.MaxPooling2D(), tf.keras.layers.Dropout(0.25), tf.keras.layers.Conv2D(64, 3, activation='relu'), tf.keras.layers.BatchNormalization(), tf.keras.layers.MaxPooling2D(), tf.keras.layers.Dropout(0.25), tf.keras.layers.Conv2D(128, 3, activation='relu'), tf.keras.layers.BatchNormalization(), tf.keras.layers.MaxPooling2D(), tf.keras.layers.Dropout(0.25), tf.keras.layers.Flatten(), tf.keras.layers.Dense(256, activation='relu'), tf.keras.layers.BatchNormalization(), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(num_classes, activation='softmax') ]) model.compile( optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'] ) # Train model start_time = time.time() history = model.fit( train_generator, steps_per_epoch=train_generator.samples // train_generator.batch_size, epochs=epochs, validation_data=val_generator, validation_steps=val_generator.samples // val_generator.batch_size, verbose=0 ) training_time = time.time() - start_time # Save models model_dir = os.path.join(MODEL_DIR, uid) os.makedirs(model_dir, exist_ok=True) # Save H5 model h5_path = os.path.join(model_dir, f"{model_name}.h5") model.save(h5_path) # Save SavedModel savedmodel_path = os.path.join(model_dir, "savedmodel") model.save(savedmodel_path) # Convert to TensorFlow.js tfjs_path = os.path.join(model_dir, "tfjs") try: subprocess.run([ "tensorflowjs_converter", "--input_format=tf_saved_model", savedmodel_path, tfjs_path ], check=True) except Exception: # Install tensorflowjs if not available subprocess.run([sys.executable, "-m", "pip", "install", "tensorflowjs"], check=True) subprocess.run([ "tensorflowjs_converter", "--input_format=tf_saved_model", savedmodel_path, tfjs_path ], check=True) # Calculate model size model_size = 0 for dirpath, _, filenames in os.walk(model_dir): for f in filenames: fp = os.path.join(dirpath, f) model_size += os.path.getsize(fp) model_size_mb = model_size / (1024 * 1024) # Prepare results result_text = f""" ✅ Training completed successfully! ⏱️ Training time: {training_time:.2f} seconds 📊 Best validation accuracy: {max(history.history['val_accuracy']):.4f} 📦 Model size: {model_size_mb:.2f} MB 🗂️ Number of classes: {num_classes} Download links available below ⬇️ """ # Return paths for download return result_text, h5_path, savedmodel_path, tfjs_path except Exception as e: return f"❌ Training failed: {str(e)}", None, None, None # Gradio interface with gr.Blocks(title="AI Image Classifier Trainer") as demo: gr.Markdown("# 🖼️ AI Image Classifier Trainer") gr.Markdown(""" Upload your dataset (ZIP file containing `train/` and `validation/` folders), configure training parameters, and download models in multiple formats. """) with gr.Row(): with gr.Column(): dataset = gr.File(label="Dataset ZIP File", file_types=[".zip"]) model_name = gr.Textbox(label="Model Name", value="my_classifier") num_classes = gr.Slider(2, 100, value=5, step=1, label="Number of Classes") epochs = gr.Slider(5, 200, value=30, step=1, label="Training Epochs") batch_size = gr.Radio([16, 32, 64], value=32, label="Batch Size") image_size = gr.Radio([128, 224, 256], value=224, label="Image Size (px)") train_btn = gr.Button("🚀 Train Model", variant="primary") with gr.Column(): output = gr.Textbox(label="Training Results", interactive=False) with gr.Column(visible=False) as download_col: h5_download = gr.File(label="H5 Model Download") savedmodel_download = gr.File(label="SavedModel Download") tfjs_download = gr.File(label="TensorFlow.js Download") def toggle_downloads(result, h5_path, saved_path, tfjs_path): if h5_path: return ( gr.Column(visible=True), gr.File(value=h5_path), gr.File(value=saved_path), gr.File(value=tfjs_path) ) return ( gr.Column(visible=False), gr.File(value=None), gr.File(value=None), gr.File(value=None) ) train_btn.click( fn=train_and_export, inputs=[dataset, model_name, num_classes, epochs, batch_size, image_size], outputs=[output, h5_download, savedmodel_download, tfjs_download] ).then( fn=toggle_downloads, inputs=[output, h5_download, savedmodel_download, tfjs_download], outputs=[download_col, h5_download, savedmodel_download, tfjs_download] ) # Launch settings for Hugging Face Spaces if __name__ == "__main__": demo.launch( server_name="0.0.0.0", server_port=7860, share=False, max_file_size="100mb" # Allows 100MB file uploads )