import gradio as gr import zipfile import os import uuid import shutil import subprocess import sys import time from PIL import Image 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: uid = str(uuid.uuid4()) zip_path = os.path.join(UPLOAD_DIR, f"{uid}.zip") shutil.copyfile(dataset_file.name, zip_path) 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) train_dir = os.path.join(extract_path, "train") val_dir = os.path.join(extract_path, "validation") # 🛠 Auto-generate folders and dummy images if missing if not os.path.exists(train_dir) or not os.path.exists(val_dir): os.makedirs(train_dir, exist_ok=True) os.makedirs(val_dir, exist_ok=True) for split_dir in [train_dir, val_dir]: for class_name in ["class_a", "class_b"]: class_path = os.path.join(split_dir, class_name) os.makedirs(class_path, exist_ok=True) # Generate 2 dummy images per class for i in range(2): img = Image.new('RGB', (image_size, image_size), color=(i * 50, 100, 150)) img.save(os.path.join(class_path, f"sample_{i}.jpg")) # 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' ) actual_classes = train_generator.num_classes if actual_classes != num_classes: num_classes = actual_classes 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'] ) 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 model_dir = os.path.join(MODEL_DIR, uid) os.makedirs(model_dir, exist_ok=True) h5_path = os.path.join(model_dir, f"{model_name}.h5") model.save(h5_path) savedmodel_path = os.path.join(model_dir, "savedmodel") model.save(savedmodel_path) 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: 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) model_size = 0 for dirpath, _, filenames in os.walk(model_dir): for f in filenames: model_size += os.path.getsize(os.path.join(dirpath, f)) model_size_mb = model_size / (1024 * 1024) 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} """ 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 a ZIP of `train/` and `validation/`, or leave it empty to auto-generate dummy data.") 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] ) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860, share=False, max_file_size="100mb")