Source code for tensorbay.opendataset.LeedsSportsPose.loader

#!/usr/bin/env python3
#
# Copyright 2021 Graviti. Licensed under MIT License.
#
# pylint: disable=invalid-name

"""Dataloader of LeedsSportsPose dataset."""

import os

from tensorbay.dataset import Data, Dataset
from tensorbay.exception import ModuleImportError
from tensorbay.geometry import Keypoint2D
from tensorbay.label import LabeledKeypoints2D
from tensorbay.opendataset._utility import glob

DATASET_NAME = "LeedsSportsPose"


[docs]def LeedsSportsPose(path: str) -> Dataset: """`Leeds Sports Pose <http://sam.johnson.io/research/lsp.html>`_ dataset. The folder structure should be like:: <path> joints.mat images/ im0001.jpg im0002.jpg ... Arguments: path: The root directory of the dataset. Raises: ModuleImportError: When the module "scipy" can not be found. Returns: Loaded :class:`~tensorbay.dataset.dataset.Dataset` instance. """ try: from scipy.io import loadmat # pylint: disable=import-outside-toplevel except ModuleNotFoundError as error: raise ModuleImportError(module_name=error.name) from error root_path = os.path.abspath(os.path.expanduser(path)) dataset = Dataset(DATASET_NAME) dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json")) segment = dataset.create_segment() mat = loadmat(os.path.join(root_path, "joints.mat")) joints = mat["joints"].T image_paths = glob(os.path.join(root_path, "images", "*.jpg")) for image_path in image_paths: data = Data(image_path) data.label.keypoints2d = [] index = int(os.path.basename(image_path)[2:6]) - 1 # get image index from "im0001.jpg" keypoints = LabeledKeypoints2D() for keypoint in joints[index]: keypoints.append(Keypoint2D(keypoint[0], keypoint[1], int(not keypoint[2]))) data.label.keypoints2d.append(keypoints) segment.append(data) return dataset