File size: 3,652 Bytes
d469c19 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | # SPDX-License-Identifier: AGPL-3.0-only
# Safetensors loading added 2026-10-28; inference preprocessing is unchanged.
"""Strict, direct safetensors loader for this pinned RF-DETR Small release."""
from importlib.metadata import version
import json
from pathlib import Path
import warnings
import torch
from safetensors.torch import load_file
HERE = Path(__file__).resolve().parent
def load_model(weights_path=None, config_path=None, device='cpu'):
"""Rebuild Small, load every tensor strictly, and return its RF-DETR facade.
No checkpoint deserialization, temporary pickle, pretrained download, dtype
conversion, or inference optimization is used. The CLI supplies the existing
RGB preparation, tiling, boundary rechecks and NMS via inference_engine.py.
"""
weights_path = Path(weights_path) if weights_path is not None else HERE / 'model.safetensors'
config_path = Path(config_path) if config_path is not None else weights_path.parent / 'model_config.json'
if weights_path.suffix != '.safetensors':
raise ValueError('Use a .safetensors weights file.')
config = json.loads(config_path.read_text())
expected_config = {
'variant': 'RFDETRSmall', 'rfdetr_version': '1.11.2',
'resolution': 512, 'num_classes': 1, 'id2label': {'0': 'door'},
}
if config != expected_config:
raise ValueError('model_config.json does not describe this RF-DETR Small release.')
if version('rfdetr') != config['rfdetr_version']:
raise RuntimeError('Install the pinned rfdetr==1.11.2 inference dependencies.')
if device not in ('cpu', 'mps', 'cuda'):
raise ValueError('Choose cpu, mps or cuda.')
if device == 'mps' and not torch.backends.mps.is_available():
raise RuntimeError('MPS unavailable; use --device cpu.')
if device == 'cuda' and not torch.cuda.is_available():
raise RuntimeError('CUDA unavailable; use --device cpu.')
from rfdetr import RFDETRSmall
from rfdetr.config import PretrainWeightsCompatibilityWarning
# The pinned Small constructor uses bundled backbone configuration. Patch 16
# and positional grid 32 create the backbone locally, without DINOv2 weights.
# Its expected "initialised from scratch" warning is superseded by the strict
# complete-state load below; no randomly initialized weights remain in use.
with warnings.catch_warnings():
warnings.simplefilter('ignore', PretrainWeightsCompatibilityWarning)
model = RFDETRSmall(pretrain_weights=None, num_classes=config['num_classes'],
resolution=config['resolution'], device=device)
module = model.model.model
state = load_file(str(weights_path), device='cpu')
expected = module.state_dict()
missing = sorted(expected.keys() - state.keys())
unexpected = sorted(state.keys() - expected.keys())
if missing or unexpected:
raise RuntimeError(f'Weight key mismatch: missing={missing}, unexpected={unexpected}')
for name, tensor in state.items():
if tensor.shape != expected[name].shape or tensor.dtype != expected[name].dtype:
raise RuntimeError(f'Weight shape/dtype mismatch: {name}')
module.load_state_dict(state, strict=True)
loaded_state = module.state_dict()
if not all(torch.equal(tensor, loaded_state[name]) for name, tensor in state.items()):
raise RuntimeError('Loaded model tensors do not exactly match the safetensors file.')
module.eval()
model.model.class_names = [config['id2label'][str(i)] for i in range(config['num_classes'])]
model.model.args.class_names = model.model.class_names
return model
|