1289 lines
56 KiB
Python
Executable File
1289 lines
56 KiB
Python
Executable File
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import random
|
||
from copy import copy
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.distributed as dist
|
||
import torch.nn as nn
|
||
|
||
from ultralytics.data import build_dataloader, build_yolo_dataset
|
||
from ultralytics.engine.trainer import BaseTrainer
|
||
from ultralytics.models import yolo
|
||
from ultralytics.models.yolo.detect.val import DetectionValidator
|
||
from ultralytics.nn.tasks import DetectionModel
|
||
from ultralytics.utils import DEFAULT_CFG, LOGGER, RANK
|
||
from ultralytics.utils.patches import override_configs
|
||
from ultralytics.utils.plotting import plot_images, plot_labels
|
||
from ultralytics.utils.torch_utils import intersect_dicts, torch_distributed_zero_first, unwrap_model
|
||
|
||
|
||
class GroundDetectionModel(DetectionModel):
|
||
"""YOLO detection model with ground loss for 2D detection tasks.
|
||
|
||
This model extends DetectionModel to use ground-specific loss functions
|
||
that handle difficulty weighting for ground 2D detection.
|
||
"""
|
||
|
||
def init_criterion(self):
|
||
"""Initialize the loss criterion for ground detection.
|
||
|
||
Returns:
|
||
E2EGroundLoss or v8DetectionLossGround depending on model configuration.
|
||
"""
|
||
from ultralytics.utils.loss import E2EGroundLoss, v8DetectionLossGround
|
||
|
||
return E2EGroundLoss(self) if getattr(self, "end2end", False) else v8DetectionLossGround(self)
|
||
|
||
|
||
class DetectionTrainer(BaseTrainer):
|
||
"""A class extending the BaseTrainer class for training based on a detection model.
|
||
|
||
This trainer specializes in object detection tasks, handling the specific requirements for training YOLO models for
|
||
object detection including dataset building, data loading, preprocessing, and model configuration.
|
||
|
||
Attributes:
|
||
model (DetectionModel): The YOLO detection model being trained.
|
||
data (dict): Dictionary containing dataset information including class names and number of classes.
|
||
loss_names (tuple): Names of the loss components used in training (box_loss, cls_loss, dfl_loss).
|
||
|
||
Methods:
|
||
build_dataset: Build YOLO dataset for training or validation.
|
||
get_dataloader: Construct and return dataloader for the specified mode.
|
||
preprocess_batch: Preprocess a batch of images by scaling and converting to float.
|
||
set_model_attributes: Set model attributes based on dataset information.
|
||
get_model: Return a YOLO detection model.
|
||
get_validator: Return a validator for model evaluation.
|
||
label_loss_items: Return a loss dictionary with labeled training loss items.
|
||
progress_string: Return a formatted string of training progress.
|
||
plot_training_samples: Plot training samples with their annotations.
|
||
plot_training_labels: Create a labeled training plot of the YOLO model.
|
||
auto_batch: Calculate optimal batch size based on model memory requirements.
|
||
|
||
Examples:
|
||
>>> from ultralytics.models.yolo.detect import DetectionTrainer
|
||
>>> args = dict(model="yolo26n.pt", data="coco8.yaml", epochs=3)
|
||
>>> trainer = DetectionTrainer(overrides=args)
|
||
>>> trainer.train()
|
||
"""
|
||
|
||
def __init__(self, cfg=DEFAULT_CFG, overrides: dict[str, Any] | None = None, _callbacks=None):
|
||
"""Initialize a DetectionTrainer object for training YOLO object detection models.
|
||
|
||
Args:
|
||
cfg (dict, optional): Default configuration dictionary containing training parameters.
|
||
overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
|
||
_callbacks (list, optional): List of callback functions to be executed during training.
|
||
"""
|
||
super().__init__(cfg, overrides, _callbacks)
|
||
|
||
def build_dataset(self, img_path: str, mode: str = "train", batch: int | None = None):
|
||
"""Build YOLO Dataset for training or validation.
|
||
|
||
Args:
|
||
img_path (str): Path to the folder containing images.
|
||
mode (str): 'train' mode or 'val' mode, users are able to customize different augmentations for each mode.
|
||
batch (int, optional): Size of batches, this is for 'rect' mode.
|
||
|
||
Returns:
|
||
(Dataset): YOLO dataset object configured for the specified mode.
|
||
"""
|
||
gs = max(int(unwrap_model(self.model).stride.max()), 32)
|
||
return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)
|
||
|
||
def get_dataloader(self, dataset_path: str, batch_size: int = 16, rank: int = 0, mode: str = "train"):
|
||
"""Construct and return dataloader for the specified mode.
|
||
|
||
Args:
|
||
dataset_path (str): Path to the dataset.
|
||
batch_size (int): Number of images per batch.
|
||
rank (int): Process rank for distributed training.
|
||
mode (str): 'train' for training dataloader, 'val' for validation dataloader.
|
||
|
||
Returns:
|
||
(DataLoader): PyTorch dataloader object.
|
||
"""
|
||
assert mode in {"train", "val"}, f"Mode must be 'train' or 'val', not {mode}."
|
||
with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP
|
||
dataset = self.build_dataset(dataset_path, mode, batch_size)
|
||
shuffle = mode == "train"
|
||
if getattr(dataset, "rect", False) and shuffle and not np.all(dataset.batch_shapes == dataset.batch_shapes[0]):
|
||
LOGGER.warning("'rect=True' is incompatible with DataLoader shuffle, setting shuffle=False")
|
||
shuffle = False
|
||
return build_dataloader(
|
||
dataset,
|
||
batch=batch_size,
|
||
workers=self.args.workers if mode == "train" else self.args.workers * 2,
|
||
shuffle=shuffle,
|
||
rank=rank,
|
||
drop_last=self.args.compile and mode == "train",
|
||
)
|
||
|
||
def preprocess_batch(self, batch: dict) -> dict:
|
||
"""Preprocess a batch of images by scaling and converting to float.
|
||
|
||
Args:
|
||
batch (dict): Dictionary containing batch data with 'img' tensor.
|
||
|
||
Returns:
|
||
(dict): Preprocessed batch with normalized images.
|
||
"""
|
||
for k, v in batch.items():
|
||
if isinstance(v, torch.Tensor):
|
||
batch[k] = v.to(self.device, non_blocking=self.device.type == "cuda")
|
||
batch["img"] = batch["img"].float() / 255
|
||
if self.args.multi_scale > 0.0:
|
||
imgs = batch["img"]
|
||
sz = (
|
||
random.randrange(
|
||
int(self.args.imgsz * (1.0 - self.args.multi_scale)),
|
||
int(self.args.imgsz * (1.0 + self.args.multi_scale) + self.stride),
|
||
)
|
||
// self.stride
|
||
* self.stride
|
||
) # size
|
||
sf = sz / max(imgs.shape[2:]) # scale factor
|
||
if sf != 1:
|
||
ns = [
|
||
math.ceil(x * sf / self.stride) * self.stride for x in imgs.shape[2:]
|
||
] # new shape (stretched to gs-multiple)
|
||
imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
|
||
batch["img"] = imgs
|
||
return batch
|
||
|
||
def set_model_attributes(self):
|
||
"""Set model attributes based on dataset information."""
|
||
# Nl = de_parallel(self.model).model[-1].nl # number of detection layers (to scale hyps)
|
||
# self.args.box *= 3 / nl # scale to layers
|
||
# self.args.cls *= self.data["nc"] / 80 * 3 / nl # scale to classes and layers
|
||
# self.args.cls *= (self.args.imgsz / 640) ** 2 * 3 / nl # scale to image size and layers
|
||
self.model.nc = self.data["nc"] # attach number of classes to model
|
||
self.model.names = self.data["names"] # attach class names to model
|
||
self.model.args = self.args # attach hyperparameters to model
|
||
if getattr(self.model, "end2end"):
|
||
self.model.set_head_attr(max_det=self.args.max_det)
|
||
# TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
|
||
|
||
def get_model(self, cfg: str | None = None, weights: str | None = None, verbose: bool = True):
|
||
"""Return a YOLO detection model.
|
||
|
||
Args:
|
||
cfg (str, optional): Path to model configuration file.
|
||
weights (str, optional): Path to model weights.
|
||
verbose (bool): Whether to display model information.
|
||
|
||
Returns:
|
||
(DetectionModel): YOLO detection model.
|
||
"""
|
||
model = DetectionModel(cfg, nc=self.data["nc"], ch=self.data["channels"], verbose=verbose and RANK == -1)
|
||
if weights:
|
||
model.load(weights)
|
||
return model
|
||
|
||
def get_validator(self):
|
||
"""Return a DetectionValidator for YOLO model validation."""
|
||
self.loss_names = "box_loss", "cls_loss", "dfl_loss", "diff_loss"
|
||
return yolo.detect.DetectionValidator(
|
||
self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
|
||
)
|
||
|
||
def label_loss_items(self, loss_items: list[float] | None = None, prefix: str = "train"):
|
||
"""Return a loss dict with labeled training loss items tensor.
|
||
|
||
Args:
|
||
loss_items (list[float], optional): List of loss values.
|
||
prefix (str): Prefix for keys in the returned dictionary.
|
||
|
||
Returns:
|
||
(dict | list): Dictionary of labeled loss items if loss_items is provided, otherwise list of keys.
|
||
"""
|
||
keys = [f"{prefix}/{x}" for x in self.loss_names]
|
||
if loss_items is not None:
|
||
loss_items = [round(float(x), 5) for x in loss_items] # convert tensors to 5 decimal place floats
|
||
return dict(zip(keys, loss_items))
|
||
else:
|
||
return keys
|
||
|
||
def progress_string(self):
|
||
"""Return a formatted string of training progress with epoch, GPU memory, loss, instances and size."""
|
||
leadw, colw = 7, 9
|
||
loss_names = list(self.loss_names)
|
||
extra_names = list(self.progress_extra_names())
|
||
return ("\n" + f"%{leadw}s" + f"%{colw}s" * (3 + len(loss_names) + len(extra_names))) % (
|
||
"Epoch",
|
||
"GPU_mem",
|
||
*loss_names,
|
||
"Inst",
|
||
"Size",
|
||
*extra_names,
|
||
)
|
||
|
||
def progress_extra_names(self):
|
||
"""Return extra per-batch progress columns appended after the loss items."""
|
||
return ()
|
||
|
||
def progress_extra_values(self):
|
||
"""Return extra per-batch progress values appended after the loss items."""
|
||
return ()
|
||
|
||
def plot_training_samples(self, batch: dict[str, Any], ni: int) -> None:
|
||
"""Plot training samples with their annotations.
|
||
|
||
Args:
|
||
batch (dict[str, Any]): Dictionary containing batch data.
|
||
ni (int): Batch index used for naming the output file.
|
||
"""
|
||
plot_images(
|
||
labels=batch,
|
||
paths=batch["im_file"],
|
||
fname=self.save_dir / f"train_batch{ni}.jpg",
|
||
on_plot=self.on_plot,
|
||
)
|
||
|
||
def plot_training_labels(self):
|
||
"""Create a labeled training plot of the YOLO model."""
|
||
boxes = np.concatenate([lb["bboxes"] for lb in self.train_loader.dataset.labels], 0)
|
||
cls = np.concatenate([lb["cls"] for lb in self.train_loader.dataset.labels], 0)
|
||
plot_labels(boxes, cls.squeeze(), names=self.data["names"], save_dir=self.save_dir, on_plot=self.on_plot)
|
||
|
||
def auto_batch(self):
|
||
"""Get optimal batch size by calculating memory occupation of model.
|
||
|
||
Returns:
|
||
(int): Optimal batch size.
|
||
"""
|
||
with override_configs(self.args, overrides={"cache": False}) as self.args:
|
||
train_dataset = self.build_dataset(self.data["train"], mode="train", batch=16)
|
||
max_num_obj = max(len(label["cls"]) for label in train_dataset.labels) * 4 # 4 for mosaic augmentation
|
||
del train_dataset # free memory
|
||
return super().auto_batch(max_num_obj)
|
||
|
||
|
||
class GroundDetectionValidator(DetectionValidator):
|
||
"""A validator for ground 2D detection.
|
||
|
||
This validator extends DetectionValidator for ground 2D detection tasks.
|
||
YUV444 to BGR conversion is handled in the dataloader, so no special preprocessing needed here.
|
||
"""
|
||
|
||
def _use_roi_metrics_only(self, batch: dict[str, Any]) -> bool:
|
||
"""Return whether validation metrics should ignore virtual-camera samples."""
|
||
return bool(getattr(self.args, "roi_metrics_only", False)) and isinstance(batch.get("camera_mode"), (list, tuple))
|
||
|
||
def _is_roi_metric_sample(self, batch: dict[str, Any], sample_idx: int) -> bool:
|
||
"""Return whether a sample should contribute to validation metrics."""
|
||
camera_mode = batch.get("camera_mode")
|
||
if not self._use_roi_metrics_only(batch):
|
||
return True
|
||
return sample_idx < len(camera_mode) and camera_mode[sample_idx] == "roi"
|
||
|
||
def preprocess(self, batch: dict[str, Any]) -> dict[str, Any]:
|
||
"""Preprocess batch (YUV444→BGR conversion already done in dataloader).
|
||
|
||
Args:
|
||
batch (dict[str, Any]): Batch containing images and annotations.
|
||
|
||
Returns:
|
||
(dict[str, Any]): Preprocessed batch with normalized BGR images.
|
||
"""
|
||
# Move to device
|
||
for k, v in batch.items():
|
||
if isinstance(v, torch.Tensor):
|
||
batch[k] = v.to(self.device, non_blocking=self.device.type == "cuda")
|
||
|
||
# Normalize to [0, 1] (images are already in BGR format from dataloader)
|
||
batch["img"] = (batch["img"].half() if self.args.half else batch["img"].float()) / 256
|
||
return batch
|
||
|
||
def update_metrics(self, preds: list[dict[str, torch.Tensor]], batch: dict[str, Any]) -> None:
|
||
"""Update metrics with new predictions and ground truth, optionally ROI-only."""
|
||
for si, pred in enumerate(preds):
|
||
if not self._is_roi_metric_sample(batch, si):
|
||
continue
|
||
self.seen += 1
|
||
pbatch = self._prepare_batch(si, batch)
|
||
predn = self._prepare_pred(pred)
|
||
|
||
cls = pbatch["cls"].cpu().numpy()
|
||
no_pred = predn["cls"].shape[0] == 0
|
||
self.metrics.update_stats(
|
||
{
|
||
**self._process_batch(predn, pbatch),
|
||
"target_cls": cls,
|
||
"target_img": np.unique(cls),
|
||
"conf": np.zeros(0) if no_pred else predn["conf"].cpu().numpy(),
|
||
"pred_cls": np.zeros(0) if no_pred else predn["cls"].cpu().numpy(),
|
||
}
|
||
)
|
||
if self.args.plots:
|
||
self.confusion_matrix.process_batch(predn, pbatch, conf=self.args.conf)
|
||
if self.args.visualize:
|
||
self.confusion_matrix.plot_matches(batch["img"][si], pbatch["im_file"], self.save_dir)
|
||
|
||
if no_pred:
|
||
continue
|
||
|
||
if self.args.save_json or self.args.save_txt:
|
||
predn_scaled = self.scale_preds(predn, pbatch)
|
||
if self.args.save_json:
|
||
self.pred_to_json(predn_scaled, pbatch)
|
||
if self.args.save_txt:
|
||
self.save_one_txt(
|
||
predn_scaled,
|
||
self.args.save_conf,
|
||
pbatch["ori_shape"],
|
||
self.save_dir / "labels" / f"{Path(pbatch['im_file']).stem}.txt",
|
||
)
|
||
|
||
|
||
class GroundDetectionTrainer(DetectionTrainer):
|
||
"""A class extending DetectionTrainer for training ground 2D detection with difficulty scores.
|
||
|
||
This trainer specializes in ground 2D detection tasks with custom annotation format including:
|
||
- String class names mapped to numeric IDs via class_map
|
||
- Difficulty scores for each bounding box
|
||
- Difficulty-based loss weighting
|
||
- YUV444 to BGR conversion (handled in dataloader)
|
||
|
||
Attributes:
|
||
model (DetectionModel): The YOLO detection model being trained.
|
||
data (dict): Dictionary containing dataset information including class_map and number of classes.
|
||
loss_names (tuple): Names of the loss components used in training (box_loss, cls_loss, dfl_loss).
|
||
|
||
Examples:
|
||
>>> from ultralytics.models.yolo.detect import GroundDetectionTrainer
|
||
>>> args = dict(model="yolo26n.pt", data="ground_dataset.yaml", epochs=100)
|
||
>>> trainer = GroundDetectionTrainer(overrides=args)
|
||
>>> trainer.train()
|
||
"""
|
||
|
||
def __init__(self, cfg=DEFAULT_CFG, overrides: dict[str, Any] | None = None, _callbacks=None):
|
||
"""Initialize a GroundDetectionTrainer object.
|
||
|
||
Args:
|
||
cfg (dict, optional): Default configuration dictionary containing training parameters.
|
||
overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
|
||
_callbacks (list, optional): List of callback functions to be executed during training.
|
||
"""
|
||
super().__init__(cfg, overrides, _callbacks)
|
||
# YUV444 conversion is now handled in the dataloader
|
||
|
||
def preprocess_batch(self, batch: dict) -> dict:
|
||
"""Preprocess a batch of images.
|
||
|
||
Args:
|
||
batch (dict): Dictionary containing batch data with 'img' tensor.
|
||
|
||
Returns:
|
||
(dict): Preprocessed batch with normalized images in BGR format.
|
||
"""
|
||
# Move batch to device first
|
||
for k, v in batch.items():
|
||
if isinstance(v, torch.Tensor):
|
||
batch[k] = v.to(self.device, non_blocking=self.device.type == "cuda")
|
||
|
||
# Normalize to [0, 1] (YUV444→BGR conversion already done in dataloader)
|
||
batch["img"] = batch["img"].float() / 256
|
||
|
||
# Multi-scale training
|
||
if self.args.multi_scale > 0.0:
|
||
imgs = batch["img"]
|
||
sz = (
|
||
random.randrange(
|
||
int(self.args.imgsz * (1.0 - self.args.multi_scale)),
|
||
int(self.args.imgsz * (1.0 + self.args.multi_scale) + self.stride),
|
||
)
|
||
// self.stride
|
||
* self.stride
|
||
) # size
|
||
sf = sz / max(imgs.shape[2:]) # scale factor
|
||
if sf != 1:
|
||
ns = [
|
||
math.ceil(x * sf / self.stride) * self.stride for x in imgs.shape[2:]
|
||
] # new shape (stretched to gs-multiple)
|
||
imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
|
||
batch["img"] = imgs
|
||
return batch
|
||
|
||
def build_dataset(self, img_path: str, mode: str = "train", batch: int | None = None):
|
||
"""Build YOLO Dataset for training or validation.
|
||
|
||
Args:
|
||
img_path (str): Path to the folder containing images.
|
||
mode (str): 'train' mode or 'val' mode.
|
||
batch (int, optional): Size of batches, this is for 'rect' mode.
|
||
|
||
Returns:
|
||
(Dataset): YOLO dataset object configured for the specified mode.
|
||
"""
|
||
return super().build_dataset(img_path, mode, batch)
|
||
|
||
def set_model_attributes(self):
|
||
"""Set model attributes based on ground dataset information."""
|
||
# The class_map to names conversion is now handled in check_det_dataset
|
||
# This method just calls the parent implementation
|
||
super().set_model_attributes()
|
||
|
||
def get_model(self, cfg: str | None = None, weights: str | None = None, verbose: bool = True):
|
||
"""Return a YOLO detection model with ground loss function.
|
||
|
||
Args:
|
||
cfg (str, optional): Path to model configuration file.
|
||
weights (str, optional): Path to model weights.
|
||
verbose (bool): Whether to display model information.
|
||
|
||
Returns:
|
||
(GroundDetectionModel): YOLO detection model with ground loss.
|
||
"""
|
||
model = GroundDetectionModel(cfg, nc=self.data["nc"], ch=self.data["channels"], verbose=verbose and RANK == -1)
|
||
if weights:
|
||
model.load(weights)
|
||
return model
|
||
|
||
def get_validator(self):
|
||
"""Return a GroundDetectionValidator for YOLO model validation.
|
||
|
||
Returns:
|
||
(GroundDetectionValidator): Validator for ground detection.
|
||
"""
|
||
self.loss_names = "box_loss", "cls_loss", "dfl_loss"
|
||
return GroundDetectionValidator(
|
||
self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
|
||
)
|
||
|
||
|
||
class Ground3DDetectionModel(DetectionModel):
|
||
"""YOLO detection model with joint 2D+3D ground loss."""
|
||
|
||
def init_criterion(self):
|
||
"""Initialize the loss criterion for ground 3D detection."""
|
||
from ultralytics.utils.loss import E2EGround3DLoss, v8Detection3DLoss
|
||
|
||
return E2EGround3DLoss(self) if getattr(self, "end2end", False) else v8Detection3DLoss(self)
|
||
|
||
|
||
class Ground3DDetectionTrainer(GroundDetectionTrainer):
|
||
"""Trainer for joint 2D+3D ground detection.
|
||
|
||
Extends GroundDetectionTrainer with:
|
||
- Ground3DDetectionModel with 3D loss
|
||
- YOLOGround3DDataset for on-the-fly 3D label loading
|
||
- 3D loss weight ramping during training
|
||
- Support for rectangular imgsz (e.g., [704, 352])
|
||
"""
|
||
|
||
def setup_model(self):
|
||
"""Load 3D model config and optionally pretrained weights."""
|
||
from pathlib import Path
|
||
|
||
from ultralytics.engine.trainer import load_checkpoint
|
||
|
||
cfg = self.model # yaml path, e.g. "yolo26-3d.yaml"
|
||
weights = None
|
||
explicit_scale = getattr(self, "explicit_model_scale", None) # 尝试读取显式模型 scale，比如 n/s/l/m/x。如果对象上没有这个属性，就返回 None
|
||
|
||
# Load pretrained weights if specified (e.g. "yolo26s-pretrain.pt")
|
||
if isinstance(self.args.pretrained, (str, Path)):
|
||
weights, _ = load_checkpoint(self.args.pretrained)
|
||
if explicit_scale is None:
|
||
pretrained_stem = Path(self.args.pretrained).stem
|
||
if pretrained_stem.startswith("yolo26") and len(pretrained_stem) > len("yolo26"):
|
||
scale = pretrained_stem[len("yolo26")]
|
||
if scale in "nslmx":
|
||
explicit_scale = scale
|
||
self.explicit_model_scale = scale
|
||
|
||
cfg = self._resolve_model_cfg(cfg, explicit_scale)
|
||
|
||
self.model = self.get_model(cfg=cfg, weights=weights, verbose=RANK == -1)
|
||
return None
|
||
|
||
@staticmethod
|
||
def _resolve_model_cfg(cfg, explicit_scale=None):
|
||
"""Resolve the model YAML/config dict and optionally force an explicit scale."""
|
||
from ultralytics.nn.tasks import yaml_model_load
|
||
|
||
if isinstance(cfg, (str, Path)):
|
||
cfg = yaml_model_load(cfg)
|
||
else:
|
||
cfg = dict(cfg)
|
||
if explicit_scale:
|
||
cfg["scale"] = explicit_scale
|
||
return cfg
|
||
|
||
@staticmethod
|
||
def _backbone_param_keys(model):
|
||
"""Return backbone parameter keys based on the parsed YAML backbone definition.定义静态方法，用于找出模型 backbone 部分的参数名"""
|
||
backbone = (getattr(model, "yaml", None) or {}).get("backbone")
|
||
if not isinstance(backbone, list) or not backbone:
|
||
raise RuntimeError("Unable to verify pretrained backbone load because the model YAML has no backbone.")
|
||
prefixes = tuple(f"model.{i}." for i in range(len(backbone)))
|
||
return sorted(k for k, _ in model.named_parameters() if k.startswith(prefixes))
|
||
|
||
def _verify_pretrained_backbone_load(self, model, weights):
|
||
"""Raise if any backbone parameter tensor fails to transfer from pretrained weights. 根据模型 yaml 中的 backbone 定义返回 backbone 参数 key"""
|
||
source_model = weights["model"] if isinstance(weights, dict) else weights
|
||
source_params = dict(source_model.named_parameters())
|
||
target_params = dict(model.named_parameters())
|
||
target_backbone_keys = self._backbone_param_keys(model)
|
||
if not target_backbone_keys:
|
||
raise RuntimeError("Unable to find any backbone parameters to verify in the target Ground3D model.")
|
||
|
||
source_state_dict = source_model.float().state_dict()
|
||
target_state_dict = model.state_dict()
|
||
transferred_state_dict = intersect_dicts(source_state_dict, target_state_dict)
|
||
|
||
missing = [k for k in target_backbone_keys if k not in source_params]
|
||
mismatched = [
|
||
k
|
||
for k in target_backbone_keys
|
||
if k in source_params and (source_params[k].shape != target_params[k].shape or k not in transferred_state_dict)
|
||
]
|
||
if missing or mismatched:
|
||
details = []
|
||
if missing:
|
||
details.append(f"missing={len(missing)} (e.g. {', '.join(missing[:5])})")
|
||
if mismatched:
|
||
details.append(f"shape_mismatch={len(mismatched)} (e.g. {', '.join(mismatched[:5])})")
|
||
scale = (getattr(model, "yaml", None) or {}).get("scale")
|
||
scale_msg = f" scale='{scale}'" if scale else ""
|
||
raise RuntimeError(
|
||
"Pretrained weights failed strict backbone verification for the Ground3D model"
|
||
f"{scale_msg}. "
|
||
"Use a matching '--model yolo26n/yolo26s/...' shorthand or a compatible checkpoint. "
|
||
f"Details: {' | '.join(details)}"
|
||
)
|
||
|
||
LOGGER.info(f"Verified strict pretrained backbone load for {len(target_backbone_keys)} parameter tensors")
|
||
|
||
def get_model(self, cfg=None, weights=None, verbose=True):
|
||
"""Return a YOLO detection model with 3D loss function."""
|
||
cfg = self._resolve_model_cfg(cfg, getattr(self, "explicit_model_scale", None))
|
||
model = Ground3DDetectionModel(cfg, nc=self.data["nc"], ch=self.data["channels"], verbose=verbose and RANK == -1)
|
||
if weights:
|
||
if getattr(self, "strict_backbone_pretrained", True):
|
||
self._verify_pretrained_backbone_load(model, weights)
|
||
model.load(weights)
|
||
# Attach 3D config to model for loss function
|
||
model.norm_scales_3d = self.data.get("norm_scales_3d", {})
|
||
model.face_3d_classes = set(self.data.get("face_3d_classes", []))
|
||
model.complete_3d_classes = set(self.data.get("complete_3d_classes", []))
|
||
fake_3d_classes = self.data.get("fake_3d_classes")
|
||
if fake_3d_classes is None:
|
||
class_map = self.data.get("class_map", {})
|
||
fake_3d_classes = [v for k, v in class_map.items() if isinstance(k, str) and k.endswith("_fake")]
|
||
model.fake_3d_classes = set(fake_3d_classes or [])
|
||
# Propagate norm_scales_3d to Detect3D head for denormalization
|
||
from ultralytics.nn.modules.head import Detect3D
|
||
|
||
for m in model.modules():
|
||
if isinstance(m, Detect3D):
|
||
m.norm_scales_3d = model.norm_scales_3d
|
||
return model
|
||
|
||
def plot_training_labels(self):
|
||
"""Skip label plotting because Ground3D labels are parsed on-the-fly."""
|
||
LOGGER.info("Skipping label plot (on-the-fly label loading)")
|
||
|
||
def auto_batch(self):
|
||
"""Ground3D GT-list datasets require an explicit batch size."""
|
||
raise ValueError("AutoBatch is not supported for Ground3D GT-list datasets. Please set --batch explicitly.")
|
||
|
||
def build_dataset(self, img_path, mode="train", batch=None):
|
||
"""Build YOLOGround3DDataset for 3D detection training. 定义数据集构建函数。mode 可以是 "train" 或 "val" """
|
||
from ultralytics.data.dataset import YOLOGround3DDataset
|
||
|
||
gs = max(int(unwrap_model(self.model).stride.max()), 32)
|
||
cfg = self.args
|
||
if RANK in {-1, 0}:
|
||
LOGGER.info(f"{mode}: Building Ground3D dataset from {img_path}")
|
||
dataset = YOLOGround3DDataset(
|
||
img_path=img_path,
|
||
imgsz=cfg.imgsz,
|
||
batch_size=batch,
|
||
augment=mode == "train",
|
||
hyp=cfg,
|
||
rect=cfg.rect or (mode == "val"),
|
||
cache=cfg.cache or None,
|
||
single_cls=cfg.single_cls or False,
|
||
stride=gs,
|
||
pad=0.0 if mode == "train" else 0.5,
|
||
prefix=f"{mode}: ",
|
||
task=cfg.task,
|
||
classes=cfg.classes,
|
||
data=self.data,
|
||
fraction=cfg.fraction if mode == "train" else 1.0,
|
||
)
|
||
if RANK in {-1, 0}:
|
||
LOGGER.info(f"{mode}: Ground3D dataset ready with {len(dataset):,} samples")
|
||
return dataset
|
||
|
||
def preprocess_batch(self, batch):
|
||
"""Preprocess batch and ramp 3D loss weight. 除了预处理 batch，还会动态调整 3D loss 权重"""
|
||
batch = super().preprocess_batch(batch)
|
||
model = unwrap_model(self.model) # 拿到真实模型对象，避免 DDP 包装影响访问属性
|
||
if hasattr(model, "criterion"):
|
||
criterion = model.criterion
|
||
w = self._compute_3d_weight() # 计算当前 epoch 应该使用的 3D loss 权重
|
||
if hasattr(criterion, "one2many"):
|
||
criterion.one2many.loss_3d_weight = w
|
||
criterion.one2one.loss_3d_weight = w
|
||
elif hasattr(criterion, "loss_3d_weight"):
|
||
criterion.loss_3d_weight = w
|
||
return batch
|
||
|
||
def progress_extra_names(self):
|
||
"""Expose lr and current 3D loss weight in the live training progress bar. 进度条额外显示两个值：学习率 lr 和当前 3D loss 权重 w3d"""
|
||
return ("lr", "w3d")
|
||
|
||
def progress_extra_values(self):
|
||
"""Return current lr and active 3D loss weight for the live training progress bar. 定义训练进度条额外显示字段值"""
|
||
lr = self.optimizer.param_groups[0]["lr"] if getattr(self, "optimizer", None) and self.optimizer.param_groups else 0.0
|
||
criterion = getattr(unwrap_model(self.model), "criterion", None)
|
||
if hasattr(criterion, "one2many"):
|
||
w3d = float(getattr(criterion.one2many, "loss_3d_weight", 0.0))
|
||
else:
|
||
w3d = float(getattr(criterion, "loss_3d_weight", 0.0)) if criterion is not None else 0.0
|
||
return (f"{lr:.1e}", f"{w3d:.3g}")
|
||
|
||
def _compute_3d_weight(self):
|
||
"""Progressively ramp 3D loss weight after the 2D-only warmup stage."""
|
||
warmup = getattr(self.args, "loss_3d_warmup_epochs", 10.0) # 读取 3D loss warmup epoch，默认 10。warmup 阶段 3D loss 权重为 0，只训练 2D。
|
||
ramp_epochs = getattr(self.args, "loss_3d_ramp_epochs", 10.0) # 读取 ramp 阶段长度，默认 10 个 epoch。 也就是在这个范围逐步增长
|
||
max_weight = float(getattr(self.args, "loss_3d_weight_max", 0.1)) # 读取 3D loss 最大权重，默认 0.1。
|
||
epoch = self.epoch
|
||
warmup = float(warmup)
|
||
if epoch < warmup:
|
||
return 0.0
|
||
progress = min((epoch - warmup) / ramp_epochs, 1.0)
|
||
return max_weight * progress
|
||
|
||
def get_validator(self):
|
||
"""Return validator with 3D loss names."""
|
||
self.loss_names = (
|
||
"box", "cls", "dfl", "diff",
|
||
"z3d", "uv", "size", "ycls", "ydeg",
|
||
"ccls", "fz", "fuv", "fsize", "fcls", "euv", "ez",
|
||
)
|
||
return Ground3DDetectionValidator(
|
||
self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
|
||
)
|
||
"""
|
||
box 2D bbox loss
|
||
cls 分类 loss
|
||
dfl Distribution Focal Loss
|
||
diff difficulty 难度分支 loss
|
||
z3d 3D 深度 loss
|
||
uv 3D 关键点/投影点 2D 坐标 loss
|
||
size 3D 尺寸 loss
|
||
ycls yaw 分类 loss
|
||
ydeg yaw 角度回归 loss
|
||
ccls complete/类别相关 3D 分类 loss
|
||
fz face 深度 loss
|
||
fuv face 2D 点 loss
|
||
fsize face 尺寸 loss
|
||
fcls face 类别/可见性 loss
|
||
euv edge 2D 点 loss
|
||
ez edge 深度 loss
|
||
"""
|
||
|
||
class Ground3DDetectionValidator(GroundDetectionValidator):
|
||
"""Validator for ground 3D detection with 3D metrics computation.
|
||
|
||
Extends GroundDetectionValidator to:
|
||
- Capture 3D predictions from model output during postprocess
|
||
- Compute 3D metrics (depth, orientation, center, size errors) for matched GT-pred pairs
|
||
- Include 3D metrics in validation results for TensorBoard logging
|
||
"""
|
||
|
||
def __init__(self, dataloader=None, save_dir=None, args=None, _callbacks=None):
|
||
"""Initialize with 3D metrics storage."""
|
||
super().__init__(dataloader, save_dir, args, _callbacks)
|
||
self.stats_3d = {"whole": [], "face": []}
|
||
self.metrics_3d_results = self._empty_3d_metrics()
|
||
self._preds_3d_selected = None
|
||
self._preds_edge_selected = None
|
||
self._preds_diff_selected = None
|
||
self._anchors_selected = None
|
||
self._strides_selected = None
|
||
self._current_batch_calib = None
|
||
self.face_3d_classes = set()
|
||
self.complete_3d_classes = set()
|
||
self.fake_3d_classes = set()
|
||
|
||
@staticmethod
|
||
def _empty_3d_metrics():
|
||
"""Return default grouped 3D metrics used for logging when no matches are available."""
|
||
from ultralytics.utils.metrics_3d import empty_3d_metrics
|
||
|
||
return {
|
||
"whole": empty_3d_metrics(include_orient=True, include_size=True, include_uv=True),
|
||
"face": empty_3d_metrics(include_orient=False, include_size=True, include_uv=True, include_visible_orient=True),
|
||
}
|
||
|
||
def init_metrics(self, model):
|
||
"""Initialize metrics and extract 3D class config from model."""
|
||
super().init_metrics(model)
|
||
self.stats_3d = {"whole": [], "face": []}
|
||
self.metrics_3d_results = self._empty_3d_metrics()
|
||
self._current_batch_calib = None
|
||
# Extract 3D config from model
|
||
if hasattr(model, "module"):
|
||
model = model.module
|
||
self.face_3d_classes = getattr(model, "face_3d_classes", set())
|
||
self.complete_3d_classes = getattr(model, "complete_3d_classes", set())
|
||
self.fake_3d_classes = getattr(model, "fake_3d_classes", set())
|
||
|
||
def get_desc(self) -> str:
|
||
"""Return a compact loss-aligned 3D metric title row for tqdm."""
|
||
trainer = getattr(self, "trainer", None)
|
||
loss_ref_names = getattr(trainer, "loss_names", None) or (
|
||
"box",
|
||
"cls",
|
||
"dfl",
|
||
"z3d",
|
||
"uv",
|
||
"size",
|
||
"ycls",
|
||
"ydeg",
|
||
"ccls",
|
||
"fz",
|
||
"fuv",
|
||
"fsize",
|
||
"fcls",
|
||
"euv",
|
||
"ez",
|
||
)
|
||
leadw, colw = 7, 9
|
||
desc_values = [
|
||
"",
|
||
"Metric3D",
|
||
"Wmatch",
|
||
"Fmatch",
|
||
"-",
|
||
"WdAbs",
|
||
"Wuv",
|
||
"Wsize",
|
||
"-",
|
||
"Wyaw",
|
||
"-",
|
||
"FdAbs",
|
||
"Fuv",
|
||
"Fsize",
|
||
"-",
|
||
"Vyaw",
|
||
]
|
||
desc_values.extend(["-"] * (1 + len(loss_ref_names) - len(desc_values)))
|
||
return (f"%{leadw}s" + f"%{colw}s" * len(loss_ref_names)) % tuple(desc_values[: 1 + len(loss_ref_names)])
|
||
|
||
def get_header(self) -> str:
|
||
"""Return no extra one-time header; tqdm description already contains the compact metric titles."""
|
||
return ""
|
||
|
||
def _face_desc(self) -> str:
|
||
"""Return header for the face-only 3D metric line."""
|
||
return ("%22s" + "%11s" * (2 + len(self.metrics.keys) + 6)) % (
|
||
"Class",
|
||
"",
|
||
"",
|
||
*("" for _ in self.metrics.keys),
|
||
"FdAbs",
|
||
"FdRel",
|
||
"FdRMSE",
|
||
"Fctr",
|
||
"Fuv",
|
||
"Fmatch",
|
||
)
|
||
|
||
def preprocess(self, batch):
|
||
"""Preprocess batch and keep per-image calibration for 3D prediction restoration."""
|
||
self._current_batch_calib = batch.get("calib")
|
||
return super().preprocess(batch)
|
||
|
||
def postprocess(self, preds):
|
||
"""Capture 3D predictions, restore depth once, and filter by confidence."""
|
||
# Extract 3D predictions from raw model output: (y, preds_dict)
|
||
self._preds_3d_filtered = None
|
||
self._anchors_filtered = None
|
||
self._strides_filtered = None
|
||
self._preds_edge_filtered = None
|
||
|
||
if isinstance(preds, (list, tuple)) and len(preds) >= 2:
|
||
y = preds[0] # (B, k, 6) from head's postprocess
|
||
preds_dict = preds[1]
|
||
one2one = preds_dict.get("one2one", {})
|
||
preds_3d_sel = one2one.get("preds_3d_selected") # (B, k, 41) or None
|
||
preds_3d_fake_sel = one2one.get("preds_3d_fake_selected") # (B, k, 41) or None
|
||
preds_edge_sel = one2one.get("preds_edge_selected") # (B, k, 60) or None
|
||
preds_diff_sel = one2one.get("preds_diff_selected") # (B, k, 1) or None
|
||
anchors_sel = one2one.get("anchors_selected") # (B, 2, k) or None
|
||
strides_sel = one2one.get("strides_selected") # (B, k) or None
|
||
|
||
if preds_3d_sel is not None and anchors_sel is not None and strides_sel is not None:
|
||
preds_3d_sel = preds_3d_sel.clone()
|
||
if preds_3d_fake_sel is not None and self.fake_3d_classes:
|
||
fake_cls = torch.as_tensor(sorted(self.fake_3d_classes), device=y.device, dtype=y.dtype)
|
||
fake_mask = (y[..., 5:6] == fake_cls.view(1, 1, -1)).any(dim=-1)
|
||
preds_3d_sel[fake_mask] = preds_3d_fake_sel[fake_mask]
|
||
if preds_edge_sel is not None:
|
||
preds_edge_sel = preds_edge_sel.clone()
|
||
batch_calib = self._current_batch_calib
|
||
if batch_calib is not None:
|
||
z_channels = (0, 6, 12, 18, 24)
|
||
for si in range(min(preds_3d_sel.shape[0], len(batch_calib))):
|
||
calib = batch_calib[si]
|
||
if calib is None:
|
||
continue
|
||
depth_scale = calib.get("depth_scale", 1.0)
|
||
if depth_scale == 1.0:
|
||
continue
|
||
for ch in z_channels:
|
||
preds_3d_sel[si, :, ch] *= depth_scale
|
||
if preds_edge_sel is not None:
|
||
preds_edge_sel[si, :, 2::3] *= depth_scale
|
||
|
||
# Mirror end2end NMS: filter by confidence and limit to max_det
|
||
# NMS does: pred[pred[:, 4] > conf_thres][:max_det]
|
||
self._preds_3d_filtered = []
|
||
self._preds_edge_filtered = []
|
||
self._anchors_filtered = []
|
||
self._strides_filtered = []
|
||
for si in range(y.shape[0]):
|
||
conf_mask = y[si, :, 4] > self.args.conf
|
||
p3d_f = preds_3d_sel[si][conf_mask][: self.args.max_det]
|
||
pedge_f = preds_edge_sel[si][conf_mask][: self.args.max_det] if preds_edge_sel is not None else None
|
||
anc_f = anchors_sel[si][:, conf_mask][:, : self.args.max_det]
|
||
str_f = strides_sel[si][conf_mask][: self.args.max_det]
|
||
self._preds_3d_filtered.append(p3d_f)
|
||
self._preds_edge_filtered.append(pedge_f)
|
||
self._anchors_filtered.append(anc_f)
|
||
self._strides_filtered.append(str_f)
|
||
|
||
# Also store unfiltered for TensorBoard visualization (uses its own conf filter)
|
||
self._preds_3d_selected = preds_3d_sel
|
||
self._preds_edge_selected = preds_edge_sel
|
||
self._preds_diff_selected = preds_diff_sel
|
||
self._anchors_selected = anchors_sel
|
||
self._strides_selected = strides_sel
|
||
|
||
return super().postprocess(preds)
|
||
|
||
def update_metrics(self, preds, batch):
|
||
"""Update 2D metrics and compute 3D metrics for matched GT-pred pairs."""
|
||
super().update_metrics(preds, batch)
|
||
|
||
labels_3d = batch.get("labels_3d")
|
||
if labels_3d is None or self._preds_3d_filtered is None:
|
||
return
|
||
|
||
self._compute_3d_metrics_for_batch(preds, batch)
|
||
|
||
def gather_stats(self) -> None:
|
||
"""Gather 2D and 3D stats from all GPUs."""
|
||
super().gather_stats()
|
||
if getattr(self, "_single_rank_eval", False) or not (dist.is_available() and dist.is_initialized()):
|
||
return
|
||
if RANK == 0:
|
||
gathered_stats_3d = [None] * dist.get_world_size()
|
||
dist.gather_object(self.stats_3d, gathered_stats_3d, dst=0)
|
||
merged_stats_3d = {"whole": [], "face": []}
|
||
for stats_dict in gathered_stats_3d:
|
||
for key in merged_stats_3d:
|
||
merged_stats_3d[key].extend(stats_dict.get(key, []))
|
||
self.stats_3d = merged_stats_3d
|
||
else:
|
||
dist.gather_object(self.stats_3d, None, dst=0)
|
||
self.stats_3d = {"whole": [], "face": []}
|
||
|
||
def _append_metric_attr(self, attr_store, attr):
|
||
"""Append a decoded 3D attribute dict into the metric store."""
|
||
if attr is None:
|
||
return
|
||
for key in attr_store:
|
||
attr_store[key].append(attr[key])
|
||
|
||
@staticmethod
|
||
def _metric_store():
|
||
"""Return an empty metric attribute store."""
|
||
return {"center": [], "depth": [], "yaw": [], "edge_yaw": [], "dims": [], "uv": []}
|
||
|
||
def _aggregate_metric_store(self, group, pred_store, gt_store, include_visible_orient=None):
|
||
"""Aggregate one metric group and append it to stats."""
|
||
from ultralytics.utils.metrics_3d import compute_3d_metrics_for_matched
|
||
|
||
if not pred_store["depth"]:
|
||
return
|
||
if include_visible_orient is None:
|
||
include_visible_orient = group == "face"
|
||
pred_np = {k: np.array(v) for k, v in pred_store.items()}
|
||
gt_np = {k: np.array(v) for k, v in gt_store.items()}
|
||
metrics = compute_3d_metrics_for_matched(
|
||
pred_np,
|
||
gt_np,
|
||
include_orient=group == "whole",
|
||
include_size=True,
|
||
include_uv=True,
|
||
include_visible_orient=include_visible_orient,
|
||
)
|
||
self.stats_3d[group].append(metrics)
|
||
|
||
@staticmethod
|
||
def _whole_metric_store():
|
||
"""Return empty stores for whole-position and whole-shape metrics."""
|
||
return {
|
||
"pos_pred": Ground3DDetectionValidator._metric_store(),
|
||
"pos_gt": Ground3DDetectionValidator._metric_store(),
|
||
"shape_pred": Ground3DDetectionValidator._metric_store(),
|
||
"shape_gt": Ground3DDetectionValidator._metric_store(),
|
||
}
|
||
|
||
def _aggregate_whole_metric_store(self, stores):
|
||
"""Aggregate whole-box metrics with separate validity for position and shape terms."""
|
||
from ultralytics.utils.metrics_3d import (
|
||
compute_3d_metrics_for_matched,
|
||
compute_orientation_error,
|
||
compute_size_error,
|
||
empty_3d_metrics,
|
||
)
|
||
|
||
pos_depth = stores["pos_pred"]["depth"]
|
||
shape_yaw = stores["shape_pred"]["yaw"]
|
||
if not pos_depth and not shape_yaw:
|
||
return
|
||
|
||
metrics = empty_3d_metrics(include_orient=True, include_size=True, include_uv=True)
|
||
metrics["matched"] = len(stores["shape_pred"]["depth"])
|
||
if pos_depth:
|
||
pos_pred = {k: np.array(v) for k, v in stores["pos_pred"].items()}
|
||
pos_gt = {k: np.array(v) for k, v in stores["pos_gt"].items()}
|
||
pos_metrics = compute_3d_metrics_for_matched(
|
||
pos_pred, pos_gt, include_orient=False, include_size=False, include_uv=True
|
||
)
|
||
for key in ("depth_abs", "depth_rel", "depth_rmse", "center", "uv"):
|
||
metrics[key] = pos_metrics[key]
|
||
metrics["_pos_matched"] = pos_metrics["matched"]
|
||
|
||
if shape_yaw:
|
||
shape_pred = {k: np.array(v) for k, v in stores["shape_pred"].items()}
|
||
shape_gt = {k: np.array(v) for k, v in stores["shape_gt"].items()}
|
||
metrics["orient"] = compute_orientation_error(shape_pred["yaw"], shape_gt["yaw"])
|
||
metrics["size"] = compute_size_error(shape_pred["dims"], shape_gt["dims"])
|
||
|
||
self.stats_3d["whole"].append(metrics)
|
||
|
||
def _aggregate_face_metric_store(self, pred_store, gt_store, pred_visible_yaw, pred_edge_visible_yaw, gt_visible_target_yaw):
|
||
"""Aggregate face metrics with object-wise visible-orientation accounting."""
|
||
from ultralytics.utils.metrics_3d import (
|
||
compute_3d_metrics_for_matched,
|
||
compute_visible_orientation_metrics,
|
||
empty_3d_metrics,
|
||
)
|
||
|
||
if not pred_store["depth"] and not gt_visible_target_yaw:
|
||
return
|
||
|
||
metrics = empty_3d_metrics(include_orient=False, include_size=True, include_uv=True, include_visible_orient=True)
|
||
if pred_store["depth"]:
|
||
pred_np = {k: np.array(v) for k, v in pred_store.items()}
|
||
gt_np = {k: np.array(v) for k, v in gt_store.items()}
|
||
face_metrics = compute_3d_metrics_for_matched(
|
||
pred_np,
|
||
gt_np,
|
||
include_orient=False,
|
||
include_size=True,
|
||
include_uv=True,
|
||
include_visible_orient=False,
|
||
)
|
||
for key in ("depth_abs", "depth_rel", "depth_rmse", "center", "matched", "uv", "size"):
|
||
metrics[key] = face_metrics[key]
|
||
|
||
if gt_visible_target_yaw:
|
||
metrics.update(
|
||
compute_visible_orientation_metrics(
|
||
np.asarray(pred_visible_yaw, dtype=np.float64),
|
||
np.asarray(pred_edge_visible_yaw, dtype=np.float64),
|
||
np.asarray(gt_visible_target_yaw, dtype=np.float64),
|
||
)
|
||
)
|
||
|
||
self.stats_3d["face"].append(metrics)
|
||
|
||
def _compute_3d_metrics_for_batch(self, preds, batch):
|
||
"""Compute whole-box and face-specific 3D metrics for all images in the batch."""
|
||
from ultralytics.utils.metrics import box_iou
|
||
from ultralytics.utils.plotting_3d import (
|
||
EDGE_YAW_MAX_LATERAL_DIST_M,
|
||
EDGE_YAW_VALID_VISIBILITY_SCORE_THRESH,
|
||
decode_edge_yaw_selection_from_prediction,
|
||
extract_3d_attrs_from_gt,
|
||
extract_3d_attrs_from_prediction,
|
||
is_gt_cut_object,
|
||
select_gt_visible_faces,
|
||
)
|
||
|
||
labels_3d = batch["labels_3d"]
|
||
batch_idx = batch.get("batch_idx")
|
||
batch_cls = batch.get("cls")
|
||
batch_calib = batch.get("calib") # list/tuple of dicts (one per image)
|
||
_, _, img_h, img_w = batch["img"].shape
|
||
|
||
if batch_idx is None or batch_cls is None:
|
||
return
|
||
|
||
if self._preds_3d_filtered is None:
|
||
return
|
||
|
||
face_visibility_score_thresh = float(
|
||
getattr(self.args, "face_visibility_score_thresh", DEFAULT_CFG.face_visibility_score_thresh)
|
||
)
|
||
edge_yaw_compare_score_thresh = float(EDGE_YAW_VALID_VISIBILITY_SCORE_THRESH)
|
||
edge_yaw_compare_max_lateral_dist_m = float(EDGE_YAW_MAX_LATERAL_DIST_M)
|
||
batch_idx_np = batch_idx.cpu().numpy().ravel()
|
||
cls_np = batch_cls.cpu().numpy().ravel()
|
||
labels_3d_np = labels_3d.cpu().numpy()
|
||
|
||
for si, pred in enumerate(preds):
|
||
if not self._is_roi_metric_sample(batch, si):
|
||
continue
|
||
if si >= len(self._preds_3d_filtered):
|
||
continue
|
||
|
||
calib = batch_calib[si] if batch_calib is not None and si < len(batch_calib) else None
|
||
if calib is None:
|
||
continue
|
||
|
||
gt_mask = batch_idx_np == si
|
||
gt_cls = cls_np[gt_mask]
|
||
gt_labels_3d = labels_3d_np[gt_mask]
|
||
|
||
pred_bboxes = pred["bboxes"]
|
||
pred_cls = pred["cls"]
|
||
if pred_bboxes.shape[0] == 0 or gt_cls.shape[0] == 0:
|
||
continue
|
||
|
||
p3d = self._preds_3d_filtered[si].cpu().numpy()
|
||
pedge = None
|
||
if self._preds_edge_filtered is not None:
|
||
pedge_tensor = self._preds_edge_filtered[si]
|
||
pedge = pedge_tensor.cpu().numpy() if pedge_tensor is not None else None
|
||
anchors_np = self._anchors_filtered[si].cpu().numpy()
|
||
strides_np = self._strides_filtered[si].cpu().numpy()
|
||
if p3d.shape[0] != pred_bboxes.shape[0]:
|
||
raise RuntimeError(
|
||
f"3D prediction alignment mismatch for sample {si}: {p3d.shape[0]} 3D rows vs {pred_bboxes.shape[0]} 2D preds"
|
||
)
|
||
|
||
from ultralytics.utils import ops
|
||
|
||
gt_bboxes = batch["bboxes"][gt_mask]
|
||
if gt_bboxes.shape[0] == 0:
|
||
continue
|
||
imgsz = batch["img"].shape[2:]
|
||
gt_bboxes_np = (ops.xywh2xyxy(gt_bboxes) * torch.tensor(imgsz, device=self.device)[[1, 0, 1, 0]]).cpu().numpy()
|
||
gt_bboxes_xyxy = ops.xywh2xyxy(gt_bboxes) * torch.tensor(imgsz, device=self.device)[[1, 0, 1, 0]]
|
||
|
||
iou = box_iou(gt_bboxes_xyxy, pred_bboxes)
|
||
correct_class = torch.tensor(gt_cls, device=self.device)[:, None] == pred_cls
|
||
iou = iou * correct_class
|
||
|
||
iou_np = iou.cpu().numpy()
|
||
matches = np.nonzero(iou_np >= 0.5)
|
||
matches = np.array(matches).T
|
||
if matches.shape[0] == 0:
|
||
continue
|
||
|
||
if matches.shape[0] > 1:
|
||
matches = matches[iou_np[matches[:, 0], matches[:, 1]].argsort()[::-1]]
|
||
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
|
||
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
|
||
|
||
whole_store = self._whole_metric_store()
|
||
face_pred = self._metric_store()
|
||
face_gt = self._metric_store()
|
||
visible_pred_yaw = []
|
||
visible_pred_edge_yaw = []
|
||
visible_gt_target_yaw = []
|
||
|
||
for gi, pi in matches:
|
||
cls_id = int(gt_cls[gi])
|
||
gt_bbox_xyxy = gt_bboxes_np[gi]
|
||
pred_bbox_xyxy = pred_bboxes[pi].cpu().numpy()
|
||
pred_whole = None
|
||
gt_whole = extract_3d_attrs_from_gt(
|
||
gt_labels_3d[gi],
|
||
cls_id,
|
||
calib,
|
||
img_w,
|
||
img_h,
|
||
self.face_3d_classes,
|
||
self.complete_3d_classes,
|
||
score_thr=face_visibility_score_thresh,
|
||
)
|
||
if gt_whole is not None:
|
||
pred_whole = extract_3d_attrs_from_prediction(
|
||
p3d[pi],
|
||
anchors_np[:, pi],
|
||
strides_np[pi],
|
||
calib,
|
||
pred_edge_60=pedge[pi] if pedge is not None else None,
|
||
)
|
||
self._append_metric_attr(whole_store["shape_gt"], gt_whole)
|
||
self._append_metric_attr(whole_store["shape_pred"], pred_whole)
|
||
|
||
if cls_id in self.complete_3d_classes or not is_gt_cut_object(gt_labels_3d[gi]):
|
||
self._append_metric_attr(whole_store["pos_gt"], gt_whole)
|
||
self._append_metric_attr(whole_store["pos_pred"], pred_whole)
|
||
|
||
if cls_id not in self.face_3d_classes:
|
||
continue
|
||
gt_visible_faces = select_gt_visible_faces(gt_labels_3d[gi], score_thr=face_visibility_score_thresh)
|
||
if gt_whole is not None and pred_whole is not None:
|
||
edge_selection = decode_edge_yaw_selection_from_prediction(
|
||
p3d[pi],
|
||
pedge[pi] if pedge is not None else None,
|
||
anchors_np[:, pi],
|
||
strides_np[pi],
|
||
calib,
|
||
score_thr=edge_yaw_compare_score_thresh,
|
||
bbox_xyxy=pred_bbox_xyxy,
|
||
img_w=img_w,
|
||
img_h=img_h,
|
||
max_lateral_dist_m=edge_yaw_compare_max_lateral_dist_m,
|
||
)
|
||
visible_pred_yaw.append(pred_whole["yaw"])
|
||
visible_pred_edge_yaw.append(edge_selection["yaw"] if edge_selection.get("is_valid") else float("nan"))
|
||
visible_gt_target_yaw.append(gt_whole["yaw"])
|
||
|
||
for face_type, _ in gt_visible_faces:
|
||
gt_face = extract_3d_attrs_from_gt(
|
||
gt_labels_3d[gi],
|
||
cls_id,
|
||
calib,
|
||
img_w,
|
||
img_h,
|
||
self.face_3d_classes,
|
||
self.complete_3d_classes,
|
||
face_type=face_type,
|
||
score_thr=face_visibility_score_thresh,
|
||
)
|
||
pred_face = extract_3d_attrs_from_prediction(
|
||
p3d[pi],
|
||
anchors_np[:, pi],
|
||
strides_np[pi],
|
||
calib,
|
||
face_type=face_type,
|
||
pred_edge_60=pedge[pi] if pedge is not None else None,
|
||
)
|
||
if gt_face is None or pred_face is None:
|
||
continue
|
||
self._append_metric_attr(face_gt, gt_face)
|
||
self._append_metric_attr(face_pred, pred_face)
|
||
|
||
self._aggregate_whole_metric_store(whole_store)
|
||
self._aggregate_face_metric_store(face_pred, face_gt, visible_pred_yaw, visible_pred_edge_yaw, visible_gt_target_yaw)
|
||
|
||
def get_stats(self):
|
||
"""Get 2D stats and merge grouped 3D metrics."""
|
||
from ultralytics.utils.metrics_3d import aggregate_3d_metric_groups
|
||
|
||
stats = super().get_stats()
|
||
metrics_3d = aggregate_3d_metric_groups(self.stats_3d)
|
||
self.metrics_3d_results = metrics_3d
|
||
|
||
whole = metrics_3d["whole"]
|
||
for key, value in whole.items():
|
||
stats[f"metrics_3d/{key}"] = value
|
||
|
||
face = metrics_3d["face"]
|
||
for key, value in face.items():
|
||
stats[f"metrics_3d_face/{key}"] = value
|
||
|
||
self.stats_3d = {"whole": [], "face": []} # Reset for next epoch
|
||
return stats
|
||
|
||
def print_results(self):
|
||
"""Print 3D physical metrics first, followed by the 2D detection summary."""
|
||
metrics_3d = self.metrics_3d_results
|
||
whole = metrics_3d["whole"]
|
||
face = metrics_3d["face"]
|
||
trainer = getattr(self, "trainer", None)
|
||
loss_ref_names = getattr(trainer, "loss_names", None) or (
|
||
"box",
|
||
"cls",
|
||
"dfl",
|
||
"z3d",
|
||
"uv",
|
||
"size",
|
||
"ycls",
|
||
"ydeg",
|
||
"ccls",
|
||
"fz",
|
||
"fuv",
|
||
"fsize",
|
||
"fcls",
|
||
"euv",
|
||
"ez",
|
||
)
|
||
leadw, colw = 7, 9
|
||
|
||
def _fmt_3d(value) -> str:
|
||
"""Format 3D metrics, showing '-' when the metric is not valid."""
|
||
return "-" if not np.isfinite(value) else f"{value:.3g}"
|
||
|
||
row_3d_values = [
|
||
"",
|
||
"all-3d",
|
||
f"{whole['matched']}",
|
||
f"{face['matched']}",
|
||
"-",
|
||
_fmt_3d(whole["depth_abs"]),
|
||
_fmt_3d(whole["uv"]),
|
||
_fmt_3d(whole["size"]),
|
||
"-",
|
||
_fmt_3d(whole["orient"]),
|
||
"-",
|
||
_fmt_3d(face["depth_abs"]),
|
||
_fmt_3d(face["uv"]),
|
||
_fmt_3d(face["size"]),
|
||
"-",
|
||
_fmt_3d(face["edge_orient_visible"]),
|
||
]
|
||
row_3d_values.extend(["-"] * (1 + len(loss_ref_names) - len(row_3d_values)))
|
||
row_3d = (f"%{leadw}s" + f"%{colw}s" * len(loss_ref_names)) % tuple(
|
||
row_3d_values[: 1 + len(loss_ref_names)]
|
||
)
|
||
header_2d = (f"%{leadw}s" + f"%{colw}s" * 7) % ("", "Class", "Imgs", "Inst", "P", "R", "m50", "m50-95")
|
||
row_2d = (f"%{leadw}s" + f"%{colw}s" + f"%{colw}i" * 2 + f"%{colw}.3g" * len(self.metrics.keys)) % (
|
||
"",
|
||
"all-2d",
|
||
self.seen,
|
||
self.metrics.nt_per_class.sum(),
|
||
*self.metrics.mean_results(),
|
||
)
|
||
LOGGER.info(row_3d)
|
||
LOGGER.info(header_2d)
|
||
LOGGER.info(row_2d)
|
||
if self.metrics.nt_per_class.sum() == 0:
|
||
LOGGER.warning(f"no labels found in {self.args.task} set, cannot compute metrics without labels")
|
||
|
||
if self.args.verbose and not self.training and self.nc > 1 and len(self.metrics.stats):
|
||
pf_class = "%22s" + "%11i" * 2 + "%11.3g" * len(self.metrics.keys) + "%11s" * 8
|
||
blanks = ("", "", "", "", "", "", "", "")
|
||
for i, c in enumerate(self.metrics.ap_class_index):
|
||
LOGGER.info(
|
||
pf_class
|
||
% (
|
||
self.names[c],
|
||
self.metrics.nt_per_image[c],
|
||
self.metrics.nt_per_class[c],
|
||
*self.metrics.class_result(i),
|
||
*blanks,
|
||
)
|
||
)
|