Update app.py

app.py

@@ -10,7 +10,7 @@ import matplotlib.pyplot as plt
 from roboflow import Roboflow
 from PIL import Image
 import torch
-from string import Template  # <-- NEW
+from string import Template  # <-- used by the shim
 
 # Quiet some noisy libs on Spaces (harmless locally)
 os.environ.setdefault("YOLO_CONFIG_DIR", "/tmp/Ultralytics")
@@ -459,14 +459,14 @@ def _install_supervisely_logger_shim():
     """))
     return str(root)
 
-# ---- NEW: kwargs-aware sitecustomize shim (safe, non-invasive) ---------------
+# ---- NEW: signature-agnostic sitecustomize shim (safe, non-invasive) ---------
 def _install_workspace_shim_v3(dest_dir: str, module_default: str = "rtdetrv2_pytorch.src"):
     """
-    Writes a sitecustomize.py that monkeypatches
-    rtdetrv2_pytorch.src.core.workspace.create so it works with the real signature:
-        def create(name, **kwargs):
-    The shim ensures kwargs['cfg'] is a dict, then guarantees cfg['_pymodule'] is a *module object*.
-    `dest_dir` MUST be on sys.path at interpreter startup (we'll prepend it to PYTHONPATH).
+    sitecustomize shim that patches rtdetrv2_pytorch.src.core.workspace.create
+    to be robust for BOTH call styles:
+        - create(name, cfg, **extras)           (positional cfg)
+        - create(name, cfg=<dict>, **extras)    (keyword cfg)
+    It guarantees cfg is a dict and cfg['_pymodule'] is a *module object*.
     """
     os.makedirs(dest_dir, exist_ok=True)
     sc_path = os.path.join(dest_dir, "sitecustomize.py")
@@ -475,51 +475,51 @@ def _install_workspace_shim_v3(dest_dir: str, module_default: str = "rtdetrv2_py
     tmpl = Template(r"""
 import os, importlib, types
 try:
-    mod_default = os.environ.get("RTDETR_PYMODULE", "$module_default") or "$module_default"
+    MOD_DEFAULT = os.environ.get("RTDETR_PYMODULE", "$module_default") or "$module_default"
     ws_mod = importlib.import_module("rtdetrv2_pytorch.src.core.workspace")
     _orig_create = ws_mod.create
 
-    def _ensure_pymodule_object(cfg):
-        pm = None
-        try:
-            pm = cfg.get("_pymodule", None)
-        except Exception:
-            pm = None
-        if isinstance(pm, str) or pm is None:
-            name = pm.strip() if isinstance(pm, str) and pm.strip() else mod_default
-            try:
-                mod = importlib.import_module(name)
-            except Exception:
-                mod = importlib.import_module(mod_default)
-            try:
-                cfg["_pymodule"] = mod
-            except Exception:
-                pass
-            return mod
+    def _ensure_pymodule_object(cfg: dict):
+        pm = cfg.get("_pymodule", None)
         if isinstance(pm, types.ModuleType):
             return pm
+        name = (pm or "").strip() if isinstance(pm, str) else MOD_DEFAULT
+        if not name:
+            name = MOD_DEFAULT
         try:
-            mod = importlib.import_module(mod_default)
-            cfg["_pymodule"] = mod
-            return mod
+            mod = importlib.import_module(name)
         except Exception:
-            return pm
+            mod = importlib.import_module(MOD_DEFAULT)
+        cfg["_pymodule"] = mod
+        return mod
+
+    def create(name, *args, **kwargs):
+        # Accept both positional and keyword cfg
+        if args:
+            args = list(args)
+            cfg = args[0]
+        else:
+            cfg = kwargs.get("cfg", None)
 
-    def create(name, **kwargs):
-        cfg = kwargs.get("cfg")
         if not isinstance(cfg, dict):
             cfg = {} if cfg is None else dict(cfg)
-            kwargs["cfg"] = cfg
+
         _ensure_pymodule_object(cfg)
-        return _orig_create(name, **kwargs)
+
+        if args:
+            args[0] = cfg
+            args = tuple(args)
+        else:
+            kwargs["cfg"] = cfg
+
+        return _orig_create(name, *args, **kwargs)
 
     ws_mod.create = create
 except Exception:
-    # never block training if shim fails
+    # Never block training on shim failure
     pass
 """)
     code = tmpl.substitute(module_default=module_default)
-
     with open(sc_path, "w", encoding="utf-8") as f:
         f.write(code)
     return sc_path
@@ -650,7 +650,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
 
     # Ensure the runtime knows which Python module hosts builders
     cfg["task"] = cfg.get("task", "detection")
-    cfg["_pymodule"] = cfg.get("_pymodule", "rtdetrv2_pytorch.src")  # <= HINT for loader
+    cfg["_pymodule"] = cfg.get("_pymodule", "rtdetrv2_pytorch.src")  # <= hint for loader
 
     # Disable SyncBN for single GPU/CPU runs; guard DDP flags
     cfg["sync_bn"] = False
@@ -696,7 +696,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
                 "shuffle": bool(default_shuffle),
                 "num_workers": 2,
                 "drop_last": bool(dl_key == "train_dataloader"),
-                "collate_fn": {"type": "BatchImageCollateFunction"},  # <-- FIXED name
+                "collate_fn": {"type": "BatchImageCollateFunction"},  # correct spelling
                 "total_batch_size": int(batch),
             }
             cfg[dl_key] = block
@@ -716,6 +716,16 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         block.setdefault("shuffle", bool(default_shuffle))
         block.setdefault("drop_last", bool(dl_key == "train_dataloader"))
 
+        # ---- FORCE-FIX collate name even if it existed already
+        cf = block.get("collate_fn", {})
+        if isinstance(cf, dict):
+            t = str(cf.get("type", ""))
+            if t.lower() == "batchimagecollatefuncion" or "Funcion" in t:
+                cf["type"] = "BatchImageCollateFunction"
+            block["collate_fn"] = cf
+        else:
+            block["collate_fn"] = {"type": "BatchImageCollateFunction"}
+
     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)