Update app.py

app.py

@@ -23,7 +23,7 @@ def train_and_export(dataset_file, model_name, num_classes, epochs, batch_size,
         zip_path = os.path.join(UPLOAD_DIR, f"{uid}.zip")
         
         # Copy uploaded file to our storage
-        shutil.copyfile(dataset_file, zip_path)
+        shutil.copyfile(dataset_file.name, zip_path)
         
         # Extract dataset
         extract_path = os.path.join(UPLOAD_DIR, uid)
@@ -37,7 +37,7 @@ def train_and_export(dataset_file, model_name, num_classes, epochs, batch_size,
         
         # 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"
+            return "Error: Dataset must contain 'train' and 'validation' folders", None, None, None
         
         # Create data generators
         train_datagen = ImageDataGenerator(
@@ -151,20 +151,19 @@ def train_and_export(dataset_file, model_name, num_classes, epochs, batch_size,
                 model_size += os.path.getsize(fp)
         model_size_mb = model_size / (1024 * 1024)
         
-        # Get class names
-        class_names = list(train_generator.class_indices.keys())
-        
-        # Prepare download links
-        download_info = f"""
+        # Prepare results
+        result_text = f"""
         ✅ Training completed successfully!
         ⏱️ Training time: {training_time:.2f} seconds
-        📊 Validation accuracy: {max(history.history['val_accuracy']):.4f}
+        📊 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 download_info, h5_path, savedmodel_path, tfjs_path
+        return result_text, h5_path, savedmodel_path, tfjs_path
     
     except Exception as e:
         return f"❌ Training failed: {str(e)}", None, None, None
@@ -172,7 +171,10 @@ def train_and_export(dataset_file, model_name, num_classes, epochs, batch_size,
 # Gradio interface
 with gr.Blocks(title="AI Image Classifier Trainer") as demo:
     gr.Markdown("# 🖼️ AI Image Classifier Trainer")
-    gr.Markdown("Upload your dataset (ZIP with train/validation folders), configure training, and download models in multiple formats.")
+    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():
@@ -186,16 +188,24 @@ with gr.Blocks(title="AI Image Classifier Trainer") as demo:
         
         with gr.Column():
             output = gr.Textbox(label="Training Results", interactive=False)
-            h5_download = gr.File(label="H5 Model Download", visible=False)
-            savedmodel_download = gr.File(label="SavedModel Download", visible=False)
-            tfjs_download = gr.File(label="TensorFlow.js Download", visible=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(results, h5_path, saved_path, tfjs_path):
-        downloads_visible = h5_path is not None
+    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.File(visible=downloads_visible, value=h5_path),
-            gr.File(visible=downloads_visible, value=saved_path),
-            gr.File(visible=downloads_visible, value=tfjs_path)
+            gr.Column(visible=False),
+            gr.File(value=None),
+            gr.File(value=None),
+            gr.File(value=None)
         )
     
     train_btn.click(
@@ -205,7 +215,7 @@ with gr.Blocks(title="AI Image Classifier Trainer") as demo:
     ).then(
         fn=toggle_downloads,
         inputs=[output, h5_download, savedmodel_download, tfjs_download],
-        outputs=[h5_download, savedmodel_download, tfjs_download]
+        outputs=[download_col, h5_download, savedmodel_download, tfjs_download]
     )
 
 # Launch settings for Hugging Face Spaces
@@ -214,5 +224,5 @@ if __name__ == "__main__":
         server_name="0.0.0.0",
         server_port=7860,
         share=False,
-        max_file_size=100  # 100MB file size limit
+        max_file_size="100mb"  # Allows 100MB file uploads
     )