diff --git a/tensorbay/opendataset/LIP/__init__.py b/tensorbay/opendataset/LIP/__init__.py new file mode 100644 index 000000000..eafcfdc5c --- /dev/null +++ b/tensorbay/opendataset/LIP/__init__.py @@ -0,0 +1,11 @@ +#!/usr/bin/env python3 +# +# Copyright 2021 Graviti. Licensed under MIT License. +# +# pylint: disable=invalid-name + +"""Dataloader of the LIP dataset.""" + +from .loader import LIP + +__all__ = ["LIP"] diff --git a/tensorbay/opendataset/LIP/catalog.json b/tensorbay/opendataset/LIP/catalog.json new file mode 100644 index 000000000..6c3049593 --- /dev/null +++ b/tensorbay/opendataset/LIP/catalog.json @@ -0,0 +1,52 @@ +{ + "SEMANTIC_MASK": { + "categories": [ + { "name": "Background", "categoryId": 0 }, + { "name": "Hat", "categoryId": 1 }, + { "name": "Hair", "categoryId": 2 }, + { "name": "Glove", "categoryId": 3 }, + { "name": "Sunglasses", "categoryId": 4 }, + { "name": "UpperClothes", "categoryId": 5 }, + { "name": "Dress", "categoryId": 6 }, + { "name": "Coat", "categoryId": 7 }, + { "name": "Socks", "categoryId": 8 }, + { "name": "Pants", "categoryId": 9 }, + { "name": "Jumpsuits", "categoryId": 10 }, + { "name": "Scarf", "categoryId": 11 }, + { "name": "Skirt", "categoryId": 12 }, + { "name": "Face", "categoryId": 13 }, + { "name": "Left-arm", "categoryId": 14 }, + { "name": "Right-arm", "categoryId": 15 }, + { "name": "Left-leg", "categoryId": 16 }, + { "name": "Right-leg", "categoryId": 17 }, + { "name": "Left-shoe", "categoryId": 18 }, + { "name": "Right-shoe", "categoryId": 19 } + ] + }, + "KEYPOINTS2D": { + "keypoints": [ + { + "number": 16, + "names": [ + "R_Ankle", + "R_Knee", + "R_Hip", + "L_Hip", + "L_Knee", + "L_Ankle", + "B_Pelvis", + "B_Spine", + "B_Neck", + "B_Head", + "R_Wrist", + "R_Elbow", + "R_Shoulder", + "L_Shoulder", + "L_Elbow", + "L_Wrist" + ], + "visible": "BINARY" + } + ] + } +} diff --git a/tensorbay/opendataset/LIP/loader.py b/tensorbay/opendataset/LIP/loader.py new file mode 100644 index 000000000..8f1d5cc4f --- /dev/null +++ b/tensorbay/opendataset/LIP/loader.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +# +# Copyright 2021 Graviti. Licensed under MIT License. +# +# pylint: disable=invalid-name +# pylint: disable=missing-module-docstring + +import csv +import os +from itertools import islice +from typing import List + +from ...dataset import Data, Dataset +from ...geometry import Keypoint2D +from ...label import LabeledKeypoints2D, SemanticMask +from ...utility import chunked + +DATASET_NAME = "LIP" +_SEGMENT_NAMES = ("train", "val", "test") + + +def LIP(path: str) -> Dataset: + """Dataloader of the `LIP`_ dataset. + + .. _LIP: https://github.com/Engineering-Course/LIP_SSL + + The file structure should be like:: + + + Testing_images/ + testing_images/ + 315_462476.jpg + ... + test_id.txt + TrainVal_images/ + TrainVal_images/ + train_images/ + 77_471474.jpg + ... + val_images/ + 36_453991.jpg + ... + train_id.txt + val_id.txt + TrainVal_parsing_annotations/ + TrainVal_parsing_annotations/ + train_segmentations/ + 77_471474.png + ... + val_segmentations/ + 36_453991.png + ... + TrainVal_pose_annotations/ + lip_train_set.csv + lip_val_set.csv + + Arguments: + path: The root directory of the dataset. + + Returns: + Loaded `~tensorbay.dataset.dataset.Dataset` instance. + + """ + root_path = os.path.abspath(os.path.expanduser(path)) + test_path = os.path.join(root_path, "Testing_images") + trainval_image_path = os.path.join(root_path, "TrainVal_images", "TrainVal_images") + trainval_parsing_path = os.path.join( + root_path, "TrainVal_parsing_annotations", "TrainVal_parsing_annotations" + ) + pose_path = os.path.join(root_path, "TrainVal_pose_annotations") + + dataset = Dataset(DATASET_NAME) + dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json")) + + for segment_name in _SEGMENT_NAMES: + segment = dataset.create_segment(segment_name) + if segment_name == "test": + image_path = os.path.join(test_path, "testing_images") + with open(os.path.join(test_path, "test_id.txt"), "r") as fp: + for filename in fp: + segment.append(Data(os.path.join(image_path, f"{filename.rstrip()}.jpg"))) + else: + image_path = os.path.join(trainval_image_path, f"{segment_name}_images") + parsing_path = os.path.join(trainval_parsing_path, f"{segment_name}_segmentations") + with open(os.path.join(pose_path, f"lip_{segment_name}_set.csv"), "r") as csvfile: + for keypoints_info in csv.reader(csvfile): + segment.append(_get_data(keypoints_info, image_path, parsing_path)) + return dataset + + +def _get_data(keypoints_info: List[str], image_path: str, parsing_path: str) -> Data: + filename = os.path.splitext(keypoints_info[0])[0] + data = Data(os.path.join(image_path, f"{filename}.jpg")) + label = data.label + label.semantic_mask = SemanticMask(os.path.join(parsing_path, f"{filename}.png")) + keypoints = LabeledKeypoints2D() + for x, y, v in chunked(islice(keypoints_info, 1, None), 3): + keypoints.append( + Keypoint2D(float(x), float(y), int(v)) if x.isnumeric() else Keypoint2D(0, 0, 0) + ) + label.keypoints2d = [keypoints] + return data diff --git a/tensorbay/opendataset/__init__.py b/tensorbay/opendataset/__init__.py index bd96b0ff6..7fec4b0ad 100644 --- a/tensorbay/opendataset/__init__.py +++ b/tensorbay/opendataset/__init__.py @@ -87,6 +87,7 @@ "KenyanFoodOrNonfood", "KenyanFoodType", "KylbergTexture", + "LIP", "LISATrafficLight", "LISATrafficSign", "LeedsSportsPose",