Coverage for collectra/tasks/object_detection/rfdetr_data.py: 93.06%

1"""Multi-channel TIFF support for the RF-DETR training data pipeline. 

2 

3RF-DETR's ``CocoDetection`` decodes every image with ``PIL.Image.convert("RGB")``, 

4which drops channels beyond three and cannot open most multispectral TIFFs. This 

5module swaps the built datasets for a subclass that decodes TIFFs with 

6``read_tiff_channels`` into CHW tensors, and replaces the ImageNet ``Normalize`` 

7with one that matches the image's channel count. 

8 

9Kept separate from ``rfdetr.py`` so that importing the task does not import 

10``rfdetr``, and so DataLoader worker processes can unpickle these classes. 

11""" 

12 

13from __future__ import annotations 

14 

15import contextlib 

16from pathlib import Path 

17from typing import Any, Iterator 

18 

19import numpy as np 

20import torch 

21from rfdetr.datasets.coco import CocoDetection 

22from rfdetr.datasets.transforms import AlbumentationsWrapper, Normalize 

23from torchvision import tv_tensors 

24 

25from collectra.types.images import read_tiff_channels 

26 

27TIFF_SUFFIXES = {".tif", ".tiff"} 

28IMAGENET_MEAN = (0.485, 0.456, 0.406) 

29IMAGENET_STD = (0.229, 0.224, 0.225) 

30 

31 

32def channel_stats(channels: int) -> tuple[tuple[float, ...], tuple[float, ...]]: 

33 """ImageNet statistics cycled across channels, matching ``RFDETR.predict``.""" 

34 mean = tuple(IMAGENET_MEAN[i % 3] for i in range(channels)) 

35 std = tuple(IMAGENET_STD[i % 3] for i in range(channels)) 

36 return mean, std 

37 

38 

39def scale_tiff_pixels(pixels: np.ndarray, keep_uint8: bool = False) -> np.ndarray: 

40 """Scale integer pixels to [0, 1] float32 as ``ToDtype(scale=True)`` would. 

41 

42 torch lacks CPU kernels (e.g. flip) for uint16/uint32, so wider integer 

43 types are scaled here rather than in the transform pipeline. Float 

44 pixels are assumed to be in [0, 1] already. 

45 """ 

46 if pixels.dtype == np.uint8 and keep_uint8: 

47 return pixels 

48 if np.issubdtype(pixels.dtype, np.integer): 

49 return pixels.astype(np.float32) / np.iinfo(pixels.dtype).max 

50 return pixels.astype(np.float32) 

51 

52 

53class ChannelNormalize(Normalize): 

54 """RF-DETR ``Normalize`` whose mean/std follow the image's channel count.""" 

55 

56 def __call__(self, image, target=None): 

57 channels = image.shape[-3] 

58 if channels != len(IMAGENET_MEAN): 

59 mean, std = channel_stats(channels) 

60 normalize = Normalize(mean, std) 

61 return normalize(image, target) 

62 return super().__call__(image, target) 

63 

64 

65class _ImageSize: 

66 """Stand-in exposing only ``.size``, the one attribute ``ConvertCoco`` reads.""" 

67 

68 def __init__(self, width: int, height: int): 

69 self.size = (width, height) 

70 

71 

72class MultiChannelCocoDetection(CocoDetection): 

73 """``CocoDetection`` that keeps every channel of TIFF images.""" 

74 

75 def _image_path(self, image_id: int) -> Path: 

76 return Path(self.root) / self.coco.loadImgs(image_id)[0]["file_name"] 

77 

78 def __getitem__(self, idx: int) -> tuple[Any, Any]: 

79 image_id = self.ids[idx] 

80 path = self._image_path(image_id) 

81 if path.suffix.lower() not in TIFF_SUFFIXES: 

82 return super().__getitem__(idx) 

83 

84 pixels = scale_tiff_pixels(read_tiff_channels(path), keep_uint8=True) 

85 img = tv_tensors.Image(torch.from_numpy(pixels.transpose(2, 0, 1).copy())) 

86 height, width = pixels.shape[:2] 

87 target = {"image_id": image_id, "annotations": self._load_target(image_id)} 

88 _, target = self.prepare(_ImageSize(width, height), target) 

89 if self._transforms is not None: 

90 img, target = self._transforms(img, target) 

91 return img, target 

92 

93 

94def _adapt_dataset(dataset): 

95 if type(dataset) is not CocoDetection: 

96 return dataset 

97 dataset.__class__ = MultiChannelCocoDetection 

98 transforms = getattr(dataset._transforms, "transforms", []) 

99 has_tiffs = any( 

100 Path(image["file_name"]).suffix.lower() in TIFF_SUFFIXES 

101 for image in dataset.coco.imgs.values() 

102 ) 

103 if has_tiffs and any(isinstance(t, AlbumentationsWrapper) for t in transforms): 

104 # The Albumentations wrapper round-trips images through PIL, which 

105 # cannot hold arbitrary channel counts. 

106 raise ValueError( 

107 "TIFF inputs are only supported with the default (torchvision) " 

108 "augmentation pipeline; use augmentation='default'." 

109 ) 

110 for index, transform in enumerate(transforms): 

111 if type(transform) is Normalize: 

112 transforms[index] = ChannelNormalize() 

113 return dataset 

114 

115 

116@contextlib.contextmanager 

117def multichannel_datasets() -> Iterator[None]: 

118 """Make RF-DETR training build datasets that keep all TIFF channels.""" 

119 import rfdetr.training.module_data as module_data 

120 

121 original = module_data.build_dataset 

122 

123 def build_dataset(image_set: str, args: Any, resolution: int): 

124 return _adapt_dataset(original(image_set, args, resolution)) 

125 

126 module_data.build_dataset = build_dataset 

127 try: 

128 yield 

129 finally: 

130 module_data.build_dataset = original