YOLOHelper Module

Provides helper functions for working with YOLO models.

HMB.YOLOHelper.ExtractYOLOModelSize(modelName)[source]

Extract model size from YOLO model name. If the model name contains “n”, “s”, “m”, “l”, or “x”, it maps to “Nano”, “Small”, “Medium”, “Large”, or “XLarge” respectively. If no match is found, the original model name is returned.

Parameters:

modelName (str) – The YOLO model name to extract size from.

Returns:

Human-readable model size or original model name if no match.

Return type:

str

HMB.YOLOHelper.TrainMultipleYoloClassifiers(datasetPath, baseDir, runsDir, targetModels=None, epochs=250, batchSize=128, inputShape=(512, 512), trialNum=1, exportOnnx=True, onnxOpset=11, deviceEnvVars=None, seed=None, overwriteExisting=False, enablePlots=True, enableSave=True)[source]

High-level function to train multiple YOLO classification models on a given dataset. It handles environment setup, model training, validation, and optional ONNX export. Also, it saves top-1 and top-5 metrics to text files for each model. Moreover, it saves the trained model in Keras format. Finally, it manages exceptions to ensure robustness across multiple models.

Parameters:

  • datasetPath (str) – Path to the dataset folder containing train/val splits.

  • baseDir (str) – Directory to save experiment outputs.

  • runsDir (str) – Subdirectory under baseDir to store run outputs.

  • targetModels (List[str], Optional) – List of YOLO model keywords to train. Defaults to None, which uses a predefined set.

  • epochs (int, Optional) – Number of training epochs. Defaults to 250.

  • batchSize (int, Optional) – Batch size for training. Defaults to 128.

  • inputShape (Tuple[int, int], Optional) – Input image shape (height, width). Defaults to (512, 512).

  • trialNum (int, Optional) – Trial number for naming runs. Defaults to 1.

  • exportOnnx (bool, Optional) – Whether to export trained models to ONNX format. Defaults to True.

  • onnxOpset (int, Optional) – ONNX opset version for export. Defaults to 11.

  • deviceEnvVars (dict, Optional) – Environment variables for device configuration. Defaults to None.

  • seed (int, Optional) – Random seed for reproducibility. Defaults to None.

  • overwriteExisting (bool, Optional) – Whether to overwrite existing runs. Defaults to False.

  • enablePlots (bool, Optional) – Whether to enable training plots. Defaults to True.

  • enableSave (bool, Optional) – Whether to save the best model. Defaults to True.

HMB.YOLOHelper.EvaluateAndSaveYoloClassifications(baseDir, datasetPath, runsDir, extensions=None, inputShape=(512, 512), trialNum=1, categories=None, targetModels=None)[source]

Evaluate trained YOLO classification models on dataset splits and save CSV results. It computes confusion matrices and performance metrics for each model and category. The results are saved in CSV format under the specified results subfolder. Also, it returns a summary mapping for programmatic access.

Parameters.

baseDir (str): Root experiment directory containing runs and dataset. datasetPath (str): Path to the dataset containing category subfolders. runsDir (str): Subdirectory under baseDir to store evaluation results. extensions (List[str] | None): Allowed file extensions for images. inputShape (Tuple[int,int]): Input image size as (height, width). trialNum (int): Trial number used in run naming. categories (List[str] | None): Categories to evaluate, defaults to [“val”,”test”,”train”]. targetModels (List[str] | None): Model list to evaluate, defaults to common yolo11 variants.

Returns.

dict: Mapping of model->{category->csvPath, metrics} for quick programmatic access.

HMB.YOLOHelper.EvaluatePredictPlotSubset(datasetDir, model, subset='test', prefix='', storageDir=None, heavy=True, computeECE=True, exportFailureCases=True, eps=1e-10, saveArtifacts=True, maxSamples=None, preprocessFn=None, dpi=720)[source]

Evaluate a trained Ultralytics YOLO classification model on a specified dataset subset (train/val/test/all), collect predictions, compute confusion matrix and performance metrics, and optionally save predictions to a CSV file. It also generates and saves confusion matrix, ROC AUC, and PRC plots.

Parameters:

  • datasetDir (str) – Path to the dataset directory containing train/val/test splits.

  • model (Callable) – A callable that takes an image (numpy.ndarray) and returns a 1D array of class probabilities.

  • subset (str) – Dataset subset to evaluate (“train”, “val”, “test”, or “all”). Defaults to “test”.

  • prefix (str) – Prefix for saved figure filenames. Defaults to “”.

  • storageDir (str | None) – Directory to save predictions CSV and figures. If None, uses current directory. Defaults to None.

  • heavy (bool) – Whether to compute heavy metrics and plot ROC/PRC curves. Defaults to True.

  • computeECE (bool) – Whether to compute Expected Calibration Error (ECE). Defaults to True.

  • exportFailureCases (bool) – Whether to export misclassified samples to CSV. Defaults to True.

  • eps (float) – Small epsilon value for numerical stability in metric calculations. Defaults to 1e-10.

  • saveArtifacts (bool) – Whether to save figures and artifacts. Defaults to True.

  • maxSamples (int | None) – Maximum number of samples to evaluate. If None, evaluates all samples. Defaults to None.

  • preprocessFn (Callable | None) – Optional preprocessing function to apply to each image before prediction. Defaults to None.

  • dpi (int) – DPI setting for saved figures. Defaults to 720.

Returns:

A tuple containing:

  • str|None: Path to the saved predictions CSV file (or None when not saved).

  • dict: Computed weighted performance metrics.

  • List[int]: List of predicted class indices.

  • List[int]: List of ground truth class indices.

  • List[List[float]]: List of predicted class probabilities for each sample.

  • List[Optional[float]]: List of predicted confidences for each sample.

  • List[Dict[str, Any]]: List of prediction records for each sample.

  • List[str]: List of class names.

  • numpy.ndarray|None: Confusion matrix as a 2D numpy array, or None if not computable.

Return type:

tuple

HMB.YOLOHelper.MeasureLatencyWithUltralytics(modelPath, runs=20, warmup=5, inputShape=(224, 224), exampleInput=None)[source]

Measure average inference latency (in milliseconds) for an Ultralytics YOLO model or other supported model file formats using a single synthetic or user-provided example input.

Supported model formats/paths:

  • Ultralytics-wrapped .pt weights or hub identifier: loaded via ultralytics.YOLO.

  • TorchScript file (.pt saved via torch.jit.save / scripted/traced): loaded via torch.jit.load.

  • ONNX file (.onnx): executed via onnxruntime.InferenceSession (if onnxruntime installed).

The helper will attempt to load the model using the best available backend for the provided path and run warmup + timed predictions, returning average latency in ms.

Parameters:

  • modelPath (str) – Path to the model file or Ultralytics model identifier.

  • runs (int) – Number of timed runs to average. Defaults to 20.

  • warmup (int) – Number of warmup runs to perform before timing. Defaults to 5.

  • inputShape (Tuple[int,int]) – (H, W) used to synthesize an RGB input when exampleInput is not provided.

  • exampleInput (numpy.ndarray | torch.Tensor | None) – Optional user-provided input. If given, this will be used for warmup and timed runs. Expected shape is HxWx3 or batch-like formats.

Returns:

Average latency in milliseconds on success, or None on failure.

Return type:

float | None

HMB.YOLOHelper.ExportYOLO2TorchScript(weightsPath, outPath, imgsz=224, device='cpu')[source]

Export a classification model to a TorchScript file.

This helper attempts to load a model via Ultralytics’ YOLO wrapper when available. It extracts the underlying core model, moves it to the requested device, and attempts to produce a TorchScript artifact by tracing with a dummy input first and falling back to scripting when tracing is not possible.

Parameters:

  • weightsPath (str) – Path to weights or a model identifier understood by Ultralytics.

  • outPath (str) – Desired output path for the TorchScript file (including filename).

  • imgsz (int) – Spatial size used for the dummy trace input. Defaults to 224.

  • device (str) – Torch device string to use for the conversion. Defaults to “cpu”.

Returns:

Returns the output path on success, otherwise None.

Return type:

str | None

HMB.YOLOHelper.ExportYOLO2ONNX(weightsPath, outPath, imgsz=224, opset=12)[source]

Export model to ONNX using Ultralytics’ export helper when available.

The function will attempt to call model.export(format=”onnx”) on a YOLO wrapper. If Ultralytics returns a path or list of paths, the first candidate is copied to the requested output location for convenience.

Parameters:

  • weightsPath (str) – Path or identifier understood by Ultralytics.

  • outPath (str) – Desired ONNX output path including filename.

  • imgsz (int) – Input spatial size to request from the exporter. Defaults to 224.

  • opset (int) – ONNX opset version to request. Defaults to 12.

Returns:

Returns the output path on success, otherwise None.

Return type:

str | None

HMB.YOLOHelper.ApplyYOLOPruning(weightsPath, outPath, sparsity=0.5)[source]

Apply global unstructured pruning to Conv2d and Linear weights and save state_dict.

This helper loads a classifier via the Ultralytics YOLO wrapper when available. It collects Conv2d and Linear weight tensors and applies global L1 unstructured pruning to the requested sparsity fraction. The function removes pruning reparametrizations to make weights dense and then saves the resulting state_dict to the requested output path.

Parameters:

  • weightsPath (str) – Path to weights or model identifier understood by Ultralytics.

  • outPath (str) – Path where the pruned state_dict will be saved.

  • sparsity (float) – Fraction of weights to remove in [0,1). Defaults to 0.5.

Returns:

Returns the output path on success, otherwise None.

Return type:

str | None

HMB.YOLOHelper.TrainMultipleYoloDetectors(datasetPath, baseDir, runsDir, targetModels=None, epochs=250, batchSize=128, inputShape=(512, 512), trialNum=1, exportOnnx=True, onnxOpset=11, deviceEnvVars=None, seed=None, overwriteExisting=False, enablePlots=True, enableSave=True)[source]

Train multiple YOLO detection models on a specified dataset and save results.

Parameters:

  • datasetPath (str) – Path to the dataset directory containing data.yaml.

  • baseDir (str) – Base directory where runs and outputs are stored.

  • runsDir (str) – Directory under baseDir where YOLO runs are stored.

  • targetModels (List[str] | None) – List of YOLO model keywords to train. Defaults to [“yolo26n”, “yolo26s”, “yolo26m”, “yolo26l”, “yolo26x”] if None.

  • epochs (int) – Number of training epochs. Defaults to 250.

  • batchSize (int) – Batch size for training. Defaults to 128.

  • inputShape (Tuple[int, int]) – Input image shape (height, width). Defaults to (512, 512).

  • trialNum (int) – Trial number to distinguish different runs. Defaults to 1.

  • exportOnnx (bool) – Whether to export trained models to ONNX format. Defaults to True.

  • onnxOpset (int) – ONNX opset version to use for export. Defaults to 11.

  • deviceEnvVars (Dict[str, str] | None) – Optional dictionary of environment variables for device configuration. Defaults to {“CudaVisibleDevices”: “0”, “CudaLaunchBlocking”: “1”, “TorchUseCudaDsa”: “1”} if None.

  • seed (int | None) – Optional random seed for reproducibility. If None, no seeding is applied. Defaults to None.

  • overwriteExisting (bool) – Whether to overwrite existing runs with the same name. Defaults to False.

  • enablePlots (bool) – Whether to generate and save plots during training and validation. Defaults to True.

  • enableSave (bool) – Whether to save model weights and artifacts. Defaults to True.

Return type:

None

HMB.YOLOHelper.EvaluateAndSaveYoloDetections(baseDir, datasetPath, runsDir, categories=None, targetModels=None, trialNum=1, inputShape=(512, 512), loadPt=True)[source]

Evaluate trained YOLO detection models on specified dataset categories and save metrics.

Parameters:

  • baseDir (str) – Base directory where runs and outputs are stored.

  • datasetPath (str) – Path to the dataset directory containing data.yaml.

  • runsDir (str) – Directory under baseDir where YOLO runs are stored.

  • categories (List[str] | None) – List of dataset categories to evaluate (e.g., “train”, “val”, “test”). Defaults to [“val”, “test”, “train”].

  • targetModels (List[str] | None) – List of YOLO model keywords to evaluate. Defaults to [“yolo26n”, “yolo26s”, “yolo26m”, “yolo26l”, “yolo26x”].

  • trialNum (int) – Trial number to distinguish different runs. Defaults to 1.

  • inputShape (Tuple[int, int]) – Input image shape (height, width) for evaluation. Defaults to (512, 512).

  • loadPt (bool) – Whether to load .pt weights (True) or .onnx weights (False). Defaults to True.

Return type:

Dict[str, Any]

HMB.YOLOHelper.GenerateYoloDetectionExplainability(datasetPath, modelPath, outputDir, inputShape=(512, 512), maxSamples=10)[source]

Generate explainability artifacts for a YOLO detection model by running inference on a subset of images.

Parameters:

  • datasetPath (str) – Path to the dataset directory containing images.

  • modelPath (str) – Path to the trained YOLO detection model weights (.pt).

  • outputDir (str) – Directory where explainability artifacts (plots) will be saved.

  • inputShape (Tuple[int, int]) – Input image shape (height, width) for inference. Defaults to (512, 512).

  • maxSamples (int) – Maximum number of images to process for explainability. Defaults to 10.

Return type:

None