To add a new dataset for pose estimation, subclass BaseKeypointsDataset. You must implement the __len__ method and the load_sample method. The load_sample method must return a tuple of (image, mask, joints, extras) with the following specifications:
image: Numpy array of [H, W, 3] (RGB).mask: Numpy array of [H, W] (binary mask where zero values indicate ignored regions).joints: Numpy array of [Num Instances, Num Joints, 3] representing skeletons.extras: A dictionary for additional sample information.
from super_gradients.training.datasets.pose_estimation_datasets import BaseKeypointsDataset
from super_gradients.training.datasets.pose_estimation_datasets import KeypointsTargetsGenerator
from super_gradients.training.transforms.keypoint_transforms import KeypointTransform
from typing import Tuple, Dict, Any, List
import numpy as np
import cv2
class MyNewPoseEstimationDataset(BaseKeypointsDataset):
def __init__(
self,
image_paths,
joint_paths,
target_generator: KeypointsTargetsGenerator,
transforms: List[KeypointTransform],
min_instance_area: float = 0.0,
):
super().__init__(target_generator, transforms, min_instance_area)
self.image_paths = image_paths
self.joint_paths = joint_paths
def __len__(self) -> int:
return len(self.image_paths)
def load_sample(self, index) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, Any]]:
# Read image from the disk
image = cv2.imread(self.image_paths[index])
mask = np.ones(image.shape[:2])
joints = np.loadtxt(self.joint_paths[index])
return image, mask, joints, {}
