Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """One-time: extract prebackbone weights from a full YOLO best.pt checkpoint.""" | |
| from __future__ import annotations | |
| import argparse | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| # Allow running before vendor is installed | |
| _HERE = Path(__file__).resolve().parent | |
| for base in (_HERE / "vendor", _HERE.parent / "ultralytics"): | |
| if (base / "ultralytics" / "__init__.py").exists(): | |
| sys.path.insert(0, str(base)) | |
| break | |
| def _clean_state_dict(state: dict) -> dict: | |
| """Drop thop profiling keys (total_ops, total_params) not in nn.Module.""" | |
| skip = ("total_ops", "total_params") | |
| return {k: v for k, v in state.items() if not any(k == s or k.endswith(f".{s}") for s in skip)} | |
| def extract( | |
| full_ckpt: Path, | |
| out_path: Path, | |
| prebackbone_name: str | None = None, | |
| ) -> Path: | |
| from ultralytics.nn.tasks import torch_safe_load | |
| ckpt, _ = torch_safe_load(str(full_ckpt)) | |
| train_args = ckpt.get("train_args") or {} | |
| name = (prebackbone_name or train_args.get("prebackbone") or "A11_CA").upper() | |
| channels = int(train_args.get("channels", 3) or 3) | |
| model = (ckpt.get("ema") or ckpt.get("model")).float() | |
| pb = getattr(model, "prebackbone", None) | |
| if pb is None: | |
| raise RuntimeError(f"No prebackbone in {full_ckpt}") | |
| state = _clean_state_dict(pb.state_dict()) | |
| # Save only tensors + metadata (no YOLO backbone/head) | |
| payload = { | |
| "prebackbone": name, | |
| "channels": channels, | |
| "state_dict": state, | |
| "source_checkpoint": str(full_ckpt.resolve()), | |
| } | |
| out_path.parent.mkdir(parents=True, exist_ok=True) | |
| torch.save(payload, out_path) | |
| n_params = sum(t.numel() for t in payload["state_dict"].values()) | |
| size_mb = out_path.stat().st_size / (1024 * 1024) | |
| print(f"Saved {out_path} ({size_mb:.2f} MB, {n_params:,} parameters, type={name})") | |
| return out_path | |
| def main() -> None: | |
| p = argparse.ArgumentParser(description="Extract prebackbone-only weights from best.pt") | |
| p.add_argument( | |
| "--ckpt", | |
| type=Path, | |
| default=_HERE / "weights" / "best.pt", | |
| help="Full YOLO checkpoint (best.pt)", | |
| ) | |
| p.add_argument( | |
| "--out", | |
| type=Path, | |
| default=_HERE / "weights" / "prebackbone_a11_ca.pt", | |
| help="Output path for prebackbone-only weights", | |
| ) | |
| p.add_argument("--name", type=str, default=None, help="Override prebackbone type (default: from train_args)") | |
| args = p.parse_args() | |
| if not args.ckpt.exists(): | |
| raise SystemExit(f"Checkpoint not found: {args.ckpt}") | |
| extract(args.ckpt, args.out, args.name) | |
| if __name__ == "__main__": | |
| main() | |