Reference for ultralytics/models/yolo/pose/val.py

def eval_json(self, stats):
    """Evaluate object detection model using COCO JSON format."""
    if self.args.save_json and self.is_coco and len(self.jdict):
        anno_json = self.data["path"] / "annotations/person_keypoints_val2017.json"  # annotations
        pred_json = self.save_dir / "predictions.json"  # predictions
        LOGGER.info(f"\nEvaluating pycocotools mAP using {pred_json} and {anno_json}...")
        try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
            check_requirements("pycocotools>=2.0.6")
            from pycocotools.coco import COCO  # noqa
            from pycocotools.cocoeval import COCOeval  # noqa

            for x in anno_json, pred_json:
                assert x.is_file(), f"{x} file not found"
            anno = COCO(str(anno_json))  # init annotations api
            pred = anno.loadRes(str(pred_json))  # init predictions api (must pass string, not Path)
            for i, eval in enumerate([COCOeval(anno, pred, "bbox"), COCOeval(anno, pred, "keypoints")]):
                if self.is_coco:
                    eval.params.imgIds = [int(Path(x).stem) for x in self.dataloader.dataset.im_files]  # im to eval
                eval.evaluate()
                eval.accumulate()
                eval.summarize()
                idx = i * 4 + 2
                stats[self.metrics.keys[idx + 1]], stats[self.metrics.keys[idx]] = eval.stats[
                    :2
                ]  # update mAP50-95 and mAP50
        except Exception as e:
            LOGGER.warning(f"pycocotools unable to run: {e}")
    return stats

def get_desc(self):
    """Return description of evaluation metrics in string format."""
    return ("%22s" + "%11s" * 10) % (
        "Class",
        "Images",
        "Instances",
        "Box(P",
        "R",
        "mAP50",
        "mAP50-95)",
        "Pose(P",
        "R",
        "mAP50",
        "mAP50-95)",
    )

def init_metrics(self, model):
    """Initialize pose estimation metrics for YOLO model."""
    super().init_metrics(model)
    self.kpt_shape = self.data["kpt_shape"]
    is_pose = self.kpt_shape == [17, 3]
    nkpt = self.kpt_shape[0]
    self.sigma = OKS_SIGMA if is_pose else np.ones(nkpt) / nkpt
    self.stats = dict(tp_p=[], tp=[], conf=[], pred_cls=[], target_cls=[], target_img=[])

def plot_predictions(self, batch, preds, ni):
    """
    Plot and save model predictions with bounding boxes and keypoints.

    Args:
        batch (dict): Dictionary containing batch data including images, file paths, and other metadata.
        preds (List[torch.Tensor]): List of prediction tensors from the model, each containing bounding boxes,
            confidence scores, class predictions, and keypoints.
        ni (int): Batch index used for naming the output file.

    The function extracts keypoints from predictions, converts predictions to target format, and plots them
    on the input images. The resulting visualization is saved to the specified save directory.
    """
    pred_kpts = torch.cat([p[:, 6:].view(-1, *self.kpt_shape) for p in preds], 0)
    plot_images(
        batch["img"],
        *output_to_target(preds, max_det=self.args.max_det),
        kpts=pred_kpts,
        paths=batch["im_file"],
        fname=self.save_dir / f"val_batch{ni}_pred.jpg",
        names=self.names,
        on_plot=self.on_plot,
    )  # pred

def plot_val_samples(self, batch, ni):
    """
    Plot and save validation set samples with ground truth bounding boxes and keypoints.

    Args:
        batch (dict): Dictionary containing batch data with keys:
            - img (torch.Tensor): Batch of images
            - batch_idx (torch.Tensor): Batch indices for each image
            - cls (torch.Tensor): Class labels
            - bboxes (torch.Tensor): Bounding box coordinates
            - keypoints (torch.Tensor): Keypoint coordinates
            - im_file (list): List of image file paths
        ni (int): Batch index used for naming the output file
    """
    plot_images(
        batch["img"],
        batch["batch_idx"],
        batch["cls"].squeeze(-1),
        batch["bboxes"],
        kpts=batch["keypoints"],
        paths=batch["im_file"],
        fname=self.save_dir / f"val_batch{ni}_labels.jpg",
        names=self.names,
        on_plot=self.on_plot,
    )

def pred_to_json(self, predn, filename):
    """
    Convert YOLO predictions to COCO JSON format.

    This method takes prediction tensors and a filename, converts the bounding boxes from YOLO format
    to COCO format, and appends the results to the internal JSON dictionary (self.jdict).

    Args:
        predn (torch.Tensor): Prediction tensor containing bounding boxes, confidence scores, class IDs,
            and keypoints, with shape (N, 6+K) where N is the number of predictions and K is the flattened
            keypoints dimension.
        filename (str | Path): Path to the image file for which predictions are being processed.

    Notes:
        The method extracts the image ID from the filename stem (either as an integer if numeric, or as a string),
        converts bounding boxes from xyxy to xywh format, and adjusts coordinates from center to top-left corner
        before saving to the JSON dictionary.
    """
    stem = Path(filename).stem
    image_id = int(stem) if stem.isnumeric() else stem
    box = ops.xyxy2xywh(predn[:, :4])  # xywh
    box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner
    for p, b in zip(predn.tolist(), box.tolist()):
        self.jdict.append(
            {
                "image_id": image_id,
                "category_id": self.class_map[int(p[5])],
                "bbox": [round(x, 3) for x in b],
                "keypoints": p[6:],
                "score": round(p[4], 5),
            }
        )

def preprocess(self, batch):
    """Preprocess batch by converting keypoints data to float and moving it to the device."""
    batch = super().preprocess(batch)
    batch["keypoints"] = batch["keypoints"].to(self.device).float()
    return batch

def save_one_txt(self, predn, pred_kpts, save_conf, shape, file):
    """
    Save YOLO pose detections to a text file in normalized coordinates.

    Args:
        predn (torch.Tensor): Prediction boxes and scores with shape (N, 6) for (x1, y1, x2, y2, conf, cls).
        pred_kpts (torch.Tensor): Predicted keypoints with shape (N, K, D) where K is the number of keypoints
            and D is the dimension (typically 3 for x, y, visibility).
        save_conf (bool): Whether to save confidence scores.
        shape (tuple): Original image shape (height, width).
        file (Path): Output file path to save detections.

    Notes:
        The output format is: class_id x_center y_center width height confidence keypoints where keypoints are
        normalized (x, y, visibility) values for each point.
    """
    from ultralytics.engine.results import Results

    Results(
        np.zeros((shape[0], shape[1]), dtype=np.uint8),
        path=None,
        names=self.names,
        boxes=predn[:, :6],
        keypoints=pred_kpts,
    ).save_txt(file, save_conf=save_conf)

def update_metrics(self, preds, batch):
    """
    Update metrics with new predictions and ground truth data.

    This method processes each prediction, compares it with ground truth, and updates various statistics
    for performance evaluation.

    Args:
        preds (List[torch.Tensor]): List of prediction tensors from the model.
        batch (dict): Batch data containing images and ground truth annotations.
    """
    for si, pred in enumerate(preds):
        self.seen += 1
        npr = len(pred)
        stat = dict(
            conf=torch.zeros(0, device=self.device),
            pred_cls=torch.zeros(0, device=self.device),
            tp=torch.zeros(npr, self.niou, dtype=torch.bool, device=self.device),
            tp_p=torch.zeros(npr, self.niou, dtype=torch.bool, device=self.device),
        )
        pbatch = self._prepare_batch(si, batch)
        cls, bbox = pbatch.pop("cls"), pbatch.pop("bbox")
        nl = len(cls)
        stat["target_cls"] = cls
        stat["target_img"] = cls.unique()
        if npr == 0:
            if nl:
                for k in self.stats.keys():
                    self.stats[k].append(stat[k])
                if self.args.plots:
                    self.confusion_matrix.process_batch(detections=None, gt_bboxes=bbox, gt_cls=cls)
            continue

        # Predictions
        if self.args.single_cls:
            pred[:, 5] = 0
        predn, pred_kpts = self._prepare_pred(pred, pbatch)
        stat["conf"] = predn[:, 4]
        stat["pred_cls"] = predn[:, 5]

        # Evaluate
        if nl:
            stat["tp"] = self._process_batch(predn, bbox, cls)
            stat["tp_p"] = self._process_batch(predn, bbox, cls, pred_kpts, pbatch["kpts"])
        if self.args.plots:
            self.confusion_matrix.process_batch(predn, bbox, cls)

        for k in self.stats.keys():
            self.stats[k].append(stat[k])

        # Save
        if self.args.save_json:
            self.pred_to_json(predn, batch["im_file"][si])
        if self.args.save_txt:
            self.save_one_txt(
                predn,
                pred_kpts,
                self.args.save_conf,
                pbatch["ori_shape"],
                self.save_dir / "labels" / f"{Path(batch['im_file'][si]).stem}.txt",
            )