API Reference: Engine Components

Core processing engines coordinating model training loops and inference predictions.

---

CoreTrainer

corecv.engine.trainer.CoreTrainer(model, optimizer, loss_fn, train_dataloader, val_dataloader=None, device=None, gradient_accumulation_steps=1, max_grad_norm=1.0, use_amp=None, amp_dtype=torch.float16, ema_decay=0.9999, ema_start_epoch=0, scheduler=None, scheduler_interval='epoch', log_interval=50, output_dir='./checkpoints', train_metrics=None, val_metrics=None, model_config=None)

Unified training engine for CoreCV models.

Coordinates the complete training loop with support for:

• Automatic Mixed Precision (AMP) via torch.amp.autocast
• Gradient accumulation
• Gradient clipping
• Exponential Moving Average (EMA) of model weights
• Checkpoint save/load with optimizer, scheduler, epoch state
• Metrics integration with corecv.metrics objects

Parameters:

Name	Type	Description	Default
model	Module	The CoreCV model to train (nn.Module).	required
optimizer	Optimizer	PyTorch optimizer.	required
loss_fn	object	Loss function (callable taking preds, targets and returning a scalar loss tensor).	required
train_dataloader	DataLoader	Training :class:DataLoader.	required
val_dataloader	DataLoader | None	Optional validation :class:DataLoader.	None
device	device | None	:class:torch.device for training ('cuda' or 'cpu'). If None, auto-detects 'cuda' when available.	None
gradient_accumulation_steps	int	Number of steps to accumulate gradients before each optimizer update. Default 1.	1
max_grad_norm	float | None	Max norm for gradient clipping. None disables clipping. Default 1.0.	1.0
use_amp	bool | None	Enable Automatic Mixed Precision. Defaults to True when device is CUDA, False otherwise.	None
amp_dtype	dtype	AMP computation dtype. Default torch.float16.	float16
ema_decay	float | None	EMA decay factor. None disables EMA. Default 0.9999.	0.9999
ema_start_epoch	int	Epoch at which to start updating EMA weights (allows EMA to begin after a warmup period). Default 0.	0
scheduler	object | None	Optional LR scheduler.	None
scheduler_interval	str	When to step the scheduler. One of 'step' (after each optimizer step) or 'epoch' (after each epoch). Default 'epoch'.	'epoch'
log_interval	int	Log training metrics every log_interval optimizer steps. Default 50.	50
output_dir	str	Directory for checkpoints. Default './checkpoints'.	'./checkpoints'
train_metrics	Module | None	Optional metrics object (from corecv.metrics) with update(preds, targets) and compute() methods. Accumulated during training.	None
val_metrics	Module | None	Optional metrics object (from corecv.metrics) with update(preds, targets) and compute() methods. Accumulated during validation.	None
model_config	dict[str, Any] | None	Optional dictionary containing the model configuration (e.g. architecture hyperparameters). Stored in checkpoints under the 'model_config' key for reproducibility. Default None.	None

Example::

>>> trainer = CoreTrainer(
...     model=model,
...     optimizer=optimizer,
...     loss_fn=nn.CrossEntropyLoss(),
...     train_dataloader=train_loader,
...     val_dataloader=val_loader,
...     device=torch.device("cuda"),
... )
>>> history = trainer.fit(num_epochs=100)
>>> with trainer.model_ema:
...     output = trainer.model(test_inputs)

Initialise the CoreTrainer with model, optimizer, and training configuration.

Source code in src/corecv/engine/trainer.py

def __init__(  # noqa: PLR0913
    self,
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    loss_fn: object,
    train_dataloader: DataLoader,
    val_dataloader: DataLoader | None = None,
    device: torch.device | None = None,
    gradient_accumulation_steps: int = 1,
    max_grad_norm: float | None = 1.0,
    use_amp: bool | None = None,
    amp_dtype: torch.dtype = torch.float16,
    ema_decay: float | None = 0.9999,
    ema_start_epoch: int = 0,
    scheduler: object | None = None,
    scheduler_interval: str = "epoch",
    log_interval: int = 50,
    output_dir: str = "./checkpoints",
    train_metrics: nn.Module | None = None,
    val_metrics: nn.Module | None = None,
    model_config: dict[str, Any] | None = None,
) -> None:
    """Initialise the CoreTrainer with model, optimizer, and training configuration."""
    # ---- Device ----------------------------------------------------
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    self.device: torch.device = device

    # ---- AMP -------------------------------------------------------
    if use_amp is None:
        use_amp = device.type == "cuda"
    self.use_amp: bool = use_amp
    self.amp_dtype: torch.dtype = amp_dtype

    # ---- Core components -------------------------------------------
    self.model: nn.Module = model.to(self.device)
    self.optimizer: torch.optim.Optimizer = optimizer
    self.loss_fn: object = loss_fn
    self.train_dataloader: DataLoader = train_dataloader
    self.val_dataloader: DataLoader | None = val_dataloader

    # ---- Training hyperparameters ----------------------------------
    self.gradient_accumulation_steps: int = int(gradient_accumulation_steps)
    self.max_grad_norm: float | None = max_grad_norm

    self.ema_decay: float | None = ema_decay
    self.ema_start_epoch: int = int(ema_start_epoch)

    self.scheduler: object | None = scheduler
    if scheduler_interval not in ("step", "epoch"):
        msg = (
            f"scheduler_interval must be 'step' or 'epoch', "
            f"got {scheduler_interval!r}"
        )
        raise ValueError(msg)
    self.scheduler_interval: str = scheduler_interval

    self.log_interval: int = int(log_interval)
    self.output_dir: Path = Path(output_dir)

    # ---- Metrics ---------------------------------------------------
    self.train_metrics: nn.Module | None = train_metrics
    self.val_metrics: nn.Module | None = val_metrics

    # ---- Model config ----------------------------------------------
    self.model_config: dict[str, Any] | None = model_config

    # ---- AMP gradient scaler ---------------------------------------
    self.scaler: GradScaler = GradScaler(device=device.type, enabled=use_amp)

    # ---- EMA shadow parameters -------------------------------------
    self._ema_params: dict[str, torch.Tensor] = {}
    if self.ema_decay is not None:
        self._init_ema()

    # ---- Output directory ------------------------------------------
    self.output_dir.mkdir(parents=True, exist_ok=True)

model_ema property

Get a context manager that temporarily applies EMA weights.

Use this for inference with EMA-averaged weights::

with trainer.model_ema:
    output = trainer.model(input)

The EMA weights are applied to the model upon entering the with block and automatically restored upon exit.

Returns:

Name	Type	Description
An	EMAContext	class:EMAContext context manager.

fit(num_epochs)

Run the complete training loop for a given number of epochs.

For each epoch:

1. Calls :meth:train_one_epoch
2. Calls :meth:validate (if val_dataloader is available)
3. Steps the epoch-based scheduler (if configured with scheduler_interval='epoch')
4. Saves a checkpoint

Parameters:

Name	Type	Description	Default
num_epochs	int	Number of epochs to train.	required

Returns:

Type	Description
dict[str, list]	History dictionary with keys 'train' and 'val', each
dict[str, list]	containing a list of per-epoch metric dictionaries.

Source code in src/corecv/engine/trainer.py

def fit(self, num_epochs: int) -> dict[str, list]:
    """Run the complete training loop for a given number of epochs.

    For each epoch:

    1. Calls :meth:`train_one_epoch`
    2. Calls :meth:`validate` (if ``val_dataloader`` is available)
    3. Steps the epoch-based scheduler (if configured with
       ``scheduler_interval='epoch'``)
    4. Saves a checkpoint

    Args:
        num_epochs: Number of epochs to train.

    Returns:
        History dictionary with keys ``'train'`` and ``'val'``, each
        containing a list of per-epoch metric dictionaries.
    """
    history: dict[str, list] = {
        "train": [],
        "val": [],
    }

    for epoch in range(1, num_epochs + 1):
        # ---- Training ----------------------------------------------
        train_metrics: dict[str, Any] = self.train_one_epoch(epoch)
        history["train"].append(train_metrics)

        # ---- Validation --------------------------------------------
        if self.val_dataloader is not None:
            val_metrics: dict[str, Any] = self.validate(epoch, use_ema=False)
            history["val"].append(val_metrics)
        else:
            history["val"].append({"val_loss": 0.0})

        # ---- Epoch-based scheduler step ----------------------------
        if (
            self.scheduler is not None
            and self.scheduler_interval == "epoch"
        ):
            self.scheduler.step()

        # ---- Print Epoch Summary -----------------------------------
        parts: list[str] = [f"Epoch {epoch:3d}/{num_epochs:3d}"]
        # Train metrics
        for k, v in train_metrics.items():
            if isinstance(v, float):
                parts.append(f"{k}: {v:.4f}")
            else:
                parts.append(f"{k}: {v}")
        # Validation metrics
        val_dict: dict[str, Any] = history["val"][-1]
        for k, v in val_dict.items():
            if isinstance(v, float):
                parts.append(f"{k}: {v:.4f}")
            else:
                parts.append(f"{k}: {v}")
        print(" | ".join(parts))  # noqa: T201

        # ---- Save checkpoint ---------------------------------------
        self.save_checkpoint(
            path=str(self.output_dir / f"epoch_{epoch}.pt"),
            epoch=epoch,
            metrics={
                "train": train_metrics,
                "val": history["val"][-1],
            },
        )

    # ---- Final checkpoint ------------------------------------------
    self.save_checkpoint(
        path=str(self.output_dir / "final.pt"),
        epoch=num_epochs,
        metrics=history,
    )

    return history

load_checkpoint(path, load_optimizer=True, load_scheduler=True, load_ema=True)

Load a training checkpoint from disk.

Parameters:

Name	Type	Description	Default
path	str	File path of the checkpoint.	required
load_optimizer	bool	If True, restore the optimizer state dict. Default True.	True
load_scheduler	bool	If True, restore the scheduler state dict. Default True.	True
load_ema	bool	If True, restore the EMA parameter state. Default True.	True

Returns:

Type	Description
dict[str, Any]	The full checkpoint dictionary (contains at least 'epoch',
dict[str, Any]	'model_state_dict', and 'metrics'). May also contain
dict[str, Any]	'model_config' if the checkpoint was saved by a trainer
dict[str, Any]	that was initialised with a model_config.

Raises:

Type	Description
FileNotFoundError	If the checkpoint file does not exist.
RuntimeError	If the checkpoint is missing a required key.

Source code in src/corecv/engine/trainer.py

def load_checkpoint(
    self,
    path: str,
    load_optimizer: bool = True,
    load_scheduler: bool = True,
    load_ema: bool = True,
) -> dict[str, Any]:
    """Load a training checkpoint from disk.

    Args:
        path: File path of the checkpoint.
        load_optimizer: If ``True``, restore the optimizer state dict.
            Default ``True``.
        load_scheduler: If ``True``, restore the scheduler state dict.
            Default ``True``.
        load_ema: If ``True``, restore the EMA parameter state.
            Default ``True``.

    Returns:
        The full checkpoint dictionary (contains at least ``'epoch'``,
        ``'model_state_dict'``, and ``'metrics'``).  May also contain
        ``'model_config'`` if the checkpoint was saved by a trainer
        that was initialised with a ``model_config``.

    Raises:
        FileNotFoundError: If the checkpoint file does not exist.
        RuntimeError: If the checkpoint is missing a required key.
    """
    checkpoint: dict[str, Any] = torch.load(path, map_location=self.device)

    if "model_state_dict" not in checkpoint:
        msg = f"Checkpoint at {path} is missing 'model_state_dict'"
        raise RuntimeError(msg)

    # ---- Load model state ------------------------------------------
    self.model.load_state_dict(checkpoint["model_state_dict"])

    # ---- Load optimizer state --------------------------------------
    if load_optimizer and "optimizer_state_dict" in checkpoint:
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

    # ---- Load scheduler state --------------------------------------
    if load_scheduler and self.scheduler is not None:
        sched_state = checkpoint.get("scheduler_state_dict")
        if sched_state is not None:
            self.scheduler.load_state_dict(sched_state)

    # ---- Load EMA state --------------------------------------------
    if load_ema and "ema_state_dict" in checkpoint:
        self._ema_params = checkpoint["ema_state_dict"]

    # ---- Load scaler state -----------------------------------------
    if "scaler_state_dict" in checkpoint:
        self.scaler.load_state_dict(checkpoint["scaler_state_dict"])

    epoch: int = checkpoint.get("epoch", 0)
    logger.info("Checkpoint loaded from %s (epoch %d)", path, epoch)

    return checkpoint

save_checkpoint(path, epoch, metrics)

Save a training checkpoint to disk.

The checkpoint dictionary contains the following keys:

• epoch — Current epoch number.
• model_state_dict — Model parameters.
• optimizer_state_dict — Optimizer state.
• scheduler_state_dict — Scheduler state (None if not set).
• ema_state_dict — EMA shadow parameters.
• scaler_state_dict — AMP gradient scaler state.
• metrics — User-supplied metrics dictionary.
• model_config — Model configuration dictionary (None if not provided at initialisation).

Parameters:

Name	Type	Description	Default
path	str	File path for the checkpoint.	required
epoch	int	Current epoch number.	required
metrics	dict[str, Any]	Dictionary of metrics to store in the checkpoint (e.g. loss, accuracy, etc.).	required

Source code in src/corecv/engine/trainer.py

def save_checkpoint(
    self,
    path: str,
    epoch: int,
    metrics: dict[str, Any],
) -> None:
    """Save a training checkpoint to disk.

    The checkpoint dictionary contains the following keys:

    - ``epoch`` — Current epoch number.
    - ``model_state_dict`` — Model parameters.
    - ``optimizer_state_dict`` — Optimizer state.
    - ``scheduler_state_dict`` — Scheduler state (``None`` if not set).
    - ``ema_state_dict`` — EMA shadow parameters.
    - ``scaler_state_dict`` — AMP gradient scaler state.
    - ``metrics`` — User-supplied metrics dictionary.
    - ``model_config`` — Model configuration dictionary (``None`` if
      not provided at initialisation).

    Args:
        path: File path for the checkpoint.
        epoch: Current epoch number.
        metrics: Dictionary of metrics to store in the checkpoint
            (e.g. loss, accuracy, etc.).
    """
    checkpoint: dict[str, Any] = {
        "epoch": epoch,
        "model_state_dict": self.model.state_dict(),
        "optimizer_state_dict": self.optimizer.state_dict(),
        "scheduler_state_dict": (
            self.scheduler.state_dict()
            if self.scheduler is not None
            else None
        ),
        "ema_state_dict": self._ema_params,
        "scaler_state_dict": self.scaler.state_dict(),
        "metrics": metrics,
        "model_config": self.model_config,
    }
    torch.save(checkpoint, path)
    logger.info("Checkpoint saved to %s", path)

train_one_epoch(epoch)

Run one training epoch.

Iterates over train_dataloader, computes forward / backward, accumulates gradients, applies gradient clipping, updates weights, and optionally updates EMA and scheduler (if scheduler_interval == 'step').

Parameters:

Name	Type	Description	Default
epoch	int	Current epoch number (1-indexed, used for logging and EMA start condition).	required

Returns:

Type	Description
dict[str, Any]	Dictionary of training metrics for the epoch, including at
dict[str, Any]	minimum 'loss' and 'lr', plus any metrics from
dict[str, Any]	train_metrics.compute().

Raises:

Type	Description
ValueError	If train_dataloader is empty.

Source code in src/corecv/engine/trainer.py

def train_one_epoch(self, epoch: int) -> dict[str, Any]:
    """Run one training epoch.

    Iterates over ``train_dataloader``, computes forward / backward,
    accumulates gradients, applies gradient clipping, updates weights,
    and optionally updates EMA and scheduler (if
    ``scheduler_interval == 'step'``).

    Args:
        epoch: Current epoch number (1-indexed, used for logging and
            EMA start condition).

    Returns:
        Dictionary of training metrics for the epoch, including at
        minimum ``'loss'`` and ``'lr'``, plus any metrics from
        ``train_metrics.compute()``.

    Raises:
        ValueError: If ``train_dataloader`` is empty.
    """
    self.model.train()

    num_batches: int = len(self.train_dataloader)
    if num_batches == 0:
        msg = "train_dataloader is empty"
        raise ValueError(msg)

    total_loss: float = 0.0
    log_loss: float = 0.0
    log_count: int = 0
    current_lr: float = self._get_current_lr()
    optimizer_steps: int = 0

    if self.train_metrics is not None:
        self.train_metrics.reset()

    self.optimizer.zero_grad()

    for batch_idx, batch in enumerate(self.train_dataloader):
        # --- Unpack batch -------------------------------------------
        inputs, targets = self._unpack_batch(batch)
        inputs = inputs.to(self.device)
        targets = self._targets_to_device(targets)

        # --- Forward (AMP autocast) ---------------------------------
        with autocast(
            device_type=self.device.type,
            dtype=self.amp_dtype,
            enabled=self.use_amp,
        ):
            outputs = self.model(inputs)
            loss = self.loss_fn(outputs, targets)

        # --- Backward (scale for accumulation) ----------------------
        scaled_loss = loss / self.gradient_accumulation_steps
        self.scaler.scale(scaled_loss).backward()

        # --- Accumulators for logging / epoch stats -----------------
        loss_item: float = loss.item()
        total_loss += loss_item
        log_loss += loss_item
        log_count += 1

        # --- Optimizer step at accumulation boundary ----------------
        if (batch_idx + 1) % self.gradient_accumulation_steps == 0:
            self._optimizer_step(epoch)
            optimizer_steps += 1
            current_lr = self._get_current_lr()

            # Logging
            if optimizer_steps % self.log_interval == 0:
                avg_loss: float = log_loss / max(log_count, 1)
                logger.info(
                    "Epoch [%d] Step [%d/%d]  Loss: %.4f  LR: %.6f",
                    epoch,
                    batch_idx + 1,
                    num_batches,
                    avg_loss,
                    current_lr,
                )
                log_loss = 0.0
                log_count = 0

        # --- Training metrics ---------------------------------------
        if self.train_metrics is not None:
            self.train_metrics.update(outputs.detach(), targets)

    # Handle remaining gradients when accumulation boundary not reached
    last_batch: int = batch_idx + 1  # noqa: F841  # safe after non-empty loop
    if last_batch % self.gradient_accumulation_steps != 0:
        self._optimizer_step(epoch)
        current_lr = self._get_current_lr()

    # ---- Epoch-level metrics ---------------------------------------
    avg_epoch_loss: float = total_loss / num_batches
    metrics: dict[str, Any] = {
        "loss": avg_epoch_loss,
        "lr": current_lr,
    }

    if self.train_metrics is not None:
        train_results: dict[str, Any] = self.train_metrics.compute()
        for k, v in train_results.items():
            metrics[k] = v.item() if isinstance(v, torch.Tensor) else v

    return metrics

validate(epoch, use_ema=False)

Run validation on val_dataloader.

Parameters:

Name	Type	Description	Default
epoch	int	Current epoch number (used for logging).	required
use_ema	bool	If True, temporarily apply EMA weights for the validation run. Default False.	False

Returns:

Type	Description
dict[str, Any]	Dictionary of validation metrics, including at least
dict[str, Any]	'val_loss', plus any metrics from val_metrics.compute().

Raises:

Type	Description
RuntimeError	If val_dataloader is None.

Source code in src/corecv/engine/trainer.py

def validate(
    self,
    epoch: int,
    use_ema: bool = False,
) -> dict[str, Any]:
    """Run validation on ``val_dataloader``.

    Args:
        epoch: Current epoch number (used for logging).
        use_ema: If ``True``, temporarily apply EMA weights for the
            validation run.  Default ``False``.

    Returns:
        Dictionary of validation metrics, including at least
        ``'val_loss'``, plus any metrics from ``val_metrics.compute()``.

    Raises:
        RuntimeError: If ``val_dataloader`` is ``None``.
    """
    if self.val_dataloader is None:
        msg = "val_dataloader is not configured"
        raise RuntimeError(msg)

    # Optionally apply EMA weights
    ema_ctx: EMAContext | None = None
    if use_ema and self.ema_decay is not None:
        ema_ctx = EMAContext(self)
        ema_ctx.__enter__()

    self.model.eval()

    total_loss: float = 0.0
    num_batches: int = len(self.val_dataloader)

    if num_batches == 0:
        if ema_ctx is not None:
            ema_ctx.__exit__()
        return {"val_loss": 0.0}

    if self.val_metrics is not None:
        self.val_metrics.reset()

    with torch.no_grad():
        for batch in self.val_dataloader:
            inputs, targets = self._unpack_batch(batch)
            inputs = inputs.to(self.device)
            targets = self._targets_to_device(targets)

            with autocast(
                device_type=self.device.type,
                dtype=self.amp_dtype,
                enabled=self.use_amp,
            ):
                outputs = self.model(inputs)
                loss = self.loss_fn(outputs, targets)

            total_loss += loss.item()

            if self.val_metrics is not None:
                self.val_metrics.update(outputs, targets)

    # Restore original weights if EMA was applied
    if ema_ctx is not None:
        ema_ctx.__exit__()

    avg_val_loss: float = total_loss / num_batches
    metrics: dict[str, Any] = {
        "val_loss": avg_val_loss,
    }

    if self.val_metrics is not None:
        val_results: dict[str, Any] = self.val_metrics.compute()
        for k, v in val_results.items():
            metrics[k] = v.item() if isinstance(v, torch.Tensor) else v

    logger.info("Epoch [%d]  Validation Loss: %.4f", epoch, avg_val_loss)
    return metrics

---

CorePredictor

corecv.engine.predictor.CorePredictor(model, task='detection', input_size=(640, 640), mean=_MEAN_DEFAULT, std=_STD_DEFAULT, conf_threshold=0.25, iou_threshold=0.45, topk=5, half_precision=False, compile_model=False, batch_size=8, num_classes=None)

Accelerated inference engine for CoreCV models.

Wraps a CoreCV model and provides a clean predict() / predict_batch() API with optimised preprocessing, GPU-native post-processing, and configurable acceleration features.

Parameters:

Name	Type	Description	Default
model	Module	A CoreCV model (nn.Module) -- e.g. CoreObjectDetector, a segmentation model, or a classification model.	required
task	Literal['classification', 'segmentation', 'detection']	The task type. One of "classification", "segmentation", or "detection".	'detection'
input_size	tuple[int, int]	Target (height, width) for preprocessing. Default (640, 640).	(640, 640)
mean	tuple[float, float, float]	Per-channel normalisation mean. Default ImageNet mean.	_MEAN_DEFAULT
std	tuple[float, float, float]	Per-channel normalisation standard deviation. Default ImageNet std.	_STD_DEFAULT
conf_threshold	float	Minimum confidence score for detection predictions. Default 0.25.	0.25
iou_threshold	float	IoU threshold for detection NMS. Default 0.45.	0.45
topk	int	Number of top predictions for classification. Default 5.	5
half_precision	bool	Enable FP16 inference via torch.autocast. Default False.	False
compile_model	bool	Enable torch.compile on the model (requires PyTorch >= 2.0). Default False.	False
batch_size	int	Maximum batch size for list/folder inference. Default 8.	8
num_classes	int | None	Number of classes for detection heads. If None, inferred from the model (if possible).	None

Raises:

Type	Description
ValueError	If task is not a supported value.

Example

predictor = CorePredictor( ... model=my_model, ... task="detection", ... input_size=(640, 640), ... half_precision=True, ... ) results = predictor.predict("photo.jpg") results[0].detection.boxes.shape[1] 4

Initialise the CorePredictor with model and inference configuration.

Source code in src/corecv/engine/predictor.py

def __init__(  # noqa: PLR0913
    self,
    model: nn.Module,
    task: Literal["classification", "segmentation", "detection"] = "detection",
    input_size: tuple[int, int] = (640, 640),
    mean: tuple[float, float, float] = _MEAN_DEFAULT,
    std: tuple[float, float, float] = _STD_DEFAULT,
    conf_threshold: float = 0.25,
    iou_threshold: float = 0.45,
    topk: int = 5,
    half_precision: bool = False,
    compile_model: bool = False,
    batch_size: int = 8,
    num_classes: int | None = None,
) -> None:
    """Initialise the CorePredictor with model and inference configuration."""
    if task not in _SUPPORTED_TASKS:
        msg = (
            f"Unsupported task {task!r}. "
            f"Must be one of {sorted(_SUPPORTED_TASKS)}."
        )
        raise ValueError(msg)

    # ---- Task configuration ---------------------------------------
    self.task: str = task
    self.input_size: tuple[int, int] = (int(input_size[0]), int(input_size[1]))
    self.conf_threshold: float = conf_threshold
    self.iou_threshold: float = iou_threshold
    self.topk: int = topk
    self.batch_size: int = batch_size
    self.half_precision: bool = half_precision

    # ---- Normalisation tensors ------------------------------------
    self._mean: Tensor = torch.tensor(mean, dtype=torch.float32).view(
        1, _CHANNEL_COUNT_RGB, 1, 1,
    )
    self._std: Tensor = torch.tensor(std, dtype=torch.float32).view(
        1, _CHANNEL_COUNT_RGB, 1, 1,
    )

    # ---- Model setup ----------------------------------------------
    self.model: nn.Module = model.eval()

    # Move model to device
    self.device: torch.device = (
        next(model.parameters()).device
        if list(model.parameters())
        else torch.device("cpu")
    )

    # Optionally compile the model
    if compile_model:
        self.model = torch.compile(self.model)  # type: ignore[assignment]
        logger.info("Model compiled with torch.compile().")

    # ---- Class labels (optional) ----------------------------------
    self._class_labels: list[str] | None = None

    # ---- Detect num_classes from model if not provided ------------
    if num_classes is None:
        num_classes = self._infer_num_classes()
    self.num_classes: int | None = num_classes

predict(source)

Run inference on one or more images.

Accepts a wide variety of input types:

• str or Path: interpreted as a single image file path.
• np.ndarray: a single HWC or CHW image (uint8 or float32).
• Tensor: a single CHW image tensor.
• list: a list of any of the above; also supports a list of Tensor objects for pre-batched input.

When a directory path is given, all images with known extensions inside it are loaded and processed in batches.

Parameters:

Name	Type	Description	Default
source	str | Path | ndarray | Tensor | list[str | Path | ndarray | Tensor]	Image source(s) as described above.	required

Returns:

Type	Description
list[Prediction]	A list of :class:Prediction objects, one per input image.

Raises:

Type	Description
FileNotFoundError	If a file path does not exist.
TypeError	If the input type is unsupported.

Source code in src/corecv/engine/predictor.py

def predict(
    self,
    source: str | Path | np.ndarray | Tensor | list[str | Path | np.ndarray | Tensor],
) -> list[Prediction]:
    """Run inference on one or more images.

    Accepts a wide variety of input types:

    - ``str`` or ``Path``: interpreted as a single image file path.
    - ``np.ndarray``: a single HWC or CHW image (uint8 or float32).
    - ``Tensor``: a single CHW image tensor.
    - ``list``: a list of any of the above; also supports a list of
      ``Tensor`` objects for pre-batched input.

    When a directory path is given, all images with known extensions
    inside it are loaded and processed in batches.

    Args:
        source: Image source(s) as described above.

    Returns:
        A list of :class:`Prediction` objects, one per input image.

    Raises:
        FileNotFoundError: If a file path does not exist.
        TypeError: If the input type is unsupported.
    """
    image_items: list[tuple[np.ndarray, str | None]] = self._resolve_source(
        source,
    )

    all_predictions: list[Prediction] = []
    for batch_start in range(0, len(image_items), self.batch_size):
        batch_items = image_items[batch_start : batch_start + self.batch_size]
        batch_predictions = self._predict_batch_items(batch_items)
        all_predictions.extend(batch_predictions)

    return all_predictions

predict_batch(images)

Run inference on a list of pre-loaded tensors.

Each tensor should be in (C, H, W) format with float values in [0, 1] or [0, 255]. Tensors are automatically normalised and batched.

Parameters:

Name	Type	Description	Default
images	list[Tensor]	List of image tensors.	required

Returns:

Type	Description
list[Prediction]	A list of :class:Prediction objects, one per tensor.

Source code in src/corecv/engine/predictor.py

def predict_batch(self, images: list[Tensor]) -> list[Prediction]:
    """Run inference on a list of pre-loaded tensors.

    Each tensor should be in ``(C, H, W)`` format with float values
    in ``[0, 1]`` or ``[0, 255]``.  Tensors are automatically
    normalised and batched.

    Args:
        images: List of image tensors.

    Returns:
        A list of :class:`Prediction` objects, one per tensor.
    """
    items: list[tuple[Tensor, None]] = [(img, None) for img in images]
    return self._predict_tensor_items(items)

register_normalization(mean, std)

Register normalisation constants as module buffers.

The tensors are stored on CPU and broadcast to the input device during preprocessing to avoid repeated allocations.

Parameters:

Name	Type	Description	Default
mean	tuple[float, float, float]	Per-channel mean (RGB).	required
std	tuple[float, float, float]	Per-channel standard deviation (RGB).	required

Source code in src/corecv/engine/predictor.py

def register_normalization(
    self,
    mean: tuple[float, float, float],
    std: tuple[float, float, float],
) -> None:
    """Register normalisation constants as module buffers.

    The tensors are stored on CPU and broadcast to the input device
    during preprocessing to avoid repeated allocations.

    Args:
        mean: Per-channel mean (RGB).
        std: Per-channel standard deviation (RGB).
    """
    self._mean = torch.tensor(mean, dtype=torch.float32).view(
        1, _CHANNEL_COUNT_RGB, 1, 1,
    )
    self._std = torch.tensor(std, dtype=torch.float32).view(
        1, _CHANNEL_COUNT_RGB, 1, 1,
    )

set_class_labels(labels)

Set human-readable class labels for predictions.

Parameters:

Name	Type	Description	Default
labels	list[str]	List of label strings, one per class index.	required

Source code in src/corecv/engine/predictor.py

def set_class_labels(self, labels: list[str]) -> None:
    """Set human-readable class labels for predictions.

    Args:
        labels: List of label strings, one per class index.
    """
    self._class_labels = list(labels)