import os
import numpy as np
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from HMB.Utils import DumpJsonFile
from HMB.Initializations import IgnoreWarnings, SeedEverything
from HMB.PerformanceMetrics import CalculatePerformanceMetrics
[docs]
def 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,
):
r'''
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.
'''
from ultralytics import YOLO
# Ensure default model list is provided when None.
if (targetModels is None):
targetModels = ["yolo11n", "yolo11s", "yolo11m", "yolo11l", "yolo11x"]
# Set the default device environment variables if none are provided.
if (deviceEnvVars is None):
# Define the environment variables.
deviceEnvVars = {
"KmpDuplicateLibOk" : "TRUE",
"CudaVisibleDevices": "0",
"CudaLaunchBlocking": "1",
"TorchUseCudaDsa" : "1",
}
# Apply the environment variables to the current process.
for envKey, envVal in deviceEnvVars.items():
# Convert the CamelCase key to uppercase for the OS environment.
os.environ[envKey.upper()] = str(envVal)
# Optionally seed randomness for reproducibility.
if (seed is not None):
SeedEverything(seed=seed)
# Suppress warnings if the project's helper exists.
IgnoreWarnings()
# Iterate over requested models and train each one.
for modelKeyword in targetModels:
print(f"\nTraining Model: {modelKeyword} ...")
try:
expOutputDir = os.path.join(baseDir, runsDir, f"{modelKeyword}-cls-{trialNum}")
os.makedirs(expOutputDir, exist_ok=True)
# Instantiate the YOLO classification model from the Ultralytics hub.
model = YOLO(f"{modelKeyword}-cls.pt", task="classify")
# Train the model with provided settings.
model.train(
data=datasetPath, # Path to dataset splits.
batch=batchSize, # Batch size for training.
epochs=epochs, # Number of epochs.
imgsz=inputShape[0], # Input image size (height).
plots=True, # Enable training plots.
save=True, # Save the best model.
name=f"{modelKeyword}-cls-{trialNum}", # Naming convention for the run.
project=runsDir, # Project directory for outputs.
exist_ok=overwriteExisting, # Overwrite existing runs if specified.
)
# Retrieve class names from the trained model.
classes = model.names
# Validate the model on the splits folder.
metrics = model.val(
data=datasetPath, # Path to dataset splits.
plots=enablePlots, # Enable validation plots.
save=enableSave, # Save predictions.
name=f"{modelKeyword}-cls-{trialNum}", # Naming convention for the run.
project=runsDir, # Project directory for outputs.
exist_ok=overwriteExisting, # Overwrite existing runs if specified.
)
# Store top-1 and top-5 metrics for later use if needed as text files.
# Ensure the output directory exists.
os.makedirs(expOutputDir, exist_ok=True)
# Save top-1 metrics to a text file.
with open(os.path.join(expOutputDir, "Top1-Metrics.txt"), "w") as f:
f.write(str(metrics.top1))
# Save top-5 metrics to a text file.
with open(os.path.join(expOutputDir, "Top5-Metrics.txt"), "w") as f:
f.write(str(metrics.top5))
print("Top-1 Metrics:", metrics.top1)
print("Top-5 Metrics:", metrics.top5)
# Save the trained model in Keras format if supported.
# model.save(os.path.join(expOutputDir, "model.keras"))
try:
model.export(format="keras", name="model", project=expOutputDir)
except Exception as e:
print(f"Warning: Keras export failed: {e}")
# Optionally export to ONNX format; wrap in try/except to avoid failing the loop.
if (exportOnnx):
try:
model.export(
format="onnx",
opset=onnxOpset,
dynamic=False,
simplify=True,
name=f"{modelKeyword}-cls-{trialNum}",
project=runsDir,
exist_ok=True,
)
except Exception as e:
print(f"Warning: ONNX export failed for {modelKeyword}: {e}")
except Exception as e:
print(f"Error training or processing model {modelKeyword}: {e}")
# Continue with next model on error.
continue
[docs]
def EvaluateAndSaveYoloClassifications(
baseDir,
datasetPath,
runsDir,
extensions=None,
inputShape=(512, 512),
trialNum=1,
categories=None,
targetModels=None,
):
r'''
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.
'''
import pandas as pd
from tqdm import tqdm
from ultralytics import YOLO
# Default allowed image extensions when None provided.
if (extensions is None):
extensions = ["tiff", "tif", "jpeg", "jpg", "png", "bmp"]
# Default categories when None provided.
if (categories is None):
categories = ["val", "test", "train"]
# Default target models when None provided.
if (targetModels is None):
targetModels = ["yolo11n", "yolo11s", "yolo11m", "yolo11l", "yolo11x"]
# Prepare a container to return run summaries.
summary = {}
# Iterate over requested categories and models.
for cat in categories:
print(f"\nEvaluating Category: {cat} ...")
# Create category-level summary mapping.
for modelKeyword in targetModels:
print(f"\nWorking on Model: {modelKeyword} ...")
# Create a per-model dictionary for results.
modelSummary = {}
# Compose the expected weights path for this run.
weights = os.path.join(
baseDir, "runs", "classify",
f"{modelKeyword}-cls-{trialNum}",
"weights", "best.pt"
)
# Check weights existence and continue if missing.
if (not os.path.exists(weights)):
# Report missing weights and continue to next model.
print(f"Weights file not found: {weights}.")
modelSummary[cat] = {"csv": None, "error": "weights_missing"}
summary[modelKeyword] = summary.get(modelKeyword, {})
summary[modelKeyword].update(modelSummary)
continue
# Load the model for classification task.
try:
# Load model in classify mode.
model = YOLO(weights, task="classify", verbose=False)
except Exception as e:
print(f"Failed to load model {weights}: {e}.")
modelSummary[cat] = {"csv": None, "error": "load_failed", "exception": str(e)}
summary[modelKeyword] = summary.get(modelKeyword, {})
summary[modelKeyword].update(modelSummary)
continue
# Prepare history list for building a DataFrame.
history = []
# Iterate over classes reported by the model.
for classKey, className in tqdm(list(model.names.items()), desc=f"Classes {modelKeyword}"):
# Compute directory containing images for this class and category.
classDir = os.path.join(datasetPath, cat, className)
# Skip if class directory does not exist.
if (not os.path.isdir(classDir)):
# Warn and continue when expected class folder is missing.
print(f"Warning: expected folder missing: {classDir}.")
continue
# List files in the class directory.
try:
files = os.listdir(classDir)
except Exception as e:
# If listing fails, skip this class.
print(f"Failed to list files in {classDir}: {e}.")
continue
# Iterate over files found in the class folder.
for fileName in tqdm(files, desc=f"Processing {className}"):
# Get file extension for filtering.
fileExt = fileName.split(".")[-1].lower()
# Skip files with unsupported extensions.
if (fileExt not in extensions):
continue
# Build the full image path.
imgPath = os.path.join(classDir, fileName)
# Run the model on the image and collect probabilities.
try:
pred = model(imgPath, imgsz=inputShape[0], verbose=False)
# Extract probabilities as a CPU numpy array when available.
toValues = None
try:
toValues = pred[0].probs.data.cpu().numpy()
except Exception:
# Fallback: try to get probs from different attribute name.
try:
toValues = pred[0].prob.data.cpu().numpy()
except Exception:
# If probabilities are not available, skip storing per-class probabilities.
toValues = None
# Determine predicted top1 index in a robust way.
try:
# predictedIndex = int(pred[0].probs.top1)
# predictedProb = float(pred[0].probs.values[0])
# Fix: Use the correct attributes.
predictedIndex = int(pred[0].probs.top1)
predictedProb = float(pred[0].probs.data[predictedIndex]) # or float(toValues[predictedIndex])
except Exception:
try:
predictedIndex = int(np.argmax(toValues)) if (toValues is not None) else -1
except Exception:
predictedIndex = -1
try:
predictedProb = (
float(toValues[predictedIndex])
if (toValues is not None and predictedIndex >= 0)
else 0.0
)
except Exception:
predictedProb = 0.0
# Compose history record with image path, actual class and predicted info.
# record = {
# "Image" : imgPath,
# "Actual" : className,
# "Predicted" : model.names.get(predictedIndex, "unknown"),
# "Actual Index" : int(classKey),
# "Predicted Index": int(predictedIndex),
# "Predicted Prob" : float(predictedProb),
# }
# Compose the history record with image path, actual class, and predicted information.
record = {
"ImagePath" : imgPath,
"ActualClass" : className,
"PredictedClass": model.names.get(predictedIndex, "Unknown"),
"ActualIndex" : int(classKey),
"PredictedIndex": int(predictedIndex),
"PredictedProb" : float(predictedProb),
}
# Add per-class probabilities when available.
if (toValues is not None):
for i in range(len(toValues)):
record[f"{model.names.get(i, str(i))} Prob"] = float(toValues[i])
# Append record to history list.
history.append(record)
except Exception as e:
# On failure to predict, log and continue to next image.
print(f"Prediction failed for image {imgPath}: {e}.")
continue
# Build a DataFrame from collected history.
df = pd.DataFrame(history)
# Compose storage path for the CSV file.
csvFileName = rf"{modelKeyword} Classification-{cat.capitalize()}.csv"
storagePath = os.path.join(baseDir, runsDir, csvFileName)
# Ensure the results folder exists.
os.makedirs(os.path.dirname(storagePath), exist_ok=True)
# Save the DataFrame to CSV.
df.to_csv(storagePath, index=False)
# Reload the CSV to ensure consistent types for metrics.
df = pd.read_csv(storagePath)
# Compute reference and prediction arrays for metrics.
references = np.array(list(df["Actual Index"].values))
predictions = np.array(list(df["Predicted Index"].values))
# Print shapes and a small sample of references and predictions.
print(f"Shape of references: {references.shape}.")
print(f"Shape of predictions: {predictions.shape}.")
print(f"Sample references: {references[:5]}.")
print(f"Sample predictions: {predictions[:5]}.")
# Compute confusion matrix for the set.
try:
from sklearn.metrics import confusion_matrix
# cm = confusion_matrix(references, predictions)
# Fix: Explicitly define all possible class labels.
numClasses = len(model.names)
cm = confusion_matrix(references, predictions, labels=range(numClasses))
except Exception as e:
# If confusion matrix computation fails, store the error and continue.
print(f"Confusion matrix computation failed: {e}.")
modelSummary[cat] = {"csv": storagePath, "error": "cm_failed", "exception": str(e)}
summary[modelKeyword] = summary.get(modelKeyword, {})
summary[modelKeyword].update(modelSummary)
continue
# Print the confusion matrix.
print("Confusion Matrix:")
print(cm)
# Compute additional performance metrics using the project's helper.
try:
metrics = CalculatePerformanceMetrics(cm, addWeightedAverage=True)
except Exception as e:
# Fallback: store the exception if metrics calculation fails.
print(f"CalculatePerformanceMetrics failed: {e}.")
modelSummary[cat] = {"csv": storagePath, "cm": cm.tolist(), "metrics_error": str(e)}
summary[modelKeyword] = summary.get(modelKeyword, {})
summary[modelKeyword].update(modelSummary)
continue
# Print the derived metrics to stdout.
print(f"Performance Metrics for {modelKeyword} on {cat} set:")
for key, value in metrics.items():
print(f"{key}: {value}")
# Populate model summary with results and metrics.
modelSummary[cat] = {"csv": storagePath, "cm": cm.tolist(), "metrics": metrics}
summary[modelKeyword] = summary.get(modelKeyword, {})
summary[modelKeyword].update(modelSummary)
# Save the summary mapping as a JSON file for reference.
summaryPath = os.path.join(baseDir, runsDir, "Classification_Summary.json")
DumpJsonFile(summaryPath, summary)
# Return the summary mapping for programmatic inspection.
return summary
[docs]
def EvaluatePredictPlotSubset(
datasetDir,
model,
subset="test",
prefix="",
storageDir=None,
heavy=True,
computeECE=True,
exportFailureCases=True,
eps=1e-10,
saveArtifacts=True,
maxSamples=None,
preprocessFn=None,
dpi: int = 720,
) -> tuple:
r'''
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:
tuple: 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.
'''
from HMB.PyTorchTrainingPipeline import GenericImageryEvaluatePredictPlotSubset
# Define a YOLO-specific prediction adapter that wraps the model to return probabilities.
def YoloPredictFn(
image: np.ndarray,
imgPath: Optional[str] = None,
cls: Optional[int | str] = None,
) -> np.ndarray:
'''
Adapter to convert YOLO model output into a 1D probability array.
Parameters:
image (numpy.ndarray): Input image as HWC NumPy array (uint8 or float32).
imgPath (str | None): Optional image path (not used here).
cls (int | str | None): Optional class index or name (not used here).
Returns:
numpy.ndarray: 1D array of class probabilities.
'''
import torch
# Run YOLO prediction (expects BGR or RGB? YOLO handles internally)
results = model.predict(image, verbose=False)
if (len(results) == 0):
raise RuntimeError("YOLO returned empty prediction list.")
result = results[0]
if (not hasattr(result, "probs") or result.probs is None):
raise RuntimeError("YOLO result missing 'probs' attribute.")
probsAttr = result.probs.data # This should be a tensor or array
if (isinstance(probsAttr, torch.Tensor)):
probs = probsAttr.detach().cpu().numpy()
elif (isinstance(probsAttr, np.ndarray)):
probs = probsAttr
else:
raise TypeError(f"Unexpected probs type: {type(probsAttr)}")
if (probs.ndim != 1):
raise ValueError(f"Expected 1D probabilities, got shape {probs.shape}")
return probs.astype(np.float32)
# Delegate all evaluation logic to the generic function.
return GenericImageryEvaluatePredictPlotSubset(
datasetDir=datasetDir,
model=YoloPredictFn, # Pass the YOLO adapter as the generic model callable.
subset=subset,
prefix=prefix,
storageDir=storageDir,
heavy=heavy,
computeECE=computeECE,
exportFailureCases=exportFailureCases,
eps=eps,
saveArtifacts=saveArtifacts,
maxSamples=maxSamples,
preprocessFn=preprocessFn,
dpi=dpi,
)
[docs]
def MeasureLatencyWithUltralytics(
modelPath: str,
runs: int = 20,
warmup: int = 5,
inputShape: Tuple[int, int] = (224, 224),
exampleInput: Optional[np.ndarray] = None,
):
r'''
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:
float | None: Average latency in milliseconds on success, or None on failure.
'''
import torch, time
from ultralytics import YOLO
try:
# Determine extension lower-case for format guessing.
_, ext = os.path.splitext(str(modelPath))
ext = ext.lower()
# Helper to build a synthetic numpy image in HWC uint8 format.
def _CreateNumpyRandomImage(h, w, backendName="ultralytics"):
if (backendName == "onnx"):
# ONNX expects NCHW float32 input; create accordingly.
return (np.random.rand(1, 3, h, w) * 255).astype("float32")
else:
# Ultralytics and torchscript expect HWC uint8 image.
return (np.random.rand(h, w, 3) * 255).astype("uint8")
# Prepare input in three possible backends: ultralytics, torchscript, onnxruntime.
backendName = None
loadedModel = None
ortSess = None
# Try ultralytics if class available and path doesn't explicitly look like an ONNX file.
if (ext != ".onnx"):
try:
loadedModel = YOLO(modelPath, task="classify", verbose=False)
backendName = "ultralytics"
except Exception:
loadedModel = None
backendName = None
# If ultralytics not selected, try TorchScript for .pt or fallback when ultralytics failed.
if (backendName is None) and (ext in [".pt", ".pth"]):
try:
tsModel = torch.jit.load(modelPath, map_location="cpu")
loadedModel = tsModel
backendName = "torchscript"
except Exception:
loadedModel = None
backendName = None
# If .onnx file, try onnxruntime.
if ((backendName is None) and (ext == ".onnx")):
try:
import onnxruntime as ort
ortSess = ort.InferenceSession(modelPath, providers=["CPUExecutionProvider"])
backendName = "onnx"
except Exception:
ortSess = None
backendName = None
# If still no backend identified, attempt a last-resort ultralytics instantiation.
if (backendName is None):
try:
loadedModel = YOLO(modelPath, task="classify", verbose=False)
backendName = "ultralytics"
except Exception:
backendName = None
if (backendName is None):
# Unsupported model type or required runtime not available.
return None
# Prepare the input for warmup and timing.
if (exampleInput is not None):
exampleInp = exampleInput
else:
h, w = inputShape
exampleInp = _CreateNumpyRandomImage(h, w, backendName=backendName)
# Warmup runs to stabilize runtime performance.
for _ in range(max(1, warmup)):
try:
if (backendName == "ultralytics"):
# Ultralytics accepts numpy HWC or path; pass imgsz when we generated the image.
loadedModel.predict(
exampleInp,
imgsz=max(inputShape) if (exampleInput is None) else None,
verbose=False
)
elif (backendName == "torchscript"):
# Ensure tensor on CPU and in NCHW float format.
if (not isinstance(exampleInp, torch.Tensor)):
if (isinstance(exampleInp, np.ndarray) and exampleInp.ndim == 3):
tensorInp = torch.from_numpy(exampleInp).permute(2, 0, 1).unsqueeze(0).float()
elif (isinstance(exampleInp, np.ndarray) and exampleInp.ndim == 4):
tensorInp = torch.from_numpy(exampleInp).permute(0, 3, 1, 2).float() if (
exampleInp.shape[-1] == 3) else torch.from_numpy(exampleInp).float()
else:
try:
tensorInp = torch.tensor(exampleInp)
except Exception:
tensorInp = None
if (tensorInp is None):
continue
else:
tensorInp = exampleInp
# Run the torchscript model.
loadedModel(tensorInp)
elif (backendName == "onnx"):
# ONNX runtime expects numpy inputs; get input name and run.
try:
inName = ortSess.get_inputs()[0].name
ortSess.run(None, {inName: exampleInp})
except Exception:
pass
except Exception:
# Ignore per-iteration errors during warmup to keep the benchmark robust.
pass
# Timed runs to measure average latency.
startTime = time.perf_counter()
for _ in range(max(1, runs)):
try:
if (backendName == "ultralytics"):
loadedModel.predict(
exampleInp,
imgsz=max(inputShape) if (exampleInput is None) else None,
verbose=False
)
elif (backendName == "torchscript"):
import torch
if (not isinstance(exampleInp, torch.Tensor)):
if (isinstance(exampleInp, np.ndarray) and exampleInp.ndim == 3):
tensorInp = torch.from_numpy(exampleInp).permute(2, 0, 1).unsqueeze(0).float()
elif (isinstance(exampleInp, np.ndarray) and exampleInp.ndim == 4):
tensorInp = (
torch.from_numpy(exampleInp).permute(0, 3, 1, 2).float()
if (exampleInp.shape[-1] == 3)
else torch.from_numpy(exampleInp).float()
)
else:
try:
tensorInp = torch.tensor(exampleInp)
except Exception:
tensorInp = None
if (tensorInp is None):
continue
else:
tensorInp = exampleInp
loadedModel(tensorInp)
elif (backendName == "onnx"):
inName = ortSess.get_inputs()[0].name
ortSess.run(None, {inName: exampleInp})
except Exception:
# Ignore individual run errors to keep benchmark robust.
pass
duration = time.perf_counter() - startTime
avgMs = (duration / max(1, runs)) * 1000.0
return float(avgMs)
except Exception:
return None
[docs]
def ExportYOLO2TorchScript(
weightsPath: str,
outPath: str,
imgsz: int = 224,
device: str = "cpu"
) -> Optional[str]:
r'''
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:
str | None: Returns the output path on success, otherwise None.
'''
import torch
from ultralytics import YOLO
try:
# Instantiate the YOLO classifier wrapper.
try:
yoloModel = YOLO(weightsPath, task="classify", verbose=False)
except Exception:
return None
# Prepare dummy input on the requested device.
deviceT = torch.device(device)
dummy = torch.randn(1, 3, imgsz, imgsz).to(deviceT)
# Attempt to extract the core PyTorch module from the wrapper.
try:
coreModel = yoloModel.model
coreModel.to(deviceT)
coreModel.eval()
# Try tracing first, and fall back to scripting when tracing fails.
try:
traced = torch.jit.trace(coreModel, dummy, strict=False)
except Exception:
traced = torch.jit.script(coreModel)
traced.save(str(outPath))
return str(outPath)
except Exception:
return None
except Exception:
return None
[docs]
def ExportYOLO2ONNX(weightsPath: str, outPath: str, imgsz: int = 224, opset: int = 12) -> Optional[str]:
r'''
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:
str | None: Returns the output path on success, otherwise None.
'''
from ultralytics import YOLO
try:
# Instantiate the YOLO classifier wrapper.
try:
yoloModel = YOLO(weightsPath, task="classify", verbose=False)
except Exception:
return None
# Attempt export using Ultralytics' helper.
try:
exported = yoloModel.export(format="onnx", imgsz=imgsz, opset=opset)
# Normalize exported return type to a candidate path string.
if (isinstance(exported, (list, tuple))):
candidate = exported[0] if (len(exported) > 0) else None
else:
candidate = exported
if (not candidate):
return None
if (os.path.exists(str(candidate))):
import shutil
shutil.copy2(str(candidate), str(outPath))
return str(outPath)
return None
except Exception:
return None
except Exception:
return None
[docs]
def ApplyYOLOPruning(weightsPath: str, outPath: str, sparsity: float = 0.5) -> Optional[str]:
r'''
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:
str | None: Returns the output path on success, otherwise None.
'''
import torch
from ultralytics import YOLO
try:
# Instantiate the wrapper for classification.
try:
modelWrapper = YOLO(weightsPath, task="classify", verbose=False)
except Exception:
return None
# Extract the core PyTorch model from the wrapper.
try:
model = modelWrapper.model
except Exception:
return None
# Import pruning utilities from torch.
import torch.nn.utils.prune as prune
# Validate sparsity value and coerce to float.
try:
amount = float(sparsity)
if (amount < 0.0 or amount >= 1.0):
# Accept exact zero to indicate no pruning.
if (amount == 0.0):
amount = 0.0
else:
return None
except Exception:
return None
# Collect module weight parameters to prune.
paramsToPrune = []
for _, module in model.named_modules():
if (isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Linear)):
if (hasattr(module, "weight")):
paramsToPrune.append((module, "weight"))
# Abort when no prunable parameters are found.
if (len(paramsToPrune) == 0):
return None
# Apply global unstructured pruning with L1 criterion.
try:
prune.global_unstructured(paramsToPrune, pruning_method=prune.L1Unstructured, amount=amount)
except Exception:
return None
# Remove pruning reparametrizations to leave dense tensors.
for module, name in paramsToPrune:
try:
prune.remove(module, name)
except Exception:
# Ignore failures to remove reparametrization for robustness.
pass
# Ensure output directory exists and save state_dict.
try:
outP = Path(outPath)
outP.parent.mkdir(parents=True, exist_ok=True)
torch.save(model.state_dict(), str(outP))
return str(outP)
except Exception:
return None
except Exception:
return None
[docs]
def TrainMultipleYoloDetectors(
datasetPath: str,
baseDir: str,
runsDir: str,
targetModels: Optional[List[str]] = None,
epochs: int = 250,
batchSize: int = 128,
inputShape: Tuple[int, int] = (512, 512),
trialNum: int = 1,
exportOnnx: bool = True,
onnxOpset: int = 11,
deviceEnvVars: Optional[Dict[str, str]] = None,
seed: Optional[int] = None,
overwriteExisting: bool = False,
enablePlots: bool = True,
enableSave: bool = True,
) -> None:
r'''
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.
'''
from ultralytics import YOLO
# Ensure the default model list is provided when none is specified.
if (targetModels is None):
# Set the default target models to the YOLO26 variants.
targetModels = ["yolo26n", "yolo26s", "yolo26m", "yolo26l", "yolo26x"]
# Set default device environment variables if none are provided.
if (deviceEnvVars is None):
deviceEnvVars = {
"CudaVisibleDevices": "0",
"CudaLaunchBlocking": "1",
"TorchUseCudaDsa" : "1",
}
# Apply environment variables to the current process.
for envKey, envVal in deviceEnvVars.items():
os.environ[envKey] = str(envVal)
# Optionally seed randomness for reproducibility.
if (seed is not None):
SeedEverything(seed=seed)
# Suppress warnings if the project helper exists.
IgnoreWarnings()
# Resolve dataset path to data.yaml if it exists to prevent strict val() errors.
dataPath = datasetPath
if (os.path.isdir(datasetPath)):
# Find all .yaml files in the dataset directory.
yamlFiles = [f for f in os.listdir(datasetPath) if (f.endswith(".yaml"))]
if (len(yamlFiles) == 0):
print(f"Warning: No .yaml file found in dataset directory {datasetPath}.")
elif (len(yamlFiles) > 1):
print(
f"Warning: Multiple .yaml files found in dataset directory {datasetPath}. "
f"Using the first one: {yamlFiles[0]}."
)
dataPath = os.path.join(datasetPath, yamlFiles[0])
else:
dataPath = os.path.join(datasetPath, yamlFiles[0])
# Iterate over requested models and train each one.
for modelKeyword in targetModels:
print(f"\nTraining Model: {modelKeyword} ...")
try:
expOutputDir = os.path.join(baseDir, runsDir, f"{modelKeyword}-det-{trialNum}")
os.makedirs(expOutputDir, exist_ok=True)
# Instantiate the YOLO detection model from the Ultralytics hub.
detectionModel = YOLO(f"{modelKeyword}.pt", task="detect")
# Train the model with the provided settings.
detectionModel.train(
data=dataPath,
batch=batchSize,
epochs=epochs,
imgsz=inputShape[0],
plots=enablePlots,
save=enableSave,
name=f"{modelKeyword}-det-{trialNum}",
project=runsDir,
exist_ok=overwriteExisting,
)
# Validate the model on the dataset splits.
validationMetrics = detectionModel.val(
data=dataPath,
plots=enablePlots,
save=enableSave,
name=f"{modelKeyword}-det-{trialNum}",
project=runsDir,
exist_ok=overwriteExisting,
)
# Ensure the output directory exists for saving metrics.
os.makedirs(expOutputDir, exist_ok=True)
# Save mean average precision metrics to a text file.
with open(os.path.join(expOutputDir, "MapMetrics.txt"), "w") as fileObject:
fileObject.write(f"Map50: {validationMetrics.box.map50}\nMap50-95: {validationMetrics.box.map}")
# Print the derived metrics to standard output.
print(f"Map50 Metrics: {validationMetrics.box.map50}")
print(f"Map50-95 Metrics: {validationMetrics.box.map}")
# Optionally export to ONNX format with error handling.
if (exportOnnx):
try:
detectionModel.export(
format="onnx",
opset=onnxOpset,
dynamic=False,
simplify=True,
name=f"{modelKeyword}-det-{trialNum}",
project=runsDir,
exist_ok=True,
)
except Exception as exceptionObject:
print(f"Warning: ONNX export failed for {modelKeyword}: {exceptionObject}")
except Exception as exceptionObject:
print(f"Error training or processing model {modelKeyword}: {exceptionObject}")
continue
[docs]
def EvaluateAndSaveYoloDetections(
baseDir: str,
datasetPath: str,
runsDir: str,
categories: Optional[List[str]] = None,
targetModels: Optional[List[str]] = None,
trialNum: int = 1,
inputShape: Tuple[int, int] = (512, 512),
loadPt: bool = True,
) -> Dict[str, Any]:
r'''
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.
'''
from ultralytics import YOLO
# Set default categories when none are provided.
if (categories is None):
categories = ["val", "test", "train"]
# Set default target models when none are provided.
if (targetModels is None):
targetModels = ["yolo26n", "yolo26s", "yolo26m", "yolo26l", "yolo26x"]
# Prepare a container to return run summaries.
summaryDict = {}
# Iterate over requested categories and models.
for categoryName in categories:
print(f"\nEvaluating Category: {categoryName} ...")
for modelKeyword in targetModels:
print(f"\nWorking on Model: {modelKeyword} ...")
# Compose the expected weights path for this run.
# weightsPath = os.path.join(
# baseDir, "runs", "detect",
# f"{modelKeyword}-det-{trialNum}",
# "weights", "best.pt"
# )
# Find all weights folders that start with the expected prefix to handle potential variations in naming.
weightsDir = os.path.join(baseDir, "runs", "detect")
candidateWeightsPath = []
if (os.path.isdir(weightsDir)):
for entry in os.listdir(weightsDir):
if (entry.startswith(f"{modelKeyword}-det-{trialNum}")):
candidatePath = os.path.join(
weightsDir, entry, "weights",
"best.pt" if (loadPt) else "best.onnx"
)
if (os.path.exists(candidatePath)):
candidateWeights = candidatePath
candidateWeightsPath.append(candidateWeights)
if (len(candidateWeightsPath) == 0):
print(
f"Weights file not found for model {modelKeyword} in category {categoryName}. "
f"Expected pattern: {modelKeyword}-det-{trialNum}"
)
continue
elif (len(candidateWeightsPath) > 1):
print(
f"Multiple candidate weights files found for model {modelKeyword} in category {categoryName}. "
f"Using the first one: {candidateWeightsPath[0]}"
)
weightsPath = candidateWeightsPath[0]
else:
weightsPath = candidateWeightsPath[0]
# Check weights existence and continue if missing.
if (not os.path.exists(weightsPath)):
print(f"Weights file not found: {weightsPath}.")
continue
# Load the model for detection task.
try:
detectionModel = YOLO(weightsPath, task="detect", verbose=False)
except Exception as exceptionObject:
print(f"Failed to load model {weightsPath}: {exceptionObject}.")
continue
# Validate the model on the dataset to compute detection metrics.
try:
print(f"Validating model {modelKeyword} on category {categoryName} ...")
validationMetrics = detectionModel.val(
data=os.path.join(datasetPath, "data.yaml"),
split=categoryName,
plots=True,
save=True,
name=f"{modelKeyword}-det-{trialNum}-{categoryName}",
project=runsDir,
exist_ok=True,
imgsz=inputShape[0],
)
except Exception as exceptionObject:
print(f"Validation failed for {modelKeyword}: {exceptionObject}.")
continue
# Extract the mean average precision metric from the validation results.
mapMetric = float(validationMetrics.box.map)
# Extract the mean average precision at fifty percent IoU metric.
map50Metric = float(validationMetrics.box.map50)
# Calculate the mean precision across all classes to handle the array output.
precisionMetric = float(np.mean(validationMetrics.box.p))
# Calculate the mean recall across all classes to handle the array output.
recallMetric = float(np.mean(validationMetrics.box.r))
# Print the derived map metric to the standard output.
print(f"Map50-95: {mapMetric}")
# Print the derived map50 metric to the standard output.
print(f"Map50: {map50Metric}")
# Print the derived precision metric to the standard output.
print(f"Precision: {precisionMetric}")
# Print the derived recall metric to the standard output.
print(f"Recall: {recallMetric}")
# Populate the model summary dictionary with the extracted results and metrics.
summaryDict[f"{modelKeyword}_{categoryName}"] = {
"CategoryName" : categoryName,
"MapMetric" : mapMetric,
"Map50Metric" : map50Metric,
"PrecisionMetric": precisionMetric,
"RecallMetric" : recallMetric,
}
if (len(targetModels) > 1):
# Save the summary mapping as a JSON file for reference.
summaryPath = os.path.join(baseDir, "runs", "detect", "Detection_Summary.json")
DumpJsonFile(summaryPath, summaryDict)
else:
modelKeyword = targetModels[0]
summaryPath = os.path.join(baseDir, "runs", "detect", f"{modelKeyword}-det-{trialNum}", "Detection_Summary.json")
DumpJsonFile(summaryPath, summaryDict)
# Return the summary mapping for programmatic inspection.
return summaryDict
[docs]
def GenerateYoloDetectionExplainability(
datasetPath: str,
modelPath: str,
outputDir: str,
inputShape: Tuple[int, int] = (512, 512),
maxSamples: int = 10,
) -> None:
r'''
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.
'''
import cv2
from ultralytics import YOLO
from HMB.Initializations import IMAGE_SUFFIXES
# Instantiate the YOLO detection model from the provided weights.
detectionModel = YOLO(modelPath, task="detect")
# Ensure the output directory exists for saving explainability artifacts.
os.makedirs(outputDir, exist_ok=True)
# Initialize an empty list to collect image file paths.
imageFiles = []
# Iterate over the root directories, subdirectories, and files.
for rootDir, subDirs, files in os.walk(datasetPath):
# Iterate over each file in the current directory.
for fileName in files:
# Check if the file has a supported image extension.
if (fileName.lower().endswith(tuple(IMAGE_SUFFIXES))):
# Append the full image path to the list of files.
imageFiles.append(os.path.join(rootDir, fileName))
# Limit the number of samples to process for explainability.
imageFiles = imageFiles[:maxSamples]
# Iterate over the selected images to perform inference and generate explanations.
for imageIndex, imagePath in enumerate(imageFiles):
# Run the model on the image to obtain detection results.
detectionResults = detectionModel(imagePath, imgsz=inputShape[0], verbose=False)
# Compose the output path for the explainability plot.
plotPath = os.path.join(outputDir, f"ExplainabilitySample{imageIndex}.jpg")
# Extract the plotted image with detection overlays.
plottedImage = detectionResults[0].plot(conf=True, line_width=2, font_size=10)
plottedImage = cv2.cvtColor(plottedImage, cv2.COLOR_RGB2BGR) # Convert RGB to BGR.
# Save the plotted image to the specified output directory.
cv2.imwrite(plotPath, plottedImage)
# Iterate over each detection box in the results.
for detectionBox in detectionResults[0].boxes:
# Extract the class name from the model names dictionary.
className = detectionModel.names[int(detectionBox.cls)]
# Extract the confidence score as a float.
confidence = float(detectionBox.conf)
# Print the inference metrics for the current sample.
print(f"Detected: {className} with confidence {confidence:.4f}")
if __name__ == "__main__":
# Generate a random seed for this run.
rndNumber = np.random.randint(0, 10000)
# Compute splits and experiment paths used in the example.
# This should point to the folder containing "train", "val", "test" subfolders.
datasetPath = r"/path/to/dataset/splits"
expDir = r"/path/to/experiment"
# Call the high-level training function with defaults.
TrainMultipleYoloClassifiers(
baseDir=expDir,
datasetPath=datasetPath,
runsDir="runs/classify",
epochs=250,
batchSize=128,
inputShape=(512, 512),
trialNum=1,
targetModels=None,
exportOnnx=True,
onnxOpset=11,
seed=rndNumber,
)
# Call the evaluation and saving function with defaults.
summary = EvaluateAndSaveYoloClassifications(
baseDir=expDir,
datasetPath=datasetPath,
runsDir="results/classify",
extensions=None,
inputShape=(512, 512),
trialNum=1,
categories=None,
targetModels=None,
)
# Generate a random seed for this run.
rndNumber = np.random.randint(0, 10000)
# Compute splits and experiment paths used in the example.
datasetPath = r"/path/to/dataset/splits"
expDir = r"/path/to/experiment"
# Call the high-level training function for object detection.
TrainMultipleYoloDetectors(
baseDir=expDir,
datasetPath=datasetPath,
runsDir="runs/detect",
epochs=250,
batchSize=128,
inputShape=(512, 512),
trialNum=1,
targetModels=None,
exportOnnx=True,
onnxOpset=11,
seed=rndNumber,
)
# Call the evaluation and saving function for object detection.
detectionSummary = EvaluateAndSaveYoloDetections(
baseDir=expDir,
datasetPath=datasetPath,
runsDir="results/detect",
categories=None,
targetModels=None,
trialNum=1,
inputShape=(512, 512),
)
# Define the path to the best trained model weights.
bestModelPath = os.path.join(expDir, "runs", "detect", "yolo11n-det-1", "weights", "best.pt")
# Define the output directory for explainability artifacts.
explainabilityDir = os.path.join(expDir, "results", "explainability")
# Generate inference and explainability artifacts for the detection model.
GenerateYoloDetectionExplainability(
datasetPath=datasetPath,
modelPath=bestModelPath,
outputDir=explainabilityDir,
inputShape=(512, 512),
maxSamples=10,
)