Update app.py

app.py

@@ -573,7 +573,7 @@ def _maybe_set_model_field(cfg: dict, key: str, value):
         return
     cfg[key] = value  # fallback
 
-# --- CRITICAL FIX: force custom dataloaders & disable sync_bn -----------------
+# --- CRITICAL: dataset override + include cleanup + sync_bn off ---------------
 def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
                       epochs, batch, imgsz, lr, optimizer, pretrained_path: str | None):
     if not base_cfg_path or not os.path.exists(base_cfg_path):
@@ -586,9 +586,16 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         cfg = yaml.safe_load(f)
     _absify_any_paths_deep(cfg, template_dir)
 
-    # Safer on single GPU/CPU
+    # Disable SyncBN for single GPU/CPU runs
     cfg["sync_bn"] = False
 
+    # Remove COCO dataset include so it can't override our dataset paths later
+    if "__include__" in cfg and isinstance(cfg["__include__"], list):
+        cfg["__include__"] = [
+            p for p in cfg["__include__"]
+            if not (isinstance(p, str) and "configs/dataset/coco" in p.replace("\\", "/"))
+        ]
+
     ann_dir = os.path.join(merged_dir, "annotations")
     paths = {
         "train_json": os.path.abspath(os.path.join(ann_dir, "instances_train.json")),
@@ -600,15 +607,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         "out_dir":    os.path.abspath(os.path.join("runs", "train", run_name)),
     }
 
-    # Remove COCO dataset include so it can't override our dataloaders
-    inc_key = "__include__"
-    if inc_key in cfg and isinstance(cfg[inc_key], list):
-        cfg[inc_key] = [
-            p for p in cfg[inc_key]
-            if not (isinstance(p, str) and "configs/dataset/coco" in p.replace("\\", "/"))
-        ]
-
-    # Helper to ensure & patch dataloaders
+    # Ensure/patch dataloaders to point to our dataset
     def ensure_and_patch_dl(dl_key, img_key, json_key, default_shuffle):
         block = cfg.get(dl_key)
         if not isinstance(block, dict):
@@ -634,15 +633,18 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
                 "total_batch_size": int(batch),
             }
             cfg[dl_key] = block
+
+        # Patch existing block
         ds = block.get("dataset", {})
         if isinstance(ds, dict):
             ds["img_folder"] = paths[img_key]
-            ds["ann_file"] = paths[json_key]
+            ds["ann_file"]   = paths[json_key]
             for k in ("img_dir", "image_root", "data_root"):
                 if k in ds: ds[k] = paths[img_key]
             for k in ("ann_path", "annotation", "annotations"):
                 if k in ds: ds[k] = paths[json_key]
             block["dataset"] = ds
+
         block["total_batch_size"] = int(batch)
         block.setdefault("num_workers", 2)
         block.setdefault("shuffle", bool(default_shuffle))
@@ -650,57 +652,62 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
 
     ensure_and_patch_dl("train_dataloader", "train_img", "train_json", default_shuffle=True)
     ensure_and_patch_dl("val_dataloader",   "val_img",   "val_json",   default_shuffle=False)
-    # Optional test loader if needed:
+    # Optional test loader
     # ensure_and_patch_dl("test_dataloader",  "test_img",  "test_json",  default_shuffle=False)
 
-    # Classes
+    # num classes (handles model: "RTDETR")
     _set_num_classes_safely(cfg, int(class_count))
 
-    # Epochs / imgsz
+    # epochs / imgsz
     applied_epoch = False
     for key in ("epoches", "max_epoch", "epochs", "num_epochs"):
         if key in cfg:
-            cfg[key] = int(epochs); applied_epoch = True; break
+            cfg[key] = int(epochs)
+            applied_epoch = True
+            break
     if "solver" in cfg and isinstance(cfg["solver"], dict):
         for key in ("epoches", "max_epoch", "epochs", "num_epochs"):
             if key in cfg["solver"]:
-                cfg["solver"][key] = int(epochs); applied_epoch = True; break
+                cfg["solver"][key] = int(epochs)
+                applied_epoch = True
+                break
     if not applied_epoch:
         cfg["epoches"] = int(epochs)
     cfg["input_size"] = int(imgsz)
 
-    # LR / optimizer / batch fallbacks
+    # lr / optimizer / batch
     if "solver" not in cfg or not isinstance(cfg["solver"], dict):
         cfg["solver"] = {}
     sol = cfg["solver"]
     for key in ("base_lr", "lr", "learning_rate"):
         if key in sol:
-            sol[key] = float(lr); break
+            sol[key] = float(lr)
+            break
     else:
         sol["base_lr"] = float(lr)
     sol["optimizer"] = str(optimizer).lower()
     if "train_dataloader" not in cfg or not isinstance(cfg["train_dataloader"], dict):
         sol["batch_size"] = int(batch)
 
-    # Output dir
+    # output dir
     if "output_dir" in cfg:
         cfg["output_dir"] = paths["out_dir"]
     else:
         sol["output_dir"] = paths["out_dir"]
 
-    # Pretrained weights in correct block
+    # pretrained weights in the right model block
     if pretrained_path:
         p = os.path.abspath(pretrained_path)
         _maybe_set_model_field(cfg, "pretrain", p)
         _maybe_set_model_field(cfg, "pretrained", p)
 
-    # Save near the template so any remaining relative includes still resolve
+    # Save near the template so internal relative references still make sense
     cfg_out_dir = os.path.join(template_dir, "generated")
     os.makedirs(cfg_out_dir, exist_ok=True)
     out_path = os.path.join(cfg_out_dir, f"{run_name}.yaml")
 
     # Force block style for lists (no inline [a, b, c])
-    class _NoFlowDumper(yaml.SafeDumper): pass
+    class _NoFlowDumper(yaml.SafeDumper): ...
     def _repr_list_block(dumper, data):
         return dumper.represent_sequence('tag:yaml.org,2002:seq', data, flow_style=False)
     _NoFlowDumper.add_representer(list, _repr_list_block)
@@ -998,4 +1005,66 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
         with gr.TabItem("2. Manage & Merge"):
             gr.Markdown("Rename/merge/remove classes and set per-class image caps. Then finalize.")
             with gr.Row():
-                class_df = gr
+                class_df = gr.DataFrame(headers=["Original Name","Rename To","Max Images","Remove"],
+                                        datatype=["str","str","number","bool"], label="Class Config", interactive=True, scale=3)
+                with gr.Column(scale=1):
+                    class_count_summary_df = gr.DataFrame(label="Merged Class Counts Preview",
+                                                          headers=["Final Class Name","Est. Total Images"], interactive=False)
+                    update_counts_btn = gr.Button("Update Counts")
+            finalize_btn = gr.Button("Finalize Merged Dataset", variant="primary")
+            finalize_status = gr.Textbox(label="Status", interactive=False)
+
+        with gr.TabItem("3. Configure & Train"):
+            gr.Markdown("Pick RT-DETRv2 model, set hyper-params, press Start.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    model_dd = gr.Dropdown(choices=[k for k,_ in MODEL_CHOICES], value=DEFAULT_MODEL_KEY,
+                                           label="Model (RT-DETRv2)")
+                    run_name_tb = gr.Textbox(label="Run Name", value="rtdetrv2_run_1")
+                    epochs_sl = gr.Slider(1, 500, 100, step=1, label="Epochs")
+                    batch_sl = gr.Slider(1, 64, 16, step=1, label="Batch Size")
+                    imgsz_num = gr.Number(label="Image Size", value=640)
+                    lr_num = gr.Number(label="Learning Rate", value=0.001)
+                    opt_dd = gr.Dropdown(["Adam","AdamW","SGD"], value="Adam", label="Optimizer")
+                    train_btn = gr.Button("Start Training", variant="primary")
+                with gr.Column(scale=2):
+                    train_status = gr.Textbox(label="Live Logs (tail)", interactive=False, lines=12)
+                    loss_plot = gr.Plot(label="Loss")
+                    map_plot = gr.Plot(label="mAP")
+                    final_model_file = gr.File(label="Download Trained Checkpoint", interactive=False, visible=False)
+
+        with gr.TabItem("4. Upload Model"):
+            gr.Markdown("Optionally push your checkpoint to Hugging Face / GitHub.")
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("**Hugging Face**")
+                    hf_token = gr.Textbox(label="HF Token", type="password")
+                    hf_repo  = gr.Textbox(label="HF Repo (user/repo)")
+                with gr.Column():
+                    gr.Markdown("**GitHub**")
+                    gh_token = gr.Textbox(label="GitHub PAT", type="password")
+                    gh_repo  = gr.Textbox(label="GitHub Repo (user/repo)")
+            upload_btn = gr.Button("Upload", variant="primary")
+            with gr.Row():
+                hf_status = gr.Textbox(label="Hugging Face Status", interactive=False)
+                gh_status = gr.Textbox(label="GitHub Status", interactive=False)
+
+    load_btn.click(load_datasets_handler, [rf_api_key, rf_url_file],
+                   [dataset_status, dataset_info_state, class_df])
+    update_counts_btn.click(update_class_counts_handler, [class_df, dataset_info_state],
+                            [class_count_summary_df])
+    finalize_btn.click(finalize_handler, [dataset_info_state, class_df],
+                       [finalize_status, final_dataset_path_state])
+    train_btn.click(training_handler,
+                    [final_dataset_path_state, model_dd, run_name_tb, epochs_sl, batch_sl, imgsz_num, lr_num, opt_dd],
+                    [train_status, loss_plot, map_plot, final_model_file])
+    upload_btn.click(upload_handler, [final_model_file, hf_token, hf_repo, gh_token, gh_repo],
+                     [hf_status, gh_status])
+
+if __name__ == "__main__":
+    try:
+        ts = find_training_script(REPO_DIR)
+        logging.info(f"Startup check — training script at: {ts}")
+    except Exception as e:
+        logging.warning(f"Startup training-script check failed: {e}")
+    app.launch(debug=True)