Create app.py

app.py

@@ -0,0 +1,100 @@
+import gradio as gr
+import zipfile, os, uuid, shutil, subprocess, sys, time
+import tensorflow as tf
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+import numpy as np
+
+# Directory constants
+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_zip, model_name, num_classes, epochs, batch_size, image_size):
+    # Save upload
+    uid = str(uuid.uuid4())
+    zip_path = os.path.join(UPLOAD_DIR, f"{uid}.zip")
+    with open(zip_path, "wb") as f:
+        f.write(dataset_zip.read())
+    extract_path = os.path.join(UPLOAD_DIR, uid)
+    with zipfile.ZipFile(zip_path, "r") as z:
+        z.extractall(extract_path)
+    train_dir = os.path.join(extract_path, "train")
+    val_dir   = os.path.join(extract_path, "validation")
+    # Data gens
+    train_gen = ImageDataGenerator(rescale=1/255,
+                                   rotation_range=20,
+                                   width_shift_range=0.2,
+                                   height_shift_range=0.2,
+                                   horizontal_flip=True,
+                                   zoom_range=0.2) \
+        .flow_from_directory(train_dir, target_size=(image_size, image_size),
+                             batch_size=batch_size, class_mode="categorical")
+    val_gen = ImageDataGenerator(rescale=1/255) \
+        .flow_from_directory(val_dir, target_size=(image_size, image_size),
+                             batch_size=batch_size, class_mode="categorical")
+    # Build model (e.g., simplified ResNet-inspired CNN)
+    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.Conv2D(64,3,activation="relu"),
+        tf.keras.layers.BatchNormalization(),
+        tf.keras.layers.MaxPooling2D(),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(128,activation="relu"),
+        tf.keras.layers.Dropout(0.5),
+        tf.keras.layers.Dense(num_classes,activation="softmax")
+    ])
+    model.compile(optimizer="adam", loss="categorical_crossentropy",
+                  metrics=["accuracy"])
+    # Train
+    start = time.time()
+    history = model.fit(train_gen,
+                        validation_data=val_gen,
+                        epochs=epochs,
+                        verbose=0)
+    elapsed = time.time() - start
+    # Save outputs
+    space_model_dir = os.path.join(MODEL_DIR, uid)
+    os.makedirs(space_model_dir, exist_ok=True)
+    h5_path = os.path.join(space_model_dir, f"{model_name}.h5")
+    model.save(h5_path)
+    savedmodel_path = os.path.join(space_model_dir, "savedmodel")
+    model.save(savedmodel_path)
+    tfjs_dir = os.path.join(space_model_dir, "tfjs")
+    subprocess.run([
+        sys.executable, "-m", "tensorflowjs_converter",
+        "--input_format=tf_saved_model",
+        savedmodel_path,
+        tfjs_dir
+    ], check=True)
+    size_mb = sum(os.path.getsize(os.path.join(dp,fn))
+                  for dp,_,fns in os.walk(space_model_dir)
+                  for fn in fns) / (1024*1024)
+    return (f"✅ Trained in {elapsed:.1f}s\n"
+            f"Final val acc: {max(history.history['val_accuracy']):.4f}\n"
+            f"Model size: {size_mb:.2f} MB\n"
+            f"Download H5: {h5_path}\n"
+            f"Download SavedModel: {savedmodel_path}\n"
+            f"Download TF.js: {tfjs_dir}")
+
+# Gradio interface
+demo = gr.Interface(
+    fn=train_and_export,
+    inputs=[
+        gr.File(label="Dataset ZIP"),
+        gr.Textbox(label="Model Name", value="my_model"),
+        gr.Slider(2, 100, value=5, label="Number of Classes"),
+        gr.Slider(1, 200, value=50, step=1, label="Epochs"),
+        gr.Radio([16,32,64,128], value=32, label="Batch Size"),
+        gr.Radio([128,224,256], value=224, label="Image Size")
+    ],
+    outputs="text",
+    title="AI Image Classifier Trainer",
+    description="Upload train/validation dataset, train a CNN, and download H5, SavedModel, or TF.js."
+)
+
+if __name__ == "__main__":
+    demo.launch()