Coverage for collectra/tasks/image_classifier/torchvision.py: 93.48%

1"""Torchvision image classification with portable, metadata-bearing checkpoints."""

3from __future__ import annotations

5from dataclasses import dataclass

6from pathlib import Path

7from typing import Annotated

9from cappa import Arg

11from collectra.cli import command

12from collectra.logger import get_logger

13from collectra.types.images import (

14 Image,

15 ImageCrop,

16 image_channel_count,

17 read_tiff_channels,

18)

19from collectra.types.links import Link

20from collectra.utils import threading_locked

21from ..machine_learning.training import run_training_command

22from .training import (

23 ClassificationTrainResult,

24 TorchClassifierTask,

25 prepare_classification_inputs,

26)

28__all__ = ["ImageClassifierTorchvision"]

29logger = get_logger(__name__)

32def _tiff_tensor(pixels):

33 """Convert HWC TIFF samples to float32 CHW, preserving every channel."""

34 import numpy as np

35 import torch

37 values = pixels.astype(np.float32)

38 if np.issubdtype(pixels.dtype, np.unsignedinteger):

39 values /= np.iinfo(pixels.dtype).max

40 return torch.from_numpy(np.ascontiguousarray(values.transpose(2, 0, 1)))

43def _load_classifier_image(path):

44 """ImageFolder loader that bypasses Pillow for TIFF images."""

45 from torchvision.datasets.folder import pil_loader

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

48 return _tiff_tensor(read_tiff_channels(path))

49 return pil_loader(path)

52def _target_tiff_pixels(image, *, tiff_maxworkers: int | None = None):

53 import numpy as np

55 target = image.resolve() if type(image) is Link else image

56 pixels = read_tiff_channels(target.get_path(), maxworkers=tiff_maxworkers)

57 if isinstance(target, ImageCrop):

58 left, top, right, bottom = target.coordinates()

59 if right <= left or bottom <= top:

60 raise ValueError(f"ImageCrop {target.id!r} produces an empty pixel crop")

61 pixels = pixels[top:bottom, left:right, :]

62 return np.rot90(pixels, k=target.orientation.to_degree() // 90, axes=(0, 1))

65@dataclass

66class _NormalizeClassifierPixels:

67 mean: list[float]

68 std: list[float]

70 def __call__(self, pixels):

71 import torch

72 from torchvision.transforms import functional as F

74 if not isinstance(pixels, torch.Tensor):

75 pixels = F.to_tensor(pixels)

76 channels = pixels.shape[0]

77 if len(self.mean) == len(self.std) == 3 and channels != 3:

78 # RGB statistics have no defined meaning for spectral bands.

79 return pixels

80 if len(self.mean) != len(self.std) or len(self.mean) not in (1, channels):

81 raise ValueError(

82 f"Normalization statistics do not match {channels} image channels"

83 )

84 return F.normalize(pixels, self.mean, self.std)

87def _replace_head(model, num_classes: int):

88 """Replace the output projection for torchvision's classification families."""

89 from torch import nn

91 for name in ("fc", "classifier", "heads", "head"):

92 head = getattr(model, name, None)

93 if head is None:

94 continue

95 if isinstance(head, nn.Linear):

96 replacement = nn.Linear(

97 head.in_features, num_classes, bias=head.bias is not None

98 )

99 setattr(model, name, replacement)

100 return replacement

101 if isinstance(head, nn.Sequential):

102 projections = [

103 (path, module)

104 for path, module in head.named_modules()

105 if isinstance(module, (nn.Linear, nn.Conv2d))

106 ]

107 if not projections:

108 continue

109 path, layer = projections[-1]

110 if isinstance(layer, nn.Linear):

111 replacement = nn.Linear(

112 layer.in_features, num_classes, bias=layer.bias is not None

113 )

114 else: # SqueezeNet's classifier is a 1x1 convolution.

115 replacement = nn.Conv2d(

116 layer.in_channels,

117 num_classes,

118 layer.kernel_size,

119 stride=layer.stride,

120 padding=layer.padding,

121 dilation=layer.dilation,

122 groups=layer.groups,

123 bias=layer.bias is not None,

124 )

125 model.num_classes = num_classes

126 parent_path, _, child = path.rpartition(".")

127 parent = head.get_submodule(parent_path) if parent_path else head

128 setattr(parent, child, replacement)

129 return replacement

130 raise ValueError(f"Unsupported classification head in {type(model).__name__}")

131

132

133def _disable_auxiliary_heads(model) -> None:

134 # GoogLeNet and Inception should return a single logits tensor in training too.

135 if hasattr(model, "aux_logits"):

136 model.aux_logits = False

137 for name in ("AuxLogits", "aux1", "aux2"):

138 if hasattr(model, name):

139 setattr(model, name, None)

140

141

142def _ensure_supported_multichannel_architecture(

143 architecture: str, input_channels: int

144) -> None:

145 if input_channels == 3:

146 return

147 #supported = {"resnet18", "resnet34", "resnet50"}

148 supported = {"resnet18", "resnet34", "resnet50", "resnet101", "resnet152",

149 "convnext_base", "convnext_tiny", "convnext_small", "convnext_large"}

150 if architecture not in supported:

151 raise ValueError(

152 f"{input_channels}-channel input adaptation is currently supported only for "

153 "resnet18, resnet34, resnet50, resnet101, resnet152, convnext_base, convnext_tiny, convnext_small, convnext_large"

154 )

155

156

157class ImageClassifierTorchvision(TorchClassifierTask):

158 """Train torchvision classifiers and return predicted class Links."""

159

160 prepare_training_inputs = staticmethod(prepare_classification_inputs)

161

162 def __init__(self, name: str, model: str | Path = "resnet18", **kwargs):

163 super().__init__(name, model=model, **kwargs)

164 self.original_model_path: Path | None = None

165 self._categories: list[str] = []

166 self._architecture = ""

167 self._model_kwargs: dict = {}

168 self._preprocessing: dict = {}

169 self._device = "cpu"

170 self._input_channels = 3

171

172 def _transforms(

173 self, *, training: bool = False, fliplr: float = 0.0, flipud: float = 0.0

174 ):

175 from torchvision import transforms as T

176

177 config = self._preprocessing

178 interpolation = T.InterpolationMode(config["interpolation"])

179 steps = [

180 T.Resize(

181 config["resize_size"], interpolation=interpolation, antialias=True

182 ),

183 T.CenterCrop(config["crop_size"]),

184 ]

185 if training and fliplr:

186 steps.append(T.RandomHorizontalFlip(fliplr))

187 if training and flipud:

188 steps.append(T.RandomVerticalFlip(flipud))

189 steps.append(_NormalizeClassifierPixels(config["mean"], config["std"]))

190 return T.Compose(steps)

191

192 def _export_image(

193 self, target, destination, min_size, tiff_maxworkers: int | None = None

194 ):

195 if target.get_path().suffix.lower() not in {".tif", ".tiff"}:

196 return super()._export_image(

197 target,

198 destination,

199 min_size,

200 tiff_maxworkers=tiff_maxworkers,

201 )

202 import numpy as np

203 import tifffile

204

205 pixels = _target_tiff_pixels(target, tiff_maxworkers=tiff_maxworkers)

206 if min(pixels.shape[:2]) < max(1, min_size):

207 logger.warning(

208 "Skipping classifier image %s below minimum size %d",

209 target.id,

210 min_size,

211 )

212 return False

213 tifffile.imwrite(

214 destination.with_suffix(".tif"),

215 np.ascontiguousarray(pixels.transpose(2, 0, 1)),

216 photometric="minisblack",

217 metadata={"axes": "CYX"},

218 )

219 return True

220

221 def _make_datasets(self, train_dir, val_dir, augmentation):

222 from torchvision.datasets import ImageFolder

223

224 train = ImageFolder(train_dir, loader=_load_classifier_image)

225 val = ImageFolder(val_dir, loader=_load_classifier_image, allow_empty=True)

226 detected = {

227 path: (

228 image_channel_count(path)

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

230 else 3

231 )

232 for path, _ in [*train.samples, *val.samples]

233 }

234 if len(set(detected.values())) != 1:

235 details = ", ".join(

236 f"{path}: {count} channels" for path, count in detected.items()

237 )

238 raise ValueError(

239 f"Inconsistent input channel counts in classifier dataset: {details}"

240 )

241 channels = next(iter(detected.values()))

242 self._preprocessing["channels"] = channels

243 self._input_channels = channels

244 if (

245 channels != 3

246 and len(self._preprocessing["mean"]) == len(self._preprocessing["std"]) == 3

247 ):

248 if channels % 3 == 0:

249 num_views = channels // 3

250 logger.info(

251 "Detected %d RGB views (%d channels); repeating pretrained RGB normalization",

252 num_views,

253 channels,

254 )

255 self._preprocessing["mean"] = self._preprocessing["mean"] * num_views

256 self._preprocessing["std"] = self._preprocessing["std"] * num_views

257 else:

258 logger.info(

259 "Using identity normalization for %d channels; channel count is not divisible by 3",

260 channels,

261 )

262 self._preprocessing["mean"] = [0.0] * channels

263 self._preprocessing["std"] = [1.0] * channels

264 train.transform = self._transforms(training=True, **augmentation)

265 val.transform = self._transforms()

266 return train, val

267

268 def _load(self, *, pretrained: bool = True) -> None:

269 import torch

270 from torchvision import models

271

272 if isinstance(self.model, torch.nn.Module):

273 return

274 if not isinstance(self.model, (str, Path)):

275 raise ValueError(

276 "Model must be a torchvision model name or a checkpoint path"

277 )

278 path = Path(self.model).expanduser()

279 if path.is_file():

280 checkpoint = torch.load(path, map_location="cpu", weights_only=True)

281 required = {

282 "architecture",

283 "state_dict",

284 "classes",

285 "preprocessing",

286 "model_kwargs",

287 }

288 if not isinstance(checkpoint, dict) or not required.issubset(checkpoint):

289 raise ValueError(

290 "Expected a Collectra torchvision checkpoint containing architecture, "

291 "state_dict, classes, preprocessing, and model_kwargs. "

292 "A bare state_dict does not identify its architecture or class labels."

293 )

294 self._architecture = checkpoint["architecture"]

295 self._model_kwargs = checkpoint["model_kwargs"]

296 self._categories = checkpoint["classes"]

297 self._preprocessing = checkpoint["preprocessing"]

298 model = models.get_model(

299 self._architecture, weights=None, **self._model_kwargs

300 )

301 self._input_channels = checkpoint.get("input_channels", 3)

302

303 if self._input_channels != 3:

304 _ensure_supported_multichannel_architecture(

305 self._architecture, self._input_channels

306 )

307 _adapt_rgb_input_channels(

308 model,

309 self._input_channels,

310 )

311

312 _disable_auxiliary_heads(model)

313 _replace_head(model, len(self._categories))

314 model.load_state_dict(checkpoint["state_dict"])

315 self.original_model_path = path.resolve()

316 else:

317 name = str(self.model)

318 available = models.list_models(module=models)

319 if name not in available:

320 raise ValueError(

321 f"Unknown torchvision classification model or missing checkpoint: {name!r}. "

322 f"Available models: {', '.join(available)}"

323 )

324 weights = models.get_model_weights(name).DEFAULT

325 preset = weights.transforms()

326 self._preprocessing = {

327 "crop_size": preset.crop_size,

328 "resize_size": preset.resize_size,

329 "mean": preset.mean,

330 "std": preset.std,

331 "interpolation": preset.interpolation.value,

332 }

333 builder_kwargs = {}

334 if name.startswith("vit_"):

335 builder_kwargs["image_size"] = preset.crop_size[0]

336 if name in {"inception_v3", "googlenet"}:

337 builder_kwargs["init_weights"] = False

338 model = models.get_model(

339 name, weights=weights if pretrained else None, **builder_kwargs

340 )

341 _disable_auxiliary_heads(model)

342 self._architecture = name

343 self._model_kwargs = dict(builder_kwargs)

344 if hasattr(model, "transform_input"):

345 self._model_kwargs["transform_input"] = model.transform_input

346 if hasattr(model, "aux_logits"):

347 self._model_kwargs["aux_logits"] = False

348 if hasattr(model, "image_size"):

349 self._model_kwargs["image_size"] = model.image_size

350 self._categories = list(weights.meta["categories"])

351 self.original_model_path = None

352 self.model = model.to(self._device)

353 self.model.eval()

354

355 @threading_locked()

356 def run(self, *args: Image) -> Link:

357 """Classify one image or crop and link the predicted class to it."""

358 import torch

359

360 if len(args) != 1 or not isinstance(args[0], Image):

361 raise ValueError("This task requires a single Image or ImageCrop input")

362 self._device = self._select_device(getattr(self, "device", ""))

363 self._load()

364 self.model.to(self._device).eval()

365 image = args[0]

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

367 pixels = _tiff_tensor(_target_tiff_pixels(image))

368 batch = self._transforms()(pixels).unsqueeze(0).to(self._device)

369 else:

370 with image.pil() as pixels:

371 batch = (

372 self._transforms()(pixels.convert("RGB"))

373 .unsqueeze(0)

374 .to(self._device)

375 )

376 with torch.inference_mode():

377 class_index = self.model(batch).argmax(dim=1).item()

378 class_name = self._categories[class_index]

379 outputs = getattr(self, "output", [])

380 outputs = [outputs] if isinstance(outputs, str) else outputs

381 if outputs and class_name not in outputs:

382 raise ValueError(

383 f"Predicted class {class_name!r} is not a configured output: {outputs}"

384 )

385 return Link(name=class_name, target=image)

386

387 @command

388 def train(

389 self,

390 inputs: Annotated[

391 list[str],

392 Arg(

393 help="Pipeline data files or directories containing classification labels."

394 ),

395 ],

396 model: Annotated[

397 str,

398 Arg(

399 help="Torchvision classification model name (e.g. resnet18, efficientnet_b0, vit_b_16) or a Collectra torchvision checkpoint. Omit to keep the configured model."

400 ),

401 ] = "",

402 output: Annotated[

403 Path | None,

404 Arg(

405 help="Directory for extracted train/val images, metrics, and checkpoints."

406 ),

407 ] = None,

408 use_existing: Annotated[

409 bool,

410 Arg(

411 help="Reuse an existing output train/ and val/ dataset if present; skip image export."

412 ),

413 ] = False,

414 keep_log: Annotated[

415 bool,

416 Arg(

417 help="Keep extracted images and training logs after saving the model to the pipeline."

418 ),

419 ] = True,

420 validation: Annotated[

421 str,

422 Arg(

423 help="Partition value identifying validation images; omit to train without validation."

424 ),

425 ] = "",

426 exclude: Annotated[

427 str, Arg(help="Partition value identifying images to exclude.")

428 ] = "",

429 epochs: Annotated[int, Arg(help="Number of training epochs.")] = 10,

430 batch: Annotated[int, Arg(help="Number of images per training batch.")] = 16,

431 learning_rate: Annotated[

432 float, Arg(help="Initial learning rate for AdamW.")

433 ] = 0.0001,

434 weight_decay: Annotated[

435 float, Arg(help="Weight decay used for regularization.")

436 ] = 0.0001,

437 pretrained: Annotated[

438 bool,

439 Arg(

440 help="Use DEFAULT pretrained weights for a named model; checkpoints always load their saved weights."

441 ),

442 ] = True,

443 freeze_backbone: Annotated[

444 bool,

445 Arg(

446 help="Train only the final classification layer; freeze all other parameters."

447 ),

448 ] = False,

449 early_stop: Annotated[

450 int,

451 Arg(

452 help="Stop after this many epochs without validation-loss improvement; 0 disables."

453 ),

454 ] = 10,

455 min_size: Annotated[

456 int,

457 Arg(

458 help="Skip images or crops with either dimension below this pixel threshold; 0 disables."

459 ),

460 ] = 0,

461 fliplr: Annotated[

462 float,

463 Arg(

464 help="Probability of horizontal flipping during training, from 0 to 1."

465 ),

466 ] = 0.5,

467 workers: Annotated[

468 int,

469 Arg(

470 help="Number of data-loading worker processes (default 8, capped by available CPUs)."

471 ),

472 ] = 8,

473 wandb: Annotated[

474 bool,

475 Arg(help="Log training metrics to Weights & Biases."),

476 ] = False,

477 device: Annotated[

478 str,

479 Arg(

480 help="Torch device such as cpu, cuda, cuda:0, or mps; omit for automatic selection."

481 ),

482 ] = "",

483 seed: Annotated[

484 int, Arg(help="Random seed for model initialization and data loading.")

485 ] = 0,

486 flipud: Annotated[

487 float,

488 Arg(help="Probability of vertical flipping during training, from 0 to 1."),

489 ] = 0.0,

490 ):

491 """Fine-tune a torchvision image classifier using its weights' prescribed image preprocessing."""

492 result = run_training_command(

493 self,

494 inputs,

495 self._train,

496 keep_log=keep_log,

497 output=output,

498 use_existing=use_existing,

499 validation=validation,

500 exclude=exclude,

501 prepare_inputs=self.prepare_training_inputs,

502 model=model,

503 epochs=epochs,

504 batch=batch,

505 learning_rate=learning_rate,

506 weight_decay=weight_decay,

507 pretrained=pretrained,

508 freeze_backbone=freeze_backbone,

509 early_stop=early_stop,

510 min_size=min_size,

511 fliplr=fliplr,

512 flipud=flipud,

513 workers=workers,

514 wandb=wandb,

515 device=device,

516 seed=seed,

517 )

518 if self.pipeline is not None:

519 saved_model = self.pipeline.data[self.name]["model"]

520 self.original_model_path = (self.pipeline.path / saved_model).resolve()

521 return result

522

523 def _checkpoint(self, epoch: int, metrics: dict) -> dict:

524 return {

525 "format_version": 1,

526 "architecture": self._architecture,

527 "model_kwargs": self._model_kwargs,

528 "input_channels": self._input_channels,

529 "classes": self._categories,

530 "preprocessing": self._preprocessing,

531 "state_dict": {

532 name: value.detach().cpu()

533 for name, value in self.model.state_dict().items()

534 },

535 "epoch": epoch,

536 "metrics": metrics,

537 }

538

539 def _prepare_training_model(self, classes, pretrained, freeze):

540 import torch

541

542 self._load(pretrained=pretrained)

543

544 if self._input_channels != 3:

545 _ensure_supported_multichannel_architecture(

546 self._architecture,

547 self._input_channels,

548 )

549 _adapt_rgb_input_channels(

550 self.model,

551 self._input_channels,

552 )

553

554 # Retain a fine-tuned head only when its class order still matches exactly.

555 if classes != self._categories:

556 _replace_head(self.model, len(classes))

557 self._categories = classes

558 self.model.to(self._device)

559 self.model.requires_grad_(not freeze)

560 if freeze:

561 for name in ("fc", "classifier", "heads", "head"):

562 head = getattr(self.model, name, None)

563 if head is not None:

564 layers = [

565 m

566 for m in head.modules()

567 if isinstance(m, (torch.nn.Linear, torch.nn.Conv2d))

568 ]

569 if layers:

570 layers[-1].requires_grad_(True)

571 break

572

573 def _forward(self, pixels):

574 return self.model(pixels.to(self._device))

575

576

577def _adapt_rgb_input_channels(model, input_channels: int):

578 import torch

579

580 # original RT coding

581 # if not hasattr(model, "conv1"):

582 # raise ValueError(

583 # f"{type(model).__name__} does not expose a ResNet-style conv1"

584 # )

585 # old_conv = model.conv1

586 # if not isinstance(old_conv, torch.nn.Conv2d):

587 # raise ValueError(

588 # f"Expected model.conv1 to be Conv2d, got {type(old_conv).__name__}"

589 # )

590 # ---

591 # KT addition to allow for convnext models

592 # Locate the model's first convolution.

593 if hasattr(model, "conv1"):

594 # ResNet-style models

595 old_conv = model.conv1

596 conv_location = "conv1"

597

598 elif (

599 hasattr(model, "features")

600 and len(model.features) > 0

601 and isinstance(model.features[0][0], torch.nn.Conv2d)

602 ):

603 # TorchVision ConvNeXt-style models

604 old_conv = model.features[0][0]

605 conv_location = "convnext"

606

607 else:

608 raise ValueError(

609 f"Could not find first Conv2d for model type "

610 f"{type(model).__name__}"

611 )

612

613 if not isinstance(old_conv, torch.nn.Conv2d):

614 raise ValueError(

615 f"Expected first layer to be Conv2d, "

616 f"got {type(old_conv).__name__}"

617 )

618 # end KT addition

619

620 if old_conv.in_channels == input_channels:

621 return

622

623 if old_conv.in_channels != 3:

624 raise ValueError(

625 f"Expected source model to have 3 input channels, "

626 f"got {old_conv.in_channels}"

627 )

628

629 if old_conv.groups != 1:

630 raise ValueError(

631 f"Cannot adapt grouped conv1 with groups={old_conv.groups}; expected 1"

632 # KT add

633 f"groups={old_conv.groups}; expected 1"

634 )

635

636 if input_channels % 3 != 0:

637 raise ValueError(

638 f"Cannot adapt RGB pretrained weights to "

639 f"{input_channels} channels: channel count must be "

640 f"divisible by 3."

641 )

642

643 num_views = input_channels // 3

644 # KT additional

645 logger.info(

646 "Adapting first convolution: %d -> %d input channels (%d RGB stacks)",

647 old_conv.in_channels,

648 input_channels,

649 num_views,

650 )

651 # end KT additional

652

653 new_conv = torch.nn.Conv2d(

654 in_channels=input_channels,

655 out_channels=old_conv.out_channels,

656 kernel_size=old_conv.kernel_size,

657 stride=old_conv.stride,

658 padding=old_conv.padding,

659 dilation=old_conv.dilation,

660 groups=old_conv.groups,

661 bias=old_conv.bias is not None,

662 padding_mode=old_conv.padding_mode,

663 ).to(

664 device=old_conv.weight.device,

665 dtype=old_conv.weight.dtype,

666 )

667

668 with torch.no_grad():

669 new_conv.weight.copy_(

670 old_conv.weight.repeat(1, num_views, 1, 1)

671 / num_views

672 )

673

674 if old_conv.bias is not None:

675 new_conv.bias.copy_(old_conv.bias)

676

677 # original RT coding

678 # model.conv1 = new_conv

679 # ---

680 # KT addition

681 # Put the adapted convolution back into the model.

682 if conv_location == "conv1":

683 model.conv1 = new_conv

684 else:

685 model.features[0][0] = new_conv

686 # end KT addition