Coverage for collectra/tasks/object_detection/adaptation_commands.py: 94.23%

1"""CLI commands for adapting detector checkpoints to aligned RGB views.""" 

2 

3from pathlib import Path 

4 

5from collectra.cli import command 

6 

7 

8def _output_path(input_model, output_model, num_views, suffix, overwrite): 

9 input_model = Path(input_model) 

10 if not input_model.is_file(): 

11 raise ValueError(f"Input checkpoint does not exist: {input_model}") 

12 if num_views < 1: 

13 raise ValueError("num_views must be >= 1") 

14 output_model = ( 

15 Path(output_model) 

16 if output_model 

17 else input_model.with_name( 

18 f"{input_model.stem}-{num_views}views-{suffix}{input_model.suffix or '.pth'}" 

19 ) 

20 ) 

21 if input_model.resolve() == output_model.resolve(): 

22 raise ValueError("Output must differ from the input checkpoint") 

23 if output_model.exists() and not overwrite: 

24 raise ValueError("Output already exists; use --overwrite") 

25 return output_model 

26 

27 

28def _check_fusion(num_views, reference_view, hidden_dim): 

29 if not 0 <= reference_view < num_views: 

30 raise ValueError("reference_view must be between 0 and num_views - 1") 

31 if hidden_dim < 1: 

32 raise ValueError("hidden_dim must be >= 1") 

33 

34 

35class RFDETRAdaptationCommands: 

36 @command 

37 def adapt_input_channels( 

38 self, 

39 input_model: Path, 

40 num_views: int, 

41 output: Path | None = None, 

42 variant: str | None = None, 

43 overwrite: bool = False, 

44 validate: bool = True, 

45 ) -> Path: 

46 """Expand RGB input filters across aligned views and save a checkpoint.""" 

47 destination = _output_path( 

48 input_model, output, num_views, f"{3 * num_views}ch", overwrite 

49 ) 

50 from .rfdetr_input_adaptation import ( 

51 expand_rfdetr_input_channels, 

52 validate_saved_model, 

53 ) 

54 

55 expand_rfdetr_input_channels(input_model, destination, num_views, variant) 

56 if validate: 

57 validate_saved_model(destination, 3 * num_views) 

58 print(f"Saved: {destination}") 

59 return destination 

60 

61 @command 

62 def adapt_feature_fusion( 

63 self, 

64 input_model: Path, 

65 num_views: int, 

66 output: Path | None = None, 

67 variant: str | None = None, 

68 reference_view: int = 0, 

69 hidden_dim: int = 64, 

70 overwrite: bool = False, 

71 validate: bool = True, 

72 ) -> Path: 

73 """Add shared RGB encoding and attention across aligned views.""" 

74 destination = _output_path(input_model, output, num_views, "fusion", overwrite) 

75 _check_fusion(num_views, reference_view, hidden_dim) 

76 from .rfdetr_feature_fusion import ( 

77 expand_rfdetr_feature_fusion, 

78 validate_saved_model, 

79 ) 

80 

81 expand_rfdetr_feature_fusion( 

82 input_model, destination, num_views, variant, reference_view, hidden_dim 

83 ) 

84 if validate: 

85 validate_saved_model(destination, num_views) 

86 print(f"Saved: {destination}") 

87 return destination 

88 

89 

90class YOLOAdaptationCommands: 

91 @command 

92 def adapt_feature_fusion( 

93 self, 

94 input_model: Path, 

95 num_views: int, 

96 output: Path | None = None, 

97 reference_view: int = 0, 

98 hidden_dim: int = 64, 

99 overwrite: bool = False, 

100 validate: bool = True, 

101 ) -> Path: 

102 """Add a shared RGB backbone and attention across aligned views.""" 

103 destination = _output_path(input_model, output, num_views, "fusion", overwrite) 

104 if Path(input_model).suffix.lower() != ".pt": 

105 raise ValueError("Input must be an RGB YOLO .pt checkpoint") 

106 _check_fusion(num_views, reference_view, hidden_dim) 

107 from .yolo_feature_fusion import ( 

108 expand_yolo_feature_fusion, 

109 validate_saved_model, 

110 ) 

111 

112 expand_yolo_feature_fusion( 

113 input_model, destination, num_views, reference_view, hidden_dim 

114 ) 

115 if validate: 

116 validate_saved_model(destination, num_views) 

117 print(f"Saved: {destination}") 

118 return destination