Update app.py

app.py

@@ -495,12 +495,12 @@ def _ensure_checkpoint(model_key: str, out_dir: str) -> str | None:
             pass
         return None
 
-# --- NEW: robust include absolutizer (no raw-text rewriting) ------------------
+# --- include absolutizer ------------------------------------------------------
 def _absify_any_paths_deep(node, base_dir, include_keys=("base", "_base_", "BASE", "BASE_YAML",
                                                          "includes", "include", "BASES", "__include__")):
     """
-    Walk dict/list; if a string looks like a relative YAML include (../*.yml/.yaml) or
-    appears under any of the known include keys, make it absolute against base_dir.
+    Walk dict/list; for known include keys or strings that look like ../*.yml/.yaml,
+    make them absolute against base_dir.
     """
     def _absify(s: str) -> str:
         if os.path.isabs(s):
@@ -510,7 +510,6 @@ def _absify_any_paths_deep(node, base_dir, include_keys=("base", "_base_", "BASE
         return s
 
     if isinstance(node, dict):
-        # First, handle explicit include keys
         for k in list(node.keys()):
             v = node[k]
             if k in include_keys:
@@ -518,13 +517,11 @@ def _absify_any_paths_deep(node, base_dir, include_keys=("base", "_base_", "BASE
                     node[k] = _absify(v)
                 elif isinstance(v, list):
                     node[k] = [_absify(x) if isinstance(x, str) else x for x in v]
-        # Recurse and also absify stray string values that look like includes
         for k, v in list(node.items()):
             if isinstance(v, (dict, list)):
                 _absify_any_paths_deep(v, base_dir, include_keys)
             elif isinstance(v, str):
                 node[k] = _absify(v)
-
     elif isinstance(node, list):
         for i, v in enumerate(list(node)):
             if isinstance(v, (dict, list)):
@@ -532,6 +529,50 @@ def _absify_any_paths_deep(node, base_dir, include_keys=("base", "_base_", "BASE
             elif isinstance(v, str):
                 node[i] = _absify(v)
 
+# --- NEW: safe model field setters --------------------------------------------
+def _set_num_classes_safely(cfg: dict, n: int):
+    """
+    Set class count without breaking templates that use `model: "RTDETR"` indirection.
+    """
+    def set_num_classes(node):
+        if not isinstance(node, dict):
+            return False
+        if "num_classes" in node:
+            node["num_classes"] = int(n)
+            return True
+        for k, v in node.items():
+            if isinstance(v, dict) and set_num_classes(v):
+                return True
+        return False
+
+    m = cfg.get("model", None)
+    if isinstance(m, dict):
+        if not set_num_classes(m):
+            m["num_classes"] = int(n)
+        return
+
+    if isinstance(m, str):
+        block = cfg.get(m, None)
+        if isinstance(block, dict):
+            if not set_num_classes(block):
+                block["num_classes"] = int(n)
+            return
+
+    cfg["num_classes"] = int(n)  # last resort
+
+def _maybe_set_model_field(cfg: dict, key: str, value):
+    """
+    Place fields like 'pretrain' under the proper model dict, respecting string indirection.
+    """
+    m = cfg.get("model", None)
+    if isinstance(m, dict):
+        m[key] = value
+        return
+    if isinstance(m, str) and isinstance(cfg.get(m), dict):
+        cfg[m][key] = value
+        return
+    cfg[key] = value  # fallback
+
 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):
@@ -539,7 +580,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
 
     template_dir = os.path.dirname(base_cfg_path)
 
-    # Load YAML directly (no raw-text editing), then absolutize known include keys and any '../*.yml'
+    # Load YAML then absolutize include-like paths
     with open(base_cfg_path, "r", encoding="utf-8") as f:
         cfg = yaml.safe_load(f)
     _absify_any_paths_deep(cfg, template_dir)
@@ -555,109 +596,81 @@ 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)),
     }
 
-    # dataset block
-    for root_key in ["dataset", "data"]:
-        if root_key in cfg and isinstance(cfg[root_key], dict):
-            ds = cfg[root_key]
-            for split, jf, ip in [
-                ("train", "train_json", "train_img"),
-                ("val",   "val_json",   "val_img"),
-                ("test",  "test_json",  "test_img"),
-            ]:
-                if split in ds and isinstance(ds[split], dict):
-                    node = ds[split]
-                    node["name"] = node.get("name", "coco")
-                    _set_first_existing_key(
-                        node,
-                        keys=["ann_file", "ann_path", "annotation", "annotations"],
-                        value=paths[jf],
-                        fallback_key="ann_file",
-                    )
-                    _set_first_existing_key(
-                        node,
-                        keys=["img_prefix", "img_dir", "image_root", "data_root"],
-                        value=paths[ip],
-                        fallback_key="img_prefix",
-                    )
-
-    # num_classes
-    def set_num_classes(node, n):
-        if not isinstance(node, dict):
-            return False
-        if "num_classes" in node:
-            node["num_classes"] = int(n)
-            return True
-        for k, v in node.items():
-            if isinstance(v, dict) and set_num_classes(v, n):
-                return True
-        return False
-
-    if "model" in cfg and isinstance(cfg["model"], dict):
-        if not set_num_classes(cfg["model"], class_count):
-            cfg["model"]["num_classes"] = int(class_count)
-    else:
-        cfg["model"] = {"num_classes": int(class_count)}
-
-    # epochs / imgsz
-    updated_epoch = False
-    for key in ["max_epoch", "epochs", "num_epochs"]:
+    # --- 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 ---
+    _set_num_classes_safely(cfg, int(class_count))
+
+    # --- epochs / imgsz ---
+    applied_epoch = False
+    for key in ("epoches", "max_epoch", "epochs", "num_epochs"):
         if key in cfg:
             cfg[key] = int(epochs)
-            updated_epoch = True
+            applied_epoch = True
             break
     if "solver" in cfg and isinstance(cfg["solver"], dict):
-        for key in ["max_epoch", "epochs", "num_epochs"]:
+        for key in ("epoches", "max_epoch", "epochs", "num_epochs"):
             if key in cfg["solver"]:
                 cfg["solver"][key] = int(epochs)
-                updated_epoch = True
+                applied_epoch = True
                 break
-    if not updated_epoch:
-        cfg["max_epoch"] = int(epochs)
+    if not applied_epoch:
+        cfg["epoches"] = int(epochs)  # common in this repo
 
-    for key in ["input_size", "img_size", "imgsz"]:
-        if key in cfg:
-            cfg[key] = int(imgsz)
-    if "input_size" not in cfg:
-        cfg["input_size"] = int(imgsz)
+    # image size knobs: unify on top-level input_size (respected by templates)
+    cfg["input_size"] = int(imgsz)
 
-    # lr / optimizer / batch
+    # --- 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"]:
+    for key in ("base_lr", "lr", "learning_rate"):
         if key in sol:
             sol[key] = float(lr)
             break
     else:
         sol["base_lr"] = float(lr)
     sol["optimizer"] = str(optimizer).lower()
-    if "train_dataloader" in cfg and isinstance(cfg["train_dataloader"], dict):
-        cfg["train_dataloader"]["batch_size"] = int(batch)
-    else:
+    if "train_dataloader" not in cfg or not isinstance(cfg["train_dataloader"], dict):
         sol["batch_size"] = int(batch)
 
     # output dir
     if "output_dir" in cfg:
         cfg["output_dir"] = paths["out_dir"]
-    elif "solver" in cfg:
-        sol["output_dir"] = paths["out_dir"]
     else:
-        cfg["output_dir"] = paths["out_dir"]
+        sol["output_dir"] = paths["out_dir"]
 
-    # Set pretrained weights if available; try common keys at top/model/solver
+    # pretrained weights in the right model block
     if pretrained_path:
-        _set_first_existing_key_deep(
-            cfg,
-            keys=["pretrain", "pretrained", "weight", "weights", "pretrained_path"],
-            value=os.path.abspath(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
     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 so lists don’t emit as inline [a, b, c] flow sequences
+    # Force block style for lists (no inline [a, b, c])
     class _NoFlowDumper(yaml.SafeDumper):
         pass
     def _repr_list_block(dumper, data):