Buckets:
| """Untouched official pretrained models; adaptation is explicitly separate.""" | |
| import hashlib | |
| import json | |
| import os | |
| import urllib.request | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parents[1] | |
| FRAMEWORK_VERSION = "8.4.185" | |
| CHECKPOINTS = { | |
| "yolo26n-obb.pt": "6f51c78197aacda4a33be77294065a9001675fb893f56227a179731b53dbd2b0", | |
| "yolo26n-sem-ade20k.pt": "68ac71bef2868c987fff7cd7c49cb656922f6d8447ad21af3893305577435000", | |
| } | |
| ASSET_BASE = "https://github.com/ultralytics/assets/releases/download/v8.4.0/" | |
| DOTA_NAMES = ( | |
| "plane", "baseball-diamond", "bridge", "ground-track-field", "small-vehicle", | |
| "large-vehicle", "ship", "tennis-court", "basketball-court", "storage-tank", | |
| "soccer-ball-field", "roundabout", "harbor", "swimming-pool", "helicopter", | |
| "container-crane", | |
| ) | |
| DETECTION_ALIASES = { | |
| "plane": 6, "helicopter": 6, "ship": 7, "small-vehicle": 4, | |
| "large-vehicle": 5, "bridge": 3, "storage-tank": 8, | |
| } | |
| # Broad land-cover classes must not be approximated by tree/field/earth/land. | |
| SEGMENTATION_ALIASES = {"building": 1, "road": 2, "water": 12} | |
| def configure_local_settings() -> Path: | |
| directory = ROOT / "training" / "cache" / "ultralytics" | |
| os.environ["YOLO_CONFIG_DIR"] = str(directory) | |
| os.environ["YOLO_AUTOINSTALL"] = "false" | |
| os.environ["YOLO_OFFLINE"] = "true" | |
| settings_path = directory / "Ultralytics" / "settings.json" | |
| settings_path.parent.mkdir(parents=True, exist_ok=True) | |
| settings = { | |
| "settings_version": "0.0.8", "datasets_dir": str(ROOT / "datasets"), | |
| "weights_dir": str(ROOT / "training" / "cache" / "checkpoints"), | |
| "runs_dir": str(ROOT / "training" / "cache" / "runs"), | |
| "uuid": "local-only", "api_key": "", "openai_api_key": "", | |
| **{key: False for key in ( | |
| "sync", "clearml", "comet", "dvc", "mlflow", "raytune", | |
| "tensorboard", "wandb", "vscode_msg", "openvino_msg", | |
| )}, | |
| } | |
| settings_path.write_text(json.dumps(settings, indent=2), encoding="utf-8") | |
| return directory | |
| configure_local_settings() | |
| def sha256(path: Path) -> str: | |
| return hashlib.sha256(path.read_bytes()).hexdigest() | |
| def model_state_sha256(model) -> str: | |
| """Identity of tensor values, independent of checkpoint container metadata.""" | |
| digest = hashlib.sha256() | |
| for name, tensor in sorted(model.state_dict().items()): | |
| value = tensor.detach().cpu().contiguous() | |
| digest.update(json.dumps([name, str(value.dtype), list(value.shape)]).encode()) | |
| digest.update(value.numpy().tobytes()) | |
| return digest.hexdigest() | |
| def _checkpoint(cache_dir: Path, name: str) -> Path: | |
| cache_dir = cache_dir.resolve() | |
| if not cache_dir.is_relative_to(ROOT / "training" / "cache"): | |
| raise ValueError("checkpoints must remain in the worktree training cache") | |
| cache_dir.mkdir(parents=True, exist_ok=True) | |
| path = cache_dir / name | |
| if not path.exists(): | |
| partial = path.with_suffix(".download") | |
| try: | |
| with urllib.request.urlopen(ASSET_BASE + name, timeout=60) as response: | |
| partial.write_bytes(response.read()) | |
| if sha256(partial) != CHECKPOINTS[name]: | |
| raise ValueError(f"checkpoint hash mismatch: {name}") | |
| partial.replace(path) | |
| finally: | |
| partial.unlink(missing_ok=True) | |
| if sha256(path) != CHECKPOINTS[name]: | |
| raise ValueError(f"checkpoint hash mismatch: {name}") | |
| return path | |
| def _yolo_class(): | |
| configure_local_settings() | |
| import ultralytics | |
| from ultralytics import YOLO | |
| from ultralytics.utils import SETTINGS | |
| if ultralytics.__version__ != FRAMEWORK_VERSION: | |
| raise RuntimeError(f"requires pinned ultralytics {FRAMEWORK_VERSION}") | |
| if SETTINGS["sync"]: | |
| raise RuntimeError("cloud synchronization must be disabled before model loading") | |
| return YOLO | |
| def _load(cache_dir: Path, device: str, name: str, task: str): | |
| if device != "cpu" and not device.startswith("cuda"): | |
| raise ValueError("device must be cpu or cuda") | |
| if device.startswith("cuda"): | |
| import torch | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("CUDA device requested but CUDA is unavailable") | |
| path = _checkpoint(cache_dir, name) | |
| model = _yolo_class()(str(path)) | |
| if model.task != task: | |
| raise RuntimeError(f"checkpoint task {model.task!r} differs from required {task!r}") | |
| model.to(device) | |
| return model | |
| def load_pretrained_detector(cache_dir: Path, device: str = "cpu"): | |
| return _load(cache_dir, device, "yolo26n-obb.pt", "obb") | |
| def load_pretrained_segmenter(cache_dir: Path, device: str = "cpu"): | |
| return _load(cache_dir, device, "yolo26n-sem-ade20k.pt", "semantic") | |
| def public_mapping(names: dict[int, str], task: str) -> dict[int, int]: | |
| aliases = {"detection": DETECTION_ALIASES, "segmentation": SEGMENTATION_ALIASES}[task] | |
| normalized = {index: name.replace(" ", "-") if task == "detection" else name | |
| for index, name in names.items()} | |
| return {int(index): aliases[name] for index, name in normalized.items() if name in aliases} | |
| def detector_transfer_plan(names: dict[int, str]) -> dict: | |
| indices = {name.replace(" ", "-"): int(index) for index, name in names.items()} | |
| return { | |
| "target_names": list(DOTA_NAMES), | |
| "source_indices": {name: indices.get(name) for name in DOTA_NAMES}, | |
| "missing_classes": [name for name in DOTA_NAMES if name not in indices], | |
| "missing_policy": "retain caller-provided fresh initialization; never adapt baseline", | |
| } | |
| def transfer_head_rows(source, initialized, names: dict[int, str]): | |
| """Copy class-axis rows to a fresh head; absent categories stay initialized. | |
| For later adaptation only: callers supply each classifier tensor and its | |
| corresponding newly initialized 16-class tensor. Does not mutate a model. | |
| """ | |
| if len(source) != len(names) or len(initialized) != len(DOTA_NAMES): | |
| raise ValueError("class-axis sizes differ from checkpoint/target names") | |
| result = initialized.clone() if hasattr(initialized, "clone") else initialized.copy() | |
| for target_index, name in enumerate(DOTA_NAMES): | |
| source_index = detector_transfer_plan(names)["source_indices"][name] | |
| if source_index is not None: | |
| result[target_index] = source[source_index] | |
| return result | |
Xet Storage Details
- Size:
- 6.55 kB
- Xet hash:
- 206973751d68d1f72ef6671e8194e47fe07c453e2f28e87d7a2ebd6a43cfc104
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.