Skip to content

PoseTrainer


Bases: v8.detect.DetectionTrainer

Source code in ultralytics/yolo/v8/pose/train.py
class PoseTrainer(v8.detect.DetectionTrainer):

    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
        """Initialize a PoseTrainer object with specified configurations and overrides."""
        if overrides is None:
            overrides = {}
        overrides['task'] = 'pose'
        super().__init__(cfg, overrides, _callbacks)

    def get_model(self, cfg=None, weights=None, verbose=True):
        """Get pose estimation model with specified configuration and weights."""
        model = PoseModel(cfg, ch=3, nc=self.data['nc'], data_kpt_shape=self.data['kpt_shape'], verbose=verbose)
        if weights:
            model.load(weights)

        return model

    def set_model_attributes(self):
        """Sets keypoints shape attribute of PoseModel."""
        super().set_model_attributes()
        self.model.kpt_shape = self.data['kpt_shape']

    def get_validator(self):
        """Returns an instance of the PoseValidator class for validation."""
        self.loss_names = 'box_loss', 'pose_loss', 'kobj_loss', 'cls_loss', 'dfl_loss'
        return v8.pose.PoseValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args))

    def plot_training_samples(self, batch, ni):
        """Plot a batch of training samples with annotated class labels, bounding boxes, and keypoints."""
        images = batch['img']
        kpts = batch['keypoints']
        cls = batch['cls'].squeeze(-1)
        bboxes = batch['bboxes']
        paths = batch['im_file']
        batch_idx = batch['batch_idx']
        plot_images(images,
                    batch_idx,
                    cls,
                    bboxes,
                    kpts=kpts,
                    paths=paths,
                    fname=self.save_dir / f'train_batch{ni}.jpg',
                    on_plot=self.on_plot)

    def plot_metrics(self):
        """Plots training/val metrics."""
        plot_results(file=self.csv, pose=True, on_plot=self.on_plot)  # save results.png

__init__(cfg=DEFAULT_CFG, overrides=None, _callbacks=None)

Initialize a PoseTrainer object with specified configurations and overrides.

Source code in ultralytics/yolo/v8/pose/train.py
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """Initialize a PoseTrainer object with specified configurations and overrides."""
    if overrides is None:
        overrides = {}
    overrides['task'] = 'pose'
    super().__init__(cfg, overrides, _callbacks)

get_model(cfg=None, weights=None, verbose=True)

Get pose estimation model with specified configuration and weights.

Source code in ultralytics/yolo/v8/pose/train.py
def get_model(self, cfg=None, weights=None, verbose=True):
    """Get pose estimation model with specified configuration and weights."""
    model = PoseModel(cfg, ch=3, nc=self.data['nc'], data_kpt_shape=self.data['kpt_shape'], verbose=verbose)
    if weights:
        model.load(weights)

    return model

get_validator()

Returns an instance of the PoseValidator class for validation.

Source code in ultralytics/yolo/v8/pose/train.py
def get_validator(self):
    """Returns an instance of the PoseValidator class for validation."""
    self.loss_names = 'box_loss', 'pose_loss', 'kobj_loss', 'cls_loss', 'dfl_loss'
    return v8.pose.PoseValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args))

plot_metrics()

Plots training/val metrics.

Source code in ultralytics/yolo/v8/pose/train.py
def plot_metrics(self):
    """Plots training/val metrics."""
    plot_results(file=self.csv, pose=True, on_plot=self.on_plot)  # save results.png

plot_training_samples(batch, ni)

Plot a batch of training samples with annotated class labels, bounding boxes, and keypoints.

Source code in ultralytics/yolo/v8/pose/train.py
def plot_training_samples(self, batch, ni):
    """Plot a batch of training samples with annotated class labels, bounding boxes, and keypoints."""
    images = batch['img']
    kpts = batch['keypoints']
    cls = batch['cls'].squeeze(-1)
    bboxes = batch['bboxes']
    paths = batch['im_file']
    batch_idx = batch['batch_idx']
    plot_images(images,
                batch_idx,
                cls,
                bboxes,
                kpts=kpts,
                paths=paths,
                fname=self.save_dir / f'train_batch{ni}.jpg',
                on_plot=self.on_plot)

set_model_attributes()

Sets keypoints shape attribute of PoseModel.

Source code in ultralytics/yolo/v8/pose/train.py
def set_model_attributes(self):
    """Sets keypoints shape attribute of PoseModel."""
    super().set_model_attributes()
    self.model.kpt_shape = self.data['kpt_shape']



train


Train the YOLO model on the given data and device.

Source code in ultralytics/yolo/v8/pose/train.py
def train(cfg=DEFAULT_CFG, use_python=False):
    """Train the YOLO model on the given data and device."""
    model = cfg.model or 'yolov8n-pose.yaml'
    data = cfg.data or 'coco8-pose.yaml'
    device = cfg.device if cfg.device is not None else ''

    args = dict(model=model, data=data, device=device)
    if use_python:
        from ultralytics import YOLO
        YOLO(model).train(**args)
    else:
        trainer = PoseTrainer(overrides=args)
        trainer.train()




Created 2023-04-16, Updated 2023-05-30
Authors: Glenn Jocher (4)