Coverage for collectra/tasks/object_detection/rfdetr_channels.py: 92.59%

1"""Load and train RF-DETR checkpoints whose input layer has more than three channels. 

2 

3RF-DETR loads checkpoint weights into a 3-channel model and only afterwards 

4widens the input layer to ``num_channels``, so a checkpoint whose input layer is 

5already wider (e.g. one written by the ``adapt-input-channels`` task command or by multi-channel 

6training) cannot be loaded through ``pretrain_weights``. Its Lightning training 

7module also never widens the input layer at all. 

8 

9``multichannel_checkpoints()`` wraps ``load_pretrain_weights`` so the model's 

10patch projection is resized to the checkpoint's channel count before loading, 

11and, for training, widened to ``model_config.num_channels`` afterwards using 

12RF-DETR's own filter repetition. 

13""" 

14 

15from __future__ import annotations 

16 

17import contextlib 

18import copy 

19import os 

20from typing import Iterator 

21 

22import torch 

23 

24from .rfdetr_data import channel_stats 

25 

26PROJECTION_KEY = "backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight" 

27 

28 

29MIN_RFDETR_VERSION = (1, 10, 1) 

30 

31 

32def require_multichannel_rfdetr() -> None: 

33 """Fail clearly on rfdetr releases whose data pipeline is PIL/RGB-only.""" 

34 from importlib.metadata import version 

35 

36 installed = version("rfdetr") 

37 parts = tuple(int(p) for p in installed.split(".")[:3] if p.isdigit()) 

38 if parts < MIN_RFDETR_VERSION: 

39 required = ".".join(map(str, MIN_RFDETR_VERSION)) 

40 raise RuntimeError( 

41 f"Training RF-DETR on TIFF inputs requires rfdetr>={required}, " 

42 f"but rfdetr {installed} is installed. Upgrade it, " 

43 f"e.g. `pip install 'rfdetr[metrics,train,loggers]>={required}'`." 

44 ) 

45 

46 

47def _patch_embeddings(nn_model): 

48 return nn_model.backbone[0].encoder.encoder.embeddings.patch_embeddings 

49 

50 

51def checkpoint_channels(path) -> int: 

52 """Input channels of a checkpoint's patch projection (3 if it cannot be read).""" 

53 if path is None or not os.path.isfile(str(path)): 

54 return 3 

55 # RF-DETR checkpoints store their args as Python objects, as RF-DETR's 

56 # own loader assumes when it reads the same file. 

57 checkpoint = torch.load(str(path), map_location="cpu", weights_only=False) 

58 state = checkpoint.get("model") 

59 if state is None: 

60 state = { 

61 key.removeprefix("model.").removeprefix("_orig_mod."): value 

62 for key, value in checkpoint.get("state_dict", {}).items() 

63 } 

64 weight = state.get(PROJECTION_KEY) 

65 return 3 if weight is None else int(weight.shape[1]) 

66 

67 

68def set_input_channels(nn_model, channels: int) -> None: 

69 """Widen or narrow the patch projection, repeating pretrained filters as RF-DETR does.""" 

70 from rfdetr.inference import _adapt_input_conv 

71 

72 embeddings = _patch_embeddings(nn_model) 

73 projection = embeddings.projection 

74 if projection.in_channels == channels: 

75 return 

76 if projection.in_channels != 3: 

77 raise ValueError( 

78 f"Cannot adapt a {projection.in_channels}-channel input layer to {channels} channels" 

79 ) 

80 new_projection = copy.deepcopy(projection) 

81 new_projection.in_channels = channels 

82 new_projection.weight = torch.nn.Parameter( 

83 _adapt_input_conv(channels, projection.weight), 

84 requires_grad=projection.weight.requires_grad, 

85 ) 

86 embeddings.projection = new_projection 

87 embeddings.num_channels = channels 

88 

89 

90def sync_channels(detector) -> int: 

91 """Align an RFDETR wrapper's config and normalization with its input layer. 

92 

93 Wrappers without a DINOv2 patch projection (e.g. test doubles) are left unchanged. 

94 """ 

95 try: 

96 encoder = detector.model.model.backbone[0].encoder 

97 channels = ( 

98 3 * encoder.num_views if hasattr(encoder, "fusion_spec") 

99 else _patch_embeddings(detector.model.model).projection.in_channels 

100 ) 

101 except (AttributeError, IndexError, TypeError): 

102 return getattr(getattr(detector, "model_config", None), "num_channels", 3) 

103 detector.model_config.num_channels = channels 

104 args = getattr(detector.model, "args", None) 

105 if args is not None and hasattr(args, "num_channels"): 

106 args.num_channels = channels 

107 mean, std = channel_stats(channels) 

108 detector.means, detector.stds = list(mean), list(std) 

109 return channels 

110 

111 

112def _channel_aware(load_pretrain_weights, widen_after: bool): 

113 def load(nn_model, model_config, *args, **kwargs): 

114 # Feature fusion processes each view with an RGB projection. Its builder 

115 # installs the fusion modules before loading; never widen that projection. 

116 if hasattr(nn_model.backbone[0].encoder, "fusion_spec"): 

117 return load_pretrain_weights(nn_model, model_config, *args, **kwargs) 

118 stored = checkpoint_channels(model_config.pretrain_weights) 

119 if stored != 3: 

120 set_input_channels(nn_model, stored) 

121 result = load_pretrain_weights(nn_model, model_config, *args, **kwargs) 

122 if widen_after and model_config.num_channels != 3: 

123 set_input_channels(nn_model, model_config.num_channels) 

124 return result 

125 

126 return load 

127 

128 

129@contextlib.contextmanager 

130def multichannel_checkpoints() -> Iterator[None]: 

131 """Let RF-DETR load, and train from, checkpoints with any number of input channels.""" 

132 import rfdetr.inference as inference 

133 

134 # The inference builder widens the input layer itself after loading. 

135 patches = [(inference, False)] 

136 try: 

137 import rfdetr.training.module_model as module_model 

138 

139 patches.append((module_model, True)) 

140 except ImportError: 

141 pass 

142 

143 originals = [(module, module.load_pretrain_weights) for module, _ in patches] 

144 for module, widen_after in patches: 

145 module.load_pretrain_weights = _channel_aware( 

146 module.load_pretrain_weights, widen_after 

147 ) 

148 try: 

149 yield 

150 finally: 

151 for module, original in originals: 

152 module.load_pretrain_weights = original