Update app.py

app.py

@@ -573,6 +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 -----------------
 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):
@@ -585,6 +586,9 @@ 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
+    cfg["sync_bn"] = False
+
     ann_dir = os.path.join(merged_dir, "annotations")
     paths = {
         "train_json": os.path.abspath(os.path.join(ann_dir, "instances_train.json")),
@@ -596,83 +600,107 @@ 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)),
     }
 
-    # --- Rewrite dataloaders to use your dataset ---
-    def _patch_dl(dl_key, img_key, json_key):
-        if dl_key in cfg and isinstance(cfg[dl_key], dict):
-            ds = cfg[dl_key].get("dataset", {})
-            if isinstance(ds, dict):
-                if "img_folder" in ds: ds["img_folder"] = paths[img_key]
-                if "ann_file"  in ds: ds["ann_file"]  = paths[json_key]
-                # alternative key names occasionally used
-                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]
-                cfg[dl_key]["dataset"] = ds
-            # batch size here if present
-            if "batch_size" in cfg[dl_key]:
-                cfg[dl_key]["batch_size"] = int(batch)
-
-    _patch_dl("train_dataloader", "train_img", "train_json")
-    _patch_dl("val_dataloader",   "val_img",   "val_json")
-    _patch_dl("test_dataloader",  "test_img",  "test_json")
-
-    # --- classes ---
+    # 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
+    def ensure_and_patch_dl(dl_key, img_key, json_key, default_shuffle):
+        block = cfg.get(dl_key)
+        if not isinstance(block, dict):
+            block = {
+                "type": "DataLoader",
+                "dataset": {
+                    "type": "CocoDetection",
+                    "img_folder": paths[img_key],
+                    "ann_file": paths[json_key],
+                    "return_masks": False,
+                    "transforms": {
+                        "type": "Compose",
+                        "ops": [
+                            {"type": "Resize", "size": [int(imgsz), int(imgsz)]},
+                            {"type": "ConvertPILImage", "dtype": "float32", "scale": True},
+                        ],
+                    },
+                },
+                "shuffle": bool(default_shuffle),
+                "num_workers": 2,
+                "drop_last": bool(dl_key == "train_dataloader"),
+                "collate_fn": {"type": "BatchImageCollateFuncion"},
+                "total_batch_size": int(batch),
+            }
+            cfg[dl_key] = block
+        ds = block.get("dataset", {})
+        if isinstance(ds, dict):
+            ds["img_folder"] = paths[img_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))
+        block.setdefault("drop_last", bool(dl_key == "train_dataloader"))
+
+    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:
+    # ensure_and_patch_dl("test_dataloader",  "test_img",  "test_json",  default_shuffle=False)
+
+    # Classes
     _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)  # common in this repo
-
-    # image size knobs: unify on top-level input_size (respected by templates)
+        cfg["epoches"] = int(epochs)
     cfg["input_size"] = int(imgsz)
 
-    # --- lr / optimizer / batch fallbacks ---
+    # LR / optimizer / batch fallbacks
     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 the right model block
+    # Pretrained weights in correct 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 internal relative references still make sense
+    # Save near the template so any remaining relative includes still resolve
     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): pass
     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)
@@ -970,66 +998,4 @@ 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.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)
+                class_df = gr