Coverage for collectra/tasks/image_classifier/huggingface.py: 94.67%
169 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Transformers image classifiers with portable Collectra checkpoints."""
3from __future__ import annotations
5import copy
6import json
7import tempfile
8from dataclasses import dataclass
9from pathlib import Path
10from typing import Annotated
12from cappa import Arg
14from collectra.cli import command
15from collectra.types.images import Image
16from collectra.types.links import Link
17from collectra.utils import threading_locked
18from ..machine_learning.training import run_training_command
19from .training import (
20 ClassificationTrainResult,
21 TorchClassifierTask,
22 prepare_classification_inputs,
23)
25__all__ = ["ImageClassifierHuggingFace"]
28@dataclass
29class _ProcessorTransform:
30 """A picklable per-image transform for DataLoader worker processes."""
32 processor: object
33 fliplr: float = 0.0
35 def __call__(self, image):
36 import torch
37 from PIL import ImageOps
39 image = image.convert("RGB")
40 if self.fliplr and torch.rand(()).item() < self.fliplr:
41 image = ImageOps.mirror(image)
42 inputs = self.processor(images=image, return_tensors="pt")
43 return {name: value.squeeze(0) for name, value in inputs.items()}
46def _class_names(config) -> list[str]:
47 return [
48 config.id2label.get(index, config.id2label.get(str(index), str(index)))
49 for index in range(config.num_labels)
50 ]
53def _set_classes(config, classes: list[str]) -> None:
54 config.num_labels = len(classes)
55 config.id2label = dict(enumerate(classes))
56 config.label2id = {name: index for index, name in enumerate(classes)}
57 config.problem_type = "single_label_classification"
60def _head_modules(model):
61 """Identify task-specific layers outside the Transformers base model."""
62 backbone = model.base_model
63 if backbone is model:
64 raise ValueError(
65 f"{type(model).__name__} does not expose a separate classification backbone"
66 )
67 backbone_modules = {id(module) for module in backbone.modules()}
68 return [
69 module
70 for module in model.modules()
71 if module is not model and id(module) not in backbone_modules
72 ]
75def _reset_head(model) -> None:
76 # Also reset when the new labels have the same count but different meanings.
77 reset = False
78 for module in _head_modules(model):
79 if hasattr(module, "reset_parameters"):
80 module.reset_parameters()
81 reset = True
82 if not reset:
83 raise ValueError(
84 f"Cannot reset the classification head of {type(model).__name__}"
85 )
88class ImageClassifierHuggingFace(TorchClassifierTask):
89 """Fine-tune Transformers image classifiers and return predicted class Links."""
91 prepare_training_inputs = staticmethod(prepare_classification_inputs)
93 def __init__(
94 self,
95 name: str,
96 model: str | Path = "google/vit-base-patch16-224-in21k",
97 **kwargs,
98 ):
99 super().__init__(name, model=model, **kwargs)
100 self.original_model_path: Path | None = None
101 self._categories: list[str] = []
102 self._processor = None
103 self._device = "cpu"
105 @staticmethod
106 def _config_from_dict(data):
107 from transformers import AutoConfig
109 data = copy.deepcopy(data)
110 model_type = data.pop("model_type")
111 return AutoConfig.for_model(model_type, **data)
113 def _from_state(self, config, state, classes=None):
114 from transformers import AutoModelForImageClassification
116 changed = classes is not None and classes != _class_names(config)
117 if classes is not None:
118 _set_classes(config, classes)
119 model = AutoModelForImageClassification.from_config(
120 config, trust_remote_code=False
121 )
122 if changed:
123 # Keep the trained backbone, but never reuse output rows for different labels.
124 backbone = model.base_model
125 if backbone is model:
126 raise ValueError(
127 f"Cannot adapt the classification head of {type(model).__name__}"
128 )
129 prefix = next(
130 name for name, module in model.named_modules() if module is backbone
131 )
132 prefix += "."
133 backbone.load_state_dict(
134 {
135 name[len(prefix) :]: value
136 for name, value in state.items()
137 if name.startswith(prefix)
138 }
139 )
140 else:
141 model.load_state_dict(state)
142 return model
144 def _load(
145 self, *, pretrained: bool = True, classes: list[str] | None = None
146 ) -> None:
147 import torch
148 from transformers import (
149 AutoConfig,
150 AutoImageProcessor,
151 AutoModelForImageClassification,
152 )
154 if isinstance(self.model, torch.nn.Module):
155 if classes is not None and classes != self._categories:
156 self.model = self._from_state(
157 copy.deepcopy(self.model.config), self.model.state_dict(), classes
158 )
159 if classes is not None:
160 _set_classes(self.model.config, classes)
161 self._categories = _class_names(self.model.config)
162 self.model.to(self._device).eval()
163 return
164 if not isinstance(self.model, (str, Path)):
165 raise ValueError(
166 "Model must be a Hugging Face model ID, local model directory, or Collectra checkpoint"
167 )
168 path = Path(self.model).expanduser()
169 if path.is_file():
170 checkpoint = torch.load(path, map_location="cpu", weights_only=True)
171 required = {"backend", "config", "processor", "state_dict", "classes"}
172 if (
173 not isinstance(checkpoint, dict)
174 or not required.issubset(checkpoint)
175 or checkpoint["backend"] != "huggingface"
176 ):
177 raise ValueError(
178 "Expected a Collectra Hugging Face checkpoint with model config, processor, classes, and state_dict; a bare weight file is not sufficient"
179 )
180 config = self._config_from_dict(checkpoint["config"])
181 if checkpoint["classes"] != _class_names(config):
182 raise ValueError(
183 "Checkpoint class names do not match its model configuration"
184 )
185 # AutoImageProcessor resolves built-in processor types from its standard JSON file.
186 with tempfile.TemporaryDirectory() as directory:
187 (Path(directory) / "preprocessor_config.json").write_text(
188 json.dumps(checkpoint["processor"])
189 )
190 self._processor = AutoImageProcessor.from_pretrained(
191 directory,
192 local_files_only=True,
193 trust_remote_code=False,
194 )
195 model = self._from_state(config, checkpoint["state_dict"], classes)
196 self.original_model_path = path.resolve()
197 else:
198 if isinstance(self.model, Path) and not path.is_dir():
199 raise ValueError(
200 f"Model directory or checkpoint does not exist: {path}"
201 )
202 if not path.exists() and (
203 str(self.model).startswith(("/", "./", "../", "~"))
204 or path.suffix in {".pt", ".pth", ".bin", ".safetensors"}
205 ):
206 raise ValueError(
207 f"Model directory or checkpoint does not exist: {path}"
208 )
209 source = str(path) if path.is_dir() else str(self.model)
210 config = AutoConfig.from_pretrained(source, trust_remote_code=False)
211 self._processor = AutoImageProcessor.from_pretrained(
212 source, trust_remote_code=False
213 )
214 changed = classes is not None and classes != _class_names(config)
215 if classes is not None:
216 _set_classes(config, classes)
217 # Local checkpoints always retain their learned weights, as with torchvision.
218 if pretrained or path.is_dir():
219 model = AutoModelForImageClassification.from_pretrained(
220 source,
221 config=config,
222 ignore_mismatched_sizes=changed,
223 trust_remote_code=False,
224 )
225 if changed:
226 _reset_head(model)
227 else:
228 model = AutoModelForImageClassification.from_config(
229 config, trust_remote_code=False
230 )
231 self.original_model_path = path.resolve() if path.is_dir() else None
232 self.model = model.to(self._device)
233 self._categories = _class_names(self.model.config)
234 self.model.eval()
236 def _prepare_training_model(self, classes, pretrained, freeze):
237 self._load(pretrained=pretrained, classes=classes)
238 self.model.requires_grad_(True)
239 if freeze:
240 _head_modules(
241 self.model
242 ) # Check that this architecture exposes a separate backbone.
243 self.model.base_model.requires_grad_(False)
244 if not any(
245 parameter.requires_grad for parameter in self.model.parameters()
246 ):
247 raise ValueError(
248 "No trainable classification head remains after freezing the backbone"
249 )
251 def _transforms(self, *, training=False, fliplr=0.0):
252 return _ProcessorTransform(self._processor, fliplr if training else 0.0)
254 def _forward(self, pixels):
255 inputs = {name: tensor.to(self._device) for name, tensor in pixels.items()}
256 return self.model(**inputs).logits
258 @threading_locked()
259 def run(self, *args: Image) -> Link:
260 """Classify one image or crop and link its predicted class to the input."""
261 import torch
263 if len(args) != 1 or not isinstance(args[0], Image):
264 raise ValueError("This task requires a single Image or ImageCrop input")
265 self._device = self._select_device(getattr(self, "device", ""))
266 self._load()
267 image = args[0]
268 with image.pil() as pixels:
269 inputs = {
270 name: tensor.unsqueeze(0)
271 for name, tensor in self._transforms()(pixels).items()
272 }
273 with torch.inference_mode():
274 class_index = self._forward(inputs).argmax(dim=1).item()
275 class_name = self._categories[class_index]
276 outputs = getattr(self, "output", [])
277 outputs = [outputs] if isinstance(outputs, str) else outputs
278 if outputs and class_name not in outputs:
279 raise ValueError(
280 f"Predicted class {class_name!r} is not a configured output: {outputs}"
281 )
282 return Link(name=class_name, target=image)
284 def _checkpoint(self, epoch, metrics):
285 return {
286 "format_version": 1,
287 "backend": "huggingface",
288 "config": json.loads(self.model.config.to_json_string(use_diff=False)),
289 "processor": json.loads(self._processor.to_json_string()),
290 "classes": self._categories,
291 "state_dict": {
292 name: value.detach().cpu()
293 for name, value in self.model.state_dict().items()
294 },
295 "epoch": epoch,
296 "metrics": metrics,
297 }
299 def _train(self, *images: Image, **kwargs) -> ClassificationTrainResult:
300 result = super()._train(*images, **kwargs)
301 directory = result.save_dir / "weights" / "huggingface"
302 self.model.save_pretrained(directory)
303 self._processor.save_pretrained(directory)
304 return result
306 @command
307 def train(
308 self,
309 inputs: Annotated[
310 list[str],
311 Arg(
312 help="Pipeline data files or directories containing classification labels."
313 ),
314 ],
315 model: Annotated[
316 str,
317 Arg(
318 help="Hugging Face model ID (e.g. google/vit-base-patch16-224-in21k), local save_pretrained directory, or Collectra Hugging Face checkpoint. Omit to keep the configured model."
319 ),
320 ] = "",
321 output: Annotated[
322 Path | None,
323 Arg(
324 help="Directory for extracted train/val images, metrics, and checkpoints."
325 ),
326 ] = None,
327 keep_log: Annotated[
328 bool,
329 Arg(
330 help="Keep extracted images and training logs after saving the model to the pipeline."
331 ),
332 ] = True,
333 validation: Annotated[
334 str,
335 Arg(
336 help="Partition value identifying validation images; omit to train without validation."
337 ),
338 ] = "",
339 exclude: Annotated[
340 str, Arg(help="Partition value identifying images to exclude.")
341 ] = "",
342 epochs: Annotated[int, Arg(help="Number of training epochs.")] = 10,
343 batch: Annotated[int, Arg(help="Number of images per training batch.")] = 16,
344 learning_rate: Annotated[
345 float, Arg(help="Initial learning rate for AdamW.")
346 ] = 0.00005,
347 weight_decay: Annotated[
348 float, Arg(help="Weight decay used for regularization.")
349 ] = 0.0001,
350 pretrained: Annotated[
351 bool,
352 Arg(
353 help="Load pretrained Hub weights; local directories and checkpoints always retain their saved weights. Disable to initialize from a Hub model configuration."
354 ),
355 ] = True,
356 freeze_backbone: Annotated[
357 bool,
358 Arg(help="Train the classification head and freeze the model backbone."),
359 ] = False,
360 early_stop: Annotated[
361 int,
362 Arg(
363 help="Stop after this many epochs without validation-loss improvement; 0 disables."
364 ),
365 ] = 10,
366 min_size: Annotated[
367 int,
368 Arg(
369 help="Skip images or crops with either dimension below this pixel threshold; 0 disables."
370 ),
371 ] = 0,
372 fliplr: Annotated[
373 float,
374 Arg(
375 help="Probability of horizontal flipping during training, from 0 to 1."
376 ),
377 ] = 0.5,
378 workers: Annotated[
379 int, Arg(help="Number of data-loading worker processes.")
380 ] = 0,
381 wandb: Annotated[
382 bool,
383 Arg(help="Log training metrics to Weights & Biases."),
384 ] = False,
385 device: Annotated[
386 str,
387 Arg(
388 help="Torch device such as cpu, cuda, cuda:0, or mps; omit for automatic selection."
389 ),
390 ] = "",
391 seed: Annotated[
392 int, Arg(help="Random seed for model initialization and data loading.")
393 ] = 0,
394 ):
395 """Fine-tune a Hugging Face image classifier using its saved image processor."""
396 result = run_training_command(
397 self,
398 inputs,
399 self._train,
400 keep_log=keep_log,
401 output=output,
402 validation=validation,
403 exclude=exclude,
404 prepare_inputs=self.prepare_training_inputs,
405 model=model,
406 epochs=epochs,
407 batch=batch,
408 learning_rate=learning_rate,
409 weight_decay=weight_decay,
410 pretrained=pretrained,
411 freeze_backbone=freeze_backbone,
412 early_stop=early_stop,
413 min_size=min_size,
414 fliplr=fliplr,
415 workers=workers,
416 wandb=wandb,
417 device=device,
418 seed=seed,
419 )
420 if self.pipeline is not None:
421 saved_model = self.pipeline.data[self.name]["model"]
422 self.original_model_path = (self.pipeline.path / saved_model).resolve()
423 return result