# 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,
                    )
                )