Coverage for collectra/tasks/image_classifier/training.py: 83.79%
290 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Shared preparation of labelled classifier inputs."""
3import csv
4import json
5import math
6import os
7from collections.abc import Mapping
8from concurrent.futures import ThreadPoolExecutor, as_completed
9from dataclasses import dataclass
10from pathlib import Path
11from rich.console import Console
12from rich.progress import (
13 BarColumn,
14 MofNCompleteColumn,
15 Progress,
16 SpinnerColumn,
17 TextColumn,
18 TimeElapsedColumn,
19 TimeRemainingColumn,
20)
22from collectra.logger import get_logger
23from ..base import Task
24from ..machine_learning.training import print_distribution_table
25from collectra.types.images import Image, ImageCrop
26from collectra.types.links import Link
29def prepare_classification_inputs(
30 processed_inputs: list,
31 processed_parents: list,
32 input_maps: dict[str, list],
33 parent_input_maps: dict[str, list],
34) -> list:
35 """Use parent images labelled by their child Link node names."""
36 labeled_parents = []
37 for file_name, children in input_maps.items():
38 parents = parent_input_maps.get(file_name, [])
39 artefacts = {item.id: item for item in [*parents, *children]}
40 for item in artefacts.values():
41 if type(item) is Link:
42 item.bind(artefacts)
43 for child in children:
44 if type(child) is not Link:
45 raise TypeError(
46 f"Classifier label {child.id!r} in {file_name!r} must be "
47 f"a Link, got {type(child).__name__}"
48 )
49 if not child.parents:
50 raise ValueError(
51 f"Classifier Link {child.id!r} in {file_name!r} has no parent"
52 )
53 parent_id = (
54 child.parents if isinstance(child.parents, str) else child.parents[0]
55 )
56 parent = artefacts.get(parent_id)
57 if parent is None:
58 raise ValueError(
59 f"Classifier Link {child.id!r} in {file_name!r} points to "
60 f"missing parent {parent_id!r}; available parent IDs: "
61 f"{[item.id for item in parents]}"
62 )
63 child.target = parent
64 try:
65 target = child.resolve()
66 except RuntimeError as error:
67 raise ValueError(
68 f"Classifier Link {child.id!r} in {file_name!r} cannot "
69 f"resolve parent chain: {error}"
70 ) from error
71 if not isinstance(target, Image):
72 raise TypeError(
73 f"Classifier Link {child.id!r} resolves to "
74 f"{type(target).__name__}, expected Image or ImageCrop"
75 )
76 child.partition = target.partition
77 labeled_parents.append(child)
78 return labeled_parents
81logger = get_logger(__name__)
82console = Console()
85@dataclass
86class ClassificationTrainResult:
87 save_dir: Path
88 results_dict: dict
91class TorchClassifierTask(Task):
92 """Shared export, optimization, and metrics for PyTorch classification backends."""
94 @staticmethod
95 def _select_device(device: str = "") -> str:
96 import torch
98 if device:
99 return device
100 if torch.cuda.is_available():
101 return "cuda"
102 if torch.backends.mps.is_available():
103 return "mps"
104 return "cpu"
106 def _reload(self) -> None:
107 if self.original_model_path is None:
108 raise ValueError("No checkpoint is available to reload")
109 self.model = self.original_model_path
110 self._load()
112 def _prepare_assets(
113 self,
114 classes,
115 log,
116 validation,
117 exclude,
118 *images,
119 min_size=0,
120 workers=0,
121 use_existing=False,
122 ):
123 """Materialize labelled images and crops in train/<class> and val/<class>."""
124 if min_size < 0:
125 raise ValueError("min_size cannot be negative")
126 for name in classes:
127 if not name or name in {".", ".."} or "/" in name or "\\" in name:
128 raise ValueError(f"Class name must be a directory name: {name!r}")
129 train_dir, val_dir = log / "train", log / "val"
130 if use_existing and train_dir.exists() and val_dir.exists():
131 logger.info(
132 "Reusing existing classifier assets in %s and %s",
133 train_dir,
134 val_dir,
135 )
136 counts = {split: {name: 0 for name in classes} for split in ("train", "val")}
137 for split, directory in (("train", train_dir), ("val", val_dir)):
138 for name in classes:
139 class_dir = directory / name
140 if class_dir.exists():
141 counts[split][name] = sum(
142 1 for path in class_dir.iterdir() if path.is_file()
143 )
144 print_distribution_table(
145 "Class Distribution", classes, counts["train"], counts["val"]
146 )
147 missing = [name for name in classes if not counts["train"][name]]
148 if missing:
149 raise ValueError(
150 "No training images for classes in existing dataset: "
151 f"{', '.join(missing)}"
152 )
153 return train_dir, val_dir
154 for directory in (train_dir, val_dir):
155 if directory.exists() and any(directory.iterdir()):
156 raise ValueError(
157 f"Dataset directory is not empty: {directory}. Use a new output directory."
158 )
159 counts = {split: {name: 0 for name in classes} for split in ("train", "val")}
160 for directory in (train_dir, val_dir):
161 for name in classes:
162 (directory / name).mkdir(parents=True, exist_ok=True)
163 jobs = []
164 for index, image in enumerate(images):
165 if exclude and image.partition == exclude:
166 continue
167 if image.name not in classes:
168 raise ValueError(
169 f"Training label {image.name!r} is not a configured class"
170 )
171 target = image.resolve() if type(image) is Link else image
172 if not isinstance(target, Image):
173 raise TypeError(
174 f"Classifier input {image.id!r} does not resolve to an Image"
175 )
176 if isinstance(target, ImageCrop):
177 left, top, right, bottom = target.coordinates()
178 if right <= left or bottom <= top:
179 logger.warning("Skipping empty classifier crop %s", image.id)
180 continue
181 split = "val" if validation and image.partition == validation else "train"
182 destination = log / split / image.name / f"{index:08d}"
183 jobs.append((image.id, image.name, split, target, destination))
185 export_workers = max(0, int(workers))
186 tiff_decode_workers = 1 if export_workers > 1 else None
188 def export_job(target, destination):
189 return self._export_image(
190 target,
191 destination,
192 min_size,
193 tiff_maxworkers=tiff_decode_workers,
194 )
196 with Progress(
197 SpinnerColumn(),
198 TextColumn("{task.description}"),
199 BarColumn(),
200 MofNCompleteColumn(),
201 TextColumn("•"),
202 TimeElapsedColumn(),
203 TextColumn("•"),
204 TimeRemainingColumn(),
205 ) as progress:
206 task_id = progress.add_task(
207 "Exporting classifier assets...", total=len(jobs)
208 )
209 if export_workers > 0:
210 with ThreadPoolExecutor(max_workers=export_workers) as executor:
211 futures = {
212 executor.submit(export_job, target, destination):
213 (image_id, class_name, split)
214 for image_id, class_name, split, target, destination in jobs
215 }
216 for future in as_completed(futures):
217 image_id, class_name, split = futures[future]
218 progress.update(task_id, description=f"Exporting {image_id}")
219 if future.result():
220 counts[split][class_name] += 1
221 progress.advance(task_id)
222 else:
223 for image_id, class_name, split, target, destination in jobs:
224 progress.update(task_id, description=f"Exporting {image_id}")
225 if export_job(target, destination):
226 counts[split][class_name] += 1
227 progress.advance(task_id)
228 print_distribution_table(
229 "Class Distribution", classes, counts["train"], counts["val"]
230 )
231 missing = [name for name in classes if not counts["train"][name]]
232 if missing:
233 raise ValueError(f"No training images for classes: {', '.join(missing)}")
234 return train_dir, val_dir
236 def _export_image(
237 self, target, destination, min_size, tiff_maxworkers: int | None = None
238 ):
239 with target.pil() as pixels:
240 if min(pixels.size) < max(1, min_size):
241 logger.warning(
242 "Skipping classifier image %s below minimum size %d",
243 target.id,
244 min_size,
245 )
246 return False
247 pixels.convert("RGB").save(destination.with_suffix(".png"))
248 return True
250 def _make_datasets(self, train_dir, val_dir, augmentation):
251 from torchvision.datasets import ImageFolder
253 return (
254 ImageFolder(
255 train_dir, transform=self._transforms(training=True, **augmentation)
256 ),
257 ImageFolder(val_dir, transform=self._transforms(), allow_empty=True),
258 )
260 def _train(self, *images: Image, **kwargs) -> ClassificationTrainResult:
261 import torch
262 from torch.utils.data import DataLoader
264 epochs, batch = int(kwargs.get("epochs", 10)), int(kwargs.get("batch", 16))
265 workers, patience = int(kwargs.get("workers", 8)), int(
266 kwargs.get("early_stop", 10)
267 )
268 fliplr = float(kwargs.get("fliplr", 0.5))
269 augmentation = {"fliplr": fliplr}
270 if "flipud" in kwargs:
271 flipud = float(kwargs["flipud"])
272 if not 0 <= flipud <= 1:
273 raise ValueError("flipud must be in [0, 1]")
274 augmentation["flipud"] = flipud
275 if (
276 epochs < 1
277 or batch < 1
278 or workers < 0
279 or patience < 0
280 or not 0 <= fliplr <= 1
281 ):
282 raise ValueError(
283 "epochs and batch must be positive; workers/early_stop nonnegative; fliplr in [0, 1]"
284 )
285 if "log" not in kwargs or "base_folder" not in kwargs:
286 raise ValueError("Training requires log and base_folder directories")
287 available_workers = max(1, os.cpu_count() or 1)
288 if workers > 0 and workers > available_workers:
289 logger.info(
290 "Capping workers from %d to available CPU count %d",
291 workers,
292 available_workers,
293 )
294 workers = available_workers
295 torch.manual_seed(int(kwargs.get("seed", 0)))
296 self._device = self._select_device(kwargs.get("device", ""))
297 if kwargs.get("model"):
298 self.model = kwargs["model"]
299 configured = kwargs.get("classes") or getattr(self, "output", [])
300 if isinstance(configured, str):
301 configured = [configured]
302 classes = sorted(
303 set(
304 configured
305 or [
306 image.name
307 for image in images
308 if not kwargs.get("exclude") or image.partition != kwargs["exclude"]
309 ]
310 )
311 )
312 if not classes:
313 raise ValueError("No classes provided/found for classifier training")
314 log = (Path(kwargs["base_folder"]) / kwargs["log"]).resolve()
315 logger.info(
316 "Preparing classifier assets at %s (use_existing=%s, workers=%d)",
317 log,
318 bool(kwargs.get("use_existing", False)),
319 workers,
320 )
321 train_dir, val_dir = self._prepare_assets(
322 classes,
323 log,
324 kwargs.get("validation", ""),
325 kwargs.get("exclude", ""),
326 *images,
327 min_size=kwargs.get("min_size", 0),
328 workers=workers,
329 use_existing=bool(kwargs.get("use_existing", False)),
330 )
331 logger.info("Asset preparation complete: train=%s val=%s", train_dir, val_dir)
332 # Load preprocessing config before dataset construction so channel
333 # detection and normalization logic can run on real sample files.
334 logger.info(
335 "Loading model for preprocessing (model=%s, pretrained=%s)",
336 self.model,
337 kwargs.get("pretrained", True),
338 )
339 self._load(pretrained=kwargs.get("pretrained", True))
340 logger.info("Building datasets and detecting channel count")
341 train_data, val_data = self._make_datasets(train_dir, val_dir, augmentation)
342 # Torchvision stores normalization here; other backends (e.g. Hugging
343 # Face) own their preprocessing elsewhere. Logging must not require it.
344 preprocessing = getattr(self, "_preprocessing", {})
345 logger.info(
346 "Datasets ready: train=%d, val=%d, detected_channels=%s, mean_len=%d, std_len=%d",
347 len(train_data),
348 len(val_data),
349 getattr(self, "_input_channels", "unknown"),
350 len(preprocessing.get("mean", [])),
351 len(preprocessing.get("std", [])),
352 )
353 freeze = kwargs.get("freeze_backbone", False)
354 logger.info("Preparing training model (freeze_backbone=%s)", freeze)
355 self._prepare_training_model(classes, kwargs.get("pretrained", True), freeze)
356 if hasattr(self.model, "conv1") and hasattr(self.model.conv1, "in_channels"):
357 logger.info(
358 "Model conv1 channels after preparation: %s",
359 self.model.conv1.in_channels,
360 )
361 train_loader = DataLoader(
362 train_data, batch_size=batch, shuffle=True, num_workers=workers
363 )
364 val_loader = (
365 DataLoader(val_data, batch_size=batch, num_workers=workers)
366 if len(val_data)
367 else None
368 )
369 logger.info(
370 "DataLoaders ready: train_batches=%d, val_batches=%d, batch=%d, workers=%d",
371 len(train_loader),
372 len(val_loader) if val_loader is not None else 0,
373 batch,
374 workers,
375 )
376 if val_loader is None and kwargs.get("validation"):
377 raise ValueError(
378 f"No validation images match partition {kwargs['validation']!r} "
379 "after exclusions and size filtering"
380 )
381 if val_loader is None:
382 logger.warning(
383 "No validation images; best checkpoint will use training loss, with early stopping disabled"
384 )
385 optimizer = torch.optim.AdamW(
386 [
387 parameter
388 for parameter in self.model.parameters()
389 if parameter.requires_grad
390 ],
391 lr=float(kwargs.get("learning_rate", 0.0001)),
392 weight_decay=float(kwargs.get("weight_decay", 0.0001)),
393 )
394 weights_dir = log / "weights"
395 weights_dir.mkdir(parents=True, exist_ok=True)
396 (weights_dir / "classes.json").write_text(
397 json.dumps({"classes": classes}, indent=2)
398 )
399 wandb_run = None
400 if bool(kwargs.get("wandb", False)):
401 try:
402 import wandb
404 wandb_config = {
405 key: str(value) if isinstance(value, Path) else value
406 for key, value in kwargs.items()
407 }
408 wandb_config["classes"] = classes
409 wandb_run = wandb.init(
410 project=str(kwargs.get("project", "collectra-classifier")),
411 name=log.name,
412 dir=str(kwargs["base_folder"]),
413 config=wandb_config,
414 reinit=True,
415 )
416 except Exception as error:
417 logger.warning("W&B init skipped: %s", error)
418 history, best_accuracy, best_metrics, stale = [], -math.inf, {}, 0
419 try:
420 with Progress(
421 SpinnerColumn(),
422 TextColumn("{task.description}"),
423 BarColumn(),
424 MofNCompleteColumn(),
425 TextColumn("•"),
426 TimeElapsedColumn(),
427 TextColumn("•"),
428 TimeRemainingColumn(),
429 ) as progress:
430 epoch_task_id = progress.add_task("Training epochs...", total=epochs)
431 for epoch in range(1, epochs + 1):
432 progress.update(epoch_task_id, description=f"Training epoch {epoch}/{epochs}")
433 logger.info("Starting epoch %d/%d", epoch, epochs)
434 self.model.train(
435 not freeze
436 ) # Frozen BatchNorm statistics must remain frozen too.
437 train_metrics = self._epoch(train_loader, optimizer)
438 self.model.eval()
439 with torch.inference_mode():
440 val_metrics = (
441 self._epoch(val_loader) if val_loader is not None else None
442 )
443 metrics = {
444 "epoch": epoch,
445 **{f"train_{k}": v for k, v in train_metrics.items()},
446 }
447 metrics.update({f"val_{k}": v for k, v in (val_metrics or {}).items()})
448 history.append(metrics)
449 checkpoint = self._checkpoint(epoch, metrics)
450 torch.save(checkpoint, weights_dir / "last.pt")
451 score = val_metrics["accuracy"] if val_metrics else train_metrics["accuracy"]
452 if not math.isfinite(score):
453 raise ValueError("Training produced a non-finite accuracy")
454 if score > best_accuracy:
455 best_accuracy, best_metrics, stale = score, dict(metrics), 0
456 torch.save(checkpoint, weights_dir / "best.pt")
457 best_print = " best"
458 else:
459 stale += 1
460 best_print = ""
461 if wandb_run is not None:
462 wandb_run.log(metrics)
463 logger.info("Epoch %d/%d: %s", epoch, epochs, metrics)
464 metrics_line = (
465 f"Epoch {epoch}/{epochs} "
466 f"train_loss={metrics['train_loss']:.4f} "
467 f"train_acc={metrics['train_accuracy']:.4f} "
468 f"train_top5={metrics['train_top5_accuracy']:.4f}"
469 )
470 if val_metrics is not None:
471 metrics_line += (
472 f" val_loss={metrics['val_loss']:.4f} "
473 f"val_acc={metrics['val_accuracy']:.4f} "
474 f"val_top5={metrics['val_top5_accuracy']:.4f}"
475 )
476 console.print(metrics_line + best_print)
477 with (log / "history.csv").open("w", newline="") as stream:
478 writer = csv.DictWriter(stream, fieldnames=list(metrics))
479 writer.writeheader()
480 writer.writerows(history)
481 progress.advance(epoch_task_id)
482 if val_loader is not None and patience and stale >= patience:
483 console.print(
484 f"Early stopping at epoch {epoch}: no val-accuracy improvement in {patience} epoch(s)."
485 )
486 break
487 finally:
488 if wandb_run is not None:
489 wandb_run.finish()
490 self.original_model_path = weights_dir / "best.pt"
491 self._reload()
492 best_metrics.update(classes=classes, epochs_completed=len(history))
493 (log / "metrics.json").write_text(json.dumps(best_metrics, indent=2))
494 return ClassificationTrainResult(save_dir=log, results_dict=best_metrics)
496 def _epoch(self, loader, optimizer=None) -> dict:
497 import torch
498 from torch.nn import functional as F
499 from sklearn.metrics import precision_recall_fscore_support
501 total, correct, correct5, total_loss = 0, 0, 0, 0.0
502 all_preds, all_labels = [], []
503 first_batch = True
504 phase = "train" if optimizer is not None else "val"
505 for pixels, labels in loader:
506 # Torchvision batches are tensors; Hugging Face batches are dicts
507 # of processor outputs keyed by model input name.
508 tensors = pixels if isinstance(pixels, Mapping) else {"pixels": pixels}
509 shapes = {name: tuple(tensor.shape) for name, tensor in tensors.items()}
510 if first_batch:
511 approximate_mib = sum(
512 tensor.numel() * tensor.element_size() for tensor in tensors.values()
513 ) / (1024 * 1024)
514 logger.info(
515 "%s first batch: input_shapes=%s, labels_shape=%s, dtypes=%s, approx_input_mib=%.1f",
516 phase,
517 shapes,
518 tuple(labels.shape),
519 {name: tensor.dtype for name, tensor in tensors.items()},
520 approximate_mib,
521 )
522 first_batch = False
523 labels = labels.to(self._device)
524 if optimizer is not None:
525 optimizer.zero_grad(set_to_none=True)
526 try:
527 logits = self._forward(pixels)
528 loss = F.cross_entropy(logits, labels)
529 except RuntimeError as error:
530 raise RuntimeError(
531 f"{phase} step failed for input shapes {shapes} "
532 f"on device {self._device}: {error}"
533 ) from error
534 if optimizer is not None:
535 loss.backward()
536 optimizer.step()
537 total += labels.size(0)
538 total_loss += loss.item() * labels.size(0)
539 correct += (logits.argmax(1) == labels).sum().item()
540 correct5 += (
541 (
542 logits.topk(min(5, len(self._categories)), dim=1).indices
543 == labels[:, None]
544 )
545 .any(1)
546 .sum()
547 .item()
548 )
549 all_preds.extend(logits.argmax(1).detach().cpu().tolist())
550 all_labels.extend(labels.cpu().tolist())
551 precision, recall, f1, _ = precision_recall_fscore_support(
552 all_labels, all_preds, average="macro", zero_division=0
553 )
554 return {
555 "precision": float(precision),
556 "recall": float(recall),
557 "f1": float(f1),
558 "loss": total_loss / total,
559 "accuracy": correct / total,
560 "top5_accuracy": correct5 / total,
561 }