Coverage for collectra/tasks/image_classifier/yolo.py: 85.94%
192 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1from __future__ import annotations
3import shutil
4import tempfile
5from pathlib import Path
6from typing import Annotated, TYPE_CHECKING
8from cappa import Arg
10from PIL import Image as PillowImage
12from collectra.cli import command
13from collectra.types.images import Image, ImageCrop, image_channel_count
14from collectra.types.links import Link
15from collectra.utils import change_dir
17from collectra.logger import get_logger
18from collectra.utils import threading_locked
19from ..machine_learning.training import print_distribution_table, run_training_command
20from ..machine_learning.yolo import YOLOTask, get_yolo_class
21from .training import prepare_classification_inputs
23if TYPE_CHECKING:
24 from ultralytics.engine.results import Results
25 from ultralytics.models import YOLO
26 from ultralytics.utils.metrics import ClassifyMetrics
28__all__ = ["ImageClassifierYOLO"]
30logger = get_logger(__name__)
33def _classification_channels(*roots):
34 """Inspect every prepared image using metadata, before starting training."""
35 detected = {}
36 for root in roots:
37 if root is None:
38 continue
39 for path in sorted(Path(root).rglob("*")):
40 if path.is_file() and path.suffix.lower() in Image.image_types():
41 try:
42 detected[path] = image_channel_count(path)
43 except Exception as error:
44 raise ValueError(
45 f"Cannot determine channels for {path}: {error}"
46 ) from error
47 if not detected:
48 raise ValueError("No prepared classification images found")
49 if len(set(detected.values())) != 1:
50 details = ", ".join(
51 f"{path}: {count} channels" for path, count in detected.items()
52 )
53 raise ValueError(
54 f"Inconsistent input channel counts in classifier dataset: {details}"
55 )
56 channels = next(iter(detected.values()))
57 tensor_loader = channels != 3 or any(
58 p.suffix.lower() in {".tif", ".tiff"} for p in detected
59 )
60 return channels, tensor_loader
63class ImageClassifierYOLO(YOLOTask):
64 prepare_training_inputs = staticmethod(prepare_classification_inputs)
66 @threading_locked()
67 def run(self, *args: Image, **kwargs) -> Link:
68 """Run image classification inference on the provided Image.
70 Args:
71 *args: A single Image (or ImageCrop) to classify.
73 Returns:
74 Link: A link from the predicted class node to the input image.
75 """
76 if len(args) != 1:
77 raise ValueError("This task only supports a single Image input.")
78 if not isinstance(args[0], Image):
79 raise TypeError("Input must be an instance of Image.")
80 image: Image = args[0]
81 self._init_model()
83 if (
84 getattr(
85 getattr(self.model, "model", None),
86 "collectra_tensor_classification",
87 False,
88 )
89 or image.get_path().suffix.lower() in {".tif", ".tiff"}
90 or image_channel_count(image.get_path()) != 3
91 ):
92 from .yolo_multichannel import classify_multichannel
94 predicted_class = classify_multichannel(self.model, image)
95 else:
96 predicted_class = self._predict_rgb(image)
98 outputs = self.output if isinstance(self.output, list) else [self.output]
99 if predicted_class not in outputs:
100 raise ValueError(
101 f"Predicted class {predicted_class!r} is not a configured output: "
102 f"{outputs}"
103 )
104 print(f"Classified as: {predicted_class}")
106 return Link(name=predicted_class, target=image)
108 def _predict_rgb(self, image):
109 if isinstance(image, ImageCrop):
110 with tempfile.TemporaryDirectory() as temp_dir:
111 img = image.pil()
112 img_path = Path(temp_dir) / Path(image.get_path()).name
113 img.save(img_path)
114 results: Results = (self.model(img_path))[0]
115 else:
116 results: Results = (self.model(image.get_path()))[0]
118 if results.probs is None:
119 raise ValueError("No classification probabilities returned.")
121 top1_index = results.probs.top1
122 return results.names[top1_index]
124 @command
125 def train(
126 self,
127 inputs: Annotated[
128 list[str],
129 Arg(
130 help="Pipeline data files or directories containing image classification labels."
131 ),
132 ],
133 model: Annotated[
134 str,
135 Arg(
136 help="YOLO classification weights name or checkpoint path (for example yolo26n-cls.pt), or a model YAML file. Omit to keep the configured model."
137 ),
138 ] = "",
139 output: Annotated[
140 Path | None,
141 Arg(help="Directory for training logs, dataset files, and checkpoints."),
142 ] = None,
143 keep_log: Annotated[
144 bool,
145 Arg(
146 help="Keep the training directory after saving the model to the pipeline."
147 ),
148 ] = True,
149 validation: Annotated[
150 str,
151 Arg(help="Partition value identifying validation examples."),
152 ] = "",
153 exclude: Annotated[
154 str,
155 Arg(help="Partition value identifying examples to exclude from training."),
156 ] = "",
157 epochs: Annotated[
158 int,
159 Arg(help="Number of training epochs."),
160 ] = 100,
161 batch: Annotated[
162 int,
163 Arg(help="Training batch size."),
164 ] = 16,
165 imgsz: Annotated[
166 int,
167 Arg(help="Target training image size in pixels."),
168 ] = 1280,
169 early_stop: Annotated[
170 int,
171 Arg(
172 help="Stop after this many epochs without validation improvement; 0 disables early stopping."
173 ),
174 ] = 50,
175 min_size: Annotated[
176 int,
177 Arg(
178 help="Skip images or crops with either dimension below this many pixels; 0 disables the size threshold."
179 ),
180 ] = 0,
181 learning_rate: Annotated[
182 float,
183 Arg(help="Initial learning rate."),
184 ] = 0.01,
185 weight_decay: Annotated[
186 float,
187 Arg(help="Weight decay used for regularization."),
188 ] = 0.0005,
189 dropout: Annotated[
190 float,
191 Arg(help="Dropout probability for the classification head."),
192 ] = 0.0,
193 erasing: Annotated[
194 float,
195 Arg(help="Probability of random erasing augmentation."),
196 ] = 0.4,
197 auto_augment: Annotated[
198 str,
199 Arg(
200 help="Automatic augmentation policy: randaugment, autoaugment, or augmix. Skipped for TIFF and non-RGB datasets."
201 ),
202 ] = "randaugment",
203 fliplr: Annotated[
204 float,
205 Arg(help="Probability of horizontally flipping an image."),
206 ] = 0.5,
207 flipud: Annotated[
208 float,
209 Arg(help="Probability of vertically flipping an image."),
210 ] = 0.0,
211 hsv_h: Annotated[
212 float,
213 Arg(
214 help="Hue augmentation amount as a fraction; skipped for TIFF and non-RGB datasets."
215 ),
216 ] = 0.015,
217 hsv_s: Annotated[
218 float,
219 Arg(
220 help="Saturation augmentation amount as a fraction; skipped for TIFF and non-RGB datasets."
221 ),
222 ] = 0.7,
223 hsv_v: Annotated[
224 float,
225 Arg(
226 help="Brightness augmentation amount as a fraction; skipped for TIFF and non-RGB datasets."
227 ),
228 ] = 0.4,
229 cls_pw: Annotated[
230 float,
231 Arg(
232 help="Class-imbalance weighting power; 0 disables weighting, 1 uses inverse class frequency."
233 ),
234 ] = 0.0,
235 freeze: Annotated[
236 int | None,
237 Arg(help="Freeze the first N model layers during training."),
238 ] = None,
239 wandb: Annotated[
240 bool, Arg(help="Log training metrics to Weights & Biases.")
241 ] = False,
242 ):
243 """Train this YOLO image classifier."""
244 return run_training_command(
245 self,
246 inputs,
247 self._train,
248 keep_log=keep_log,
249 output=output,
250 validation=validation,
251 exclude=exclude,
252 prepare_inputs=self.prepare_training_inputs,
253 model=model,
254 epochs=epochs,
255 batch=batch,
256 imgsz=imgsz,
257 early_stop=early_stop,
258 min_size=min_size,
259 learning_rate=learning_rate,
260 weight_decay=weight_decay,
261 dropout=dropout,
262 erasing=erasing,
263 auto_augment=auto_augment,
264 fliplr=fliplr,
265 flipud=flipud,
266 hsv_h=hsv_h,
267 hsv_s=hsv_s,
268 hsv_v=hsv_v,
269 cls_pw=cls_pw,
270 freeze=freeze,
271 wandb=wandb,
272 )
274 def _train(self, *images: Image, **kwargs) -> ClassifyMetrics | None:
275 YOLO = get_yolo_class()
277 if kwargs.get("model"):
278 self.model = kwargs["model"]
279 self._init_model()
280 log = kwargs.get("log", None)
281 if log is None:
282 raise ValueError(
283 "Log directory must be specified in kwargs with key 'log'."
284 )
285 log = kwargs["base_folder"] / log
286 log.mkdir(parents=True, exist_ok=True)
287 print(f"Training files will be saved to: {log}")
288 # Derive classes from training data, not from pipeline children nodes
289 classes = sorted(set(img.name for img in images))
290 kwargs.pop("classes", None)
291 if not isinstance(self.model, YOLO):
292 raise ValueError("Model must be a YOLO instance for training.")
294 print(f"Training with model: {self.model.model_name}")
296 validation = kwargs.get("validation", "")
297 exclude = kwargs.get("exclude", "")
298 train_dir, val_dir = self._prepare_assets(
299 classes,
300 log,
301 validation,
302 exclude,
303 *images,
304 min_size=int(kwargs.get("min_size", 0)),
305 )
306 return self._train_fold(train_dir, val_dir, classes, log, kwargs)
308 def _train_fold(
309 self,
310 train: list[str] | Path,
311 val: list[str] | Path,
312 classes: list[str],
313 log: Path,
314 kwargs: dict,
315 fold_count: int | None = None,
316 ) -> ClassifyMetrics:
317 YOLO = get_yolo_class()
319 if not isinstance(self.model, YOLO):
320 raise ValueError("Expected model to be a YOLO instance for training.")
321 _, tensor_loader = _classification_channels(train, val)
322 kwargs["config_file"] = self._prepare_yolo_config(log, classes, train, val)
323 kwargs["log"] = f"{log.name}"
324 if fold_count:
325 kwargs["log"] += f"_fold_{fold_count}"
326 params = self._prepare_params(**kwargs)
327 if tensor_loader:
328 from .yolo_multichannel import MultichannelClassificationTrainer
330 params["trainer"] = MultichannelClassificationTrainer
331 with (
332 change_dir(kwargs["base_folder"]),
333 self._wandb_logging(kwargs.get("wandb", False)),
334 ):
335 results = self.model.train(**params)
336 if results is None:
337 raise Exception("[red]Training failed, no results returned.[/red]")
338 self._reload()
339 return results
341 def _prepare_assets(
342 self,
343 classes: list[str],
344 log: Path,
345 validation_flag: str,
346 exclude_flag: str,
347 *images: Image,
348 min_size: int = 0,
349 ) -> tuple[Path, Path]:
350 """Prepare assets for YOLO classification training.
352 Creates directory structure: log/train/{class}/ and log/val/{class}/
354 Args:
355 classes: List of class names.
356 log: Path to store the assets for training.
357 validation_flag: The value of the validation flag to identify validation images.
358 exclude_flag: The value of the exclude flag to identify images to be excluded.
359 *images: Image instances to be prepared.
361 Returns:
362 Tuple of (train_dir, val_dir) paths.
363 """
364 train_dir = log / "train"
365 val_dir = log / "val"
367 if min_size < 0:
368 raise ValueError("min_size cannot be negative")
370 for cls in classes:
371 (train_dir / cls).mkdir(parents=True, exist_ok=True)
372 (val_dir / cls).mkdir(parents=True, exist_ok=True)
374 source_images: dict[Path, PillowImage.Image] = {}
375 for img in images:
376 if exclude_flag and img.partition == exclude_flag:
377 continue
379 cls_name = img.name
380 if cls_name not in classes:
381 continue
383 if img.partition == validation_flag:
384 target_dir = val_dir / cls_name
385 else:
386 target_dir = train_dir / cls_name
388 target = img.resolve() if type(img) is Link else img
389 src = target.get_path()
390 dst = target_dir / src.name
391 pixel_box = None
392 if isinstance(target, ImageCrop):
393 dst = target_dir / f"{src.stem}-{img.id}{src.suffix}"
394 left, upper, right, bottom = target.coordinates()
395 pixel_box = (left, upper, right, bottom)
396 width = right - left
397 height = bottom - upper
398 elif min_size:
399 width = int(target.width)
400 height = int(target.height)
402 if type(img) is Link:
403 dst = target_dir / f"{src.stem}-{img.id}{src.suffix}"
405 empty_crop = isinstance(target, ImageCrop) and (width <= 0 or height <= 0)
406 below_minimum = min_size and (width < min_size or height < min_size)
407 if empty_crop or below_minimum:
408 effective_minimum = min_size or 1
409 parent_id = img.parent_id if type(img) is Link else ""
410 logger.warning(
411 "Skipping classifier image %s%s from %s: image is %dx%d "
412 "pixels (minimum %d)%s",
413 img.id,
414 f" -> {parent_id}" if parent_id else "",
415 img.source_file or src,
416 width,
417 height,
418 effective_minimum,
419 f", pixel_box={pixel_box}" if pixel_box else "",
420 )
421 continue
423 if not dst.exists():
424 if isinstance(target, ImageCrop) and src.suffix.lower() in {
425 ".tif",
426 ".tiff",
427 }:
428 import numpy as np
429 import tifffile
430 from .torchvision import _target_tiff_pixels
432 pixels = _target_tiff_pixels(target)
433 tifffile.imwrite(
434 dst,
435 np.ascontiguousarray(pixels.transpose(2, 0, 1)),
436 photometric="minisblack",
437 metadata={"axes": "CYX"},
438 )
439 elif isinstance(target, ImageCrop):
440 source_path = src.resolve()
441 if source_path not in source_images:
442 with PillowImage.open(source_path) as source:
443 source.load()
444 source_images[source_path] = source.copy()
445 with source_images[source_path].crop(pixel_box) as cropped:
446 rotated = cropped.rotate(
447 target.orientation.to_degree(), expand=True
448 )
449 try:
450 rotated.save(dst)
451 finally:
452 rotated.close()
453 else:
454 shutil.copy(src, dst)
456 for source in source_images.values():
457 source.close()
459 return train_dir, val_dir
461 def _prepare_yolo_config(
462 self,
463 log: Path,
464 classes: list[str],
465 train: list[str] | Path,
466 val: list[str] | Path,
467 ) -> Path:
468 """For classification, data param is the parent directory containing train/ and val/.
470 Also prints distribution info.
471 """
472 self._check_distribution(log, classes)
473 return log
475 def _check_distribution(
476 self,
477 log: Path,
478 metadata: dict | list[str],
479 ) -> None:
480 classes = (
481 metadata
482 if isinstance(metadata, list)
483 else list(metadata.get("classes", []))
484 )
485 train_dir = log / "train"
486 val_dir = log / "val"
488 train_counts = {}
489 validation_counts = {}
490 for cls_name in classes:
491 train_counts[cls_name] = sum(
492 1
493 for ext in Image.image_types()
494 for _ in (train_dir / cls_name).glob(f"*{ext}")
495 )
496 validation_counts[cls_name] = sum(
497 1
498 for ext in Image.image_types()
499 for _ in (val_dir / cls_name).glob(f"*{ext}")
500 )
501 print_distribution_table(
502 "Classification Distribution", classes, train_counts, validation_counts
503 )
505 def _prepare_params(self, **kwargs) -> dict:
506 import platform
508 import torch
510 params = {
511 "name": kwargs["log"],
512 "data": kwargs["config_file"],
513 "project": kwargs["project"],
514 "device": (
515 "mps"
516 if platform.system() == "Darwin"
517 else "cuda" if torch.cuda.is_available() else "cpu"
518 ),
519 "epochs": kwargs.get("epochs", 100),
520 "imgsz": kwargs.get("imgsz", 640),
521 "patience": kwargs.get("early_stop", 50),
522 "batch": kwargs.get("batch", 16),
523 "lr0": kwargs.get("learning_rate", 0.01),
524 "weight_decay": kwargs.get("weight_decay", 0.0005),
525 "dropout": kwargs.get("dropout", 0.0),
526 "erasing": kwargs.get("erasing", 0.4),
527 "auto_augment": kwargs.get("auto_augment", "randaugment"),
528 "fliplr": kwargs.get("fliplr", 0.5),
529 "flipud": kwargs.get("flipud", 0.0),
530 "hsv_h": kwargs.get("hsv_h", 0.015),
531 "hsv_s": kwargs.get("hsv_s", 0.7),
532 "hsv_v": kwargs.get("hsv_v", 0.4),
533 "cls_pw": kwargs.get("cls_pw", 0.0),
534 }
535 if kwargs.get("freeze") is not None:
536 params["freeze"] = kwargs["freeze"]
537 return params