Update app.py

app.py

@@ -36,7 +36,7 @@ SESSION = {
 
 
 def run_training(train_files, test_files, input_size, epochs, batch_size,
-                  learning_rate, background_weight):
+                  learning_rate, background_weight, optimizer_name):
     if not train_files or not test_files:
         yield "Please select both a train folder and a test folder first.", None, None
         return
@@ -93,7 +93,7 @@ def run_training(train_files, test_files, input_size, epochs, batch_size,
     epoch_logs = []
     history = fu.train_model(
         model, train_ds, val_ds, int(epochs), float(learning_rate),
-        float(background_weight), epoch_logs,
+        float(background_weight), epoch_logs, optimizer_name=optimizer_name,
     )
     log_lines.extend(epoch_logs)
     yield "\n".join(log_lines), None, None
@@ -196,6 +196,9 @@ with gr.Blocks(title="Crow FOMO Trainer") as demo:
                 0.05, 1.0, value=0.3, step=0.05,
                 label="Background loss weight (lower = focus more on crows, helps with class imbalance)",
             )
+            optimizer_name = gr.Dropdown(
+                ["adamw", "adam", "sgd"], value="adamw", label="Optimizer",
+            )
         train_btn = gr.Button("Start Training", variant="primary")
         train_log = gr.Textbox(label="Training log", lines=15)
         loss_plot = gr.Image(label="Loss curve")
@@ -203,7 +206,7 @@ with gr.Blocks(title="Crow FOMO Trainer") as demo:
 
         train_btn.click(
             run_training,
-            inputs=[train_files, test_files, input_size, epochs, batch_size, learning_rate, background_weight],
+            inputs=[train_files, test_files, input_size, epochs, batch_size, learning_rate, background_weight, optimizer_name],
             outputs=[train_log, loss_plot, tflite_out],
         )