Coverage for collectra/tasks/image_classifier/yolo_multichannel.py: 94.62%

1"""Tensor-based Ultralytics classification for TIFF and non-RGB images. 

2 

3Imported only when needed so CLI task discovery does not load ML backends. 

4""" 

5 

6from copy import copy, deepcopy 

7from pathlib import Path 

8 

9import numpy as np 

10import torch 

11from PIL import Image as PillowImage 

12from torchvision import transforms as T 

13from torchvision.datasets import ImageFolder 

14from ultralytics.models.yolo.classify import ( 

15 ClassificationTrainer, 

16 ClassificationValidator, 

17) 

18from ultralytics.nn.tasks import yaml_model_load 

19from ultralytics.utils import LOGGER 

20 

21from collectra.types.images import read_tiff_channels 

22from .torchvision import _target_tiff_pixels, _tiff_tensor 

23from .yolo import _classification_channels 

24 

25 

26def load_channels(path): 

27 """Load CHW float32; unsigned integers scale by their dtype's full range.""" 

28 if Path(path).suffix.lower() in {".tif", ".tiff"}: 

29 pixels = read_tiff_channels(path) 

30 else: 

31 with PillowImage.open(path) as image: 

32 if image.mode == "P": 

33 image = image.convert("RGBA" if "transparency" in image.info else "RGB") 

34 pixels = np.array(image) 

35 if pixels.ndim == 2: 

36 pixels = pixels[..., None] 

37 return _tiff_tensor(pixels) 

38 

39 

40def channel_transforms(size, args=None): 

41 """Use shared spatial transforms on all bands, with identity normalization.""" 

42 if args is None: 

43 return T.Compose([T.Resize(size), T.CenterCrop(size)]) 

44 return T.Compose( 

45 [ 

46 T.RandomResizedCrop(size, scale=(1.0 - args.scale, 1.0)), 

47 T.RandomHorizontalFlip(args.fliplr), 

48 T.RandomVerticalFlip(args.flipud), 

49 T.RandomErasing(p=args.erasing), 

50 ] 

51 ) 

52 

53 

54class MultichannelClassificationDataset: 

55 """Ultralytics batch contract backed by ImageFolder and lossless TIFF reading.""" 

56 

57 def __init__(self, root, args, names, augment=False): 

58 self.base = ImageFolder(root, allow_empty=True) 

59 mapping = {name: index for index, name in names.items()} 

60 unknown = set(self.base.classes) - set(mapping) 

61 if unknown: 

62 raise ValueError( 

63 f"Unknown classification directories in {root}: {sorted(unknown)}" 

64 ) 

65 self.samples = [ 

66 (path, mapping[self.base.classes[index]]) 

67 for path, index in self.base.samples 

68 ] 

69 if augment and args.fraction < 1.0: 

70 self.samples = self.samples[: round(len(self.samples) * args.fraction)] 

71 self.torch_transforms = channel_transforms( 

72 args.imgsz, args if augment else None 

73 ) 

74 

75 def __len__(self): 

76 return len(self.samples) 

77 

78 def __getitem__(self, index): 

79 path, label = self.samples[index] 

80 return {"img": self.torch_transforms(load_channels(path)), "cls": label} 

81 

82 

83class MultichannelClassificationValidator(ClassificationValidator): 

84 def __call__(self, trainer=None, model=None): 

85 # Standalone/final validation rechecks classification metadata (upstream 

86 # defaults to RGB). Rebuild the loader so channels are corrected before 

87 # AutoBackend's warmup, including on GPU. 

88 if trainer is None: 

89 self.dataloader = None 

90 return super().__call__(trainer=trainer, model=model) 

91 

92 def build_dataset(self, img_path): 

93 self.data["channels"], _ = _classification_channels( 

94 self.data.get("train"), self.data.get("val"), self.data.get("test") 

95 ) 

96 return MultichannelClassificationDataset( 

97 img_path, self.args, self.data["names"] 

98 ) 

99 

100 def plot_val_samples(self, batch, ni): 

101 """RGB sample plots cannot represent spectral bands; keep metric plots.""" 

102 

103 def plot_predictions(self, batch, preds, ni): 

104 """Do not pass spectral tensors to RGB visualization routines.""" 

105 

106 

107class MultichannelClassificationTrainer(ClassificationTrainer): 

108 def get_dataset(self): 

109 data = super().get_dataset() 

110 data["channels"], _ = _classification_channels( 

111 data.get("train"), data.get("val"), data.get("test") 

112 ) 

113 LOGGER.info( 

114 f"Collectra: preserving {data['channels']} image channels; using spatial " 

115 "augmentation and erasing, identity normalization, and no RGB colour/auto augmentation." 

116 ) 

117 return data 

118 

119 def get_model(self, cfg=None, weights=None, verbose=True): 

120 # ClassificationModel otherwise lets an old YAML's channels override ch. 

121 # Dataset metadata is authoritative, including when transferring weights. 

122 cfg = deepcopy(cfg) if isinstance(cfg, dict) else yaml_model_load(cfg) 

123 cfg["channels"] = self.data["channels"] 

124 model = super().get_model(cfg=cfg, weights=weights, verbose=verbose) 

125 model.collectra_tensor_classification = True 

126 return model 

127 

128 def build_dataset(self, img_path, mode="train", batch=None): 

129 return MultichannelClassificationDataset( 

130 img_path, self.args, self.data["names"], augment=mode == "train" 

131 ) 

132 

133 def get_validator(self): 

134 self.loss_names = ["loss"] 

135 return MultichannelClassificationValidator( 

136 self.test_loader, 

137 self.save_dir, 

138 args=copy(self.args), 

139 _callbacks=self.callbacks, 

140 ) 

141 

142 def plot_training_samples(self, batch, ni): 

143 """RGB sample plots cannot represent spectral bands; keep metric plots.""" 

144 

145 

146@torch.inference_mode() 

147def classify_multichannel(yolo, image): 

148 """Classify directly from tensors, bypassing Ultralytics' RGB predictor.""" 

149 if image.get_path().suffix.lower() in {".tif", ".tiff"}: 

150 pixels = _tiff_tensor(_target_tiff_pixels(image)) 

151 else: 

152 # Pillow supports the ordinary grayscale/RGBA crop and orientation path. 

153 with image.pil() as source: 

154 if source.mode == "P": 

155 source = source.convert( 

156 "RGBA" if "transparency" in source.info else "RGB" 

157 ) 

158 values = np.array(source) 

159 pixels = _tiff_tensor(values[..., None] if values.ndim == 2 else values) 

160 model = yolo.model 

161 first_conv = next(m for m in model.modules() if isinstance(m, torch.nn.Conv2d)) 

162 if pixels.shape[0] != first_conv.in_channels: 

163 raise ValueError( 

164 f"{image.get_path()}: {pixels.shape[0]} image channels, " 

165 f"but the classification model expects {first_conv.in_channels}" 

166 ) 

167 transform = getattr(model, "transforms", None) 

168 # Models trained via Collectra store tensor transforms in the checkpoint. 

169 # A separately prepared model may still carry upstream PIL/RGB transforms. 

170 if transform is None or any( 

171 isinstance(t, T.ToTensor) for t in getattr(transform, "transforms", []) 

172 ): 

173 transform = channel_transforms(yolo.overrides.get("imgsz", 224)) 

174 parameter = next(model.parameters()) 

175 batch = ( 

176 transform(pixels) 

177 .unsqueeze(0) 

178 .to(device=parameter.device, dtype=parameter.dtype) 

179 ) 

180 model.eval() 

181 predictions = model(batch) 

182 predictions = ( 

183 predictions[0] if isinstance(predictions, (tuple, list)) else predictions 

184 ) 

185 return model.names[predictions.argmax(dim=1).item()]