| |
| |
| """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 |
|
|
| |
| |
| |
| |
| 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 |
|
|
