From d51b157ed5d4ad6b17cab5572ae7bf1cb3e09324 Mon Sep 17 00:00:00 2001 From: "xin.wang" Date: Mon, 20 Dec 2021 11:28:05 +0800 Subject: [PATCH] feat(opendataset): add dataloader for the CityscapesGTCoarse dataset PR Closed: https://github.com/Graviti-AI/tensorbay-python-sdk/pull/1167 --- docs/source/reference/api/opendataset.rst | 1 + tensorbay/opendataset/Cityscapes/__init__.py | 6 +- tensorbay/opendataset/Cityscapes/loader.py | 132 +++++++++++++++---- tensorbay/opendataset/__init__.py | 3 +- 4 files changed, 109 insertions(+), 33 deletions(-) diff --git a/docs/source/reference/api/opendataset.rst b/docs/source/reference/api/opendataset.rst index 92648a35e..6b93c0a1e 100644 --- a/docs/source/reference/api/opendataset.rst +++ b/docs/source/reference/api/opendataset.rst @@ -24,6 +24,7 @@ tensorbay.opendataset CCPD CCPDGreen CIHP + CityscapesGTCoarse CityscapesGTFine COCO2017 COVIDChestXRay diff --git a/tensorbay/opendataset/Cityscapes/__init__.py b/tensorbay/opendataset/Cityscapes/__init__.py index df92e6cf9..a081c7cf7 100644 --- a/tensorbay/opendataset/Cityscapes/__init__.py +++ b/tensorbay/opendataset/Cityscapes/__init__.py @@ -4,8 +4,8 @@ # # pylint: disable=invalid-name -"""Dataloader of the CityscapesGTFine dataset.""" +"""Dataloaders of the CityscapesGTCoarse dataset and the CityscapesGTFine dataset.""" -from tensorbay.opendataset.Cityscapes.loader import CityscapesGTFine +from tensorbay.opendataset.Cityscapes.loader import CityscapesGTCoarse, CityscapesGTFine -__all__ = ["CityscapesGTFine"] +__all__ = ["CityscapesGTCoarse", "CityscapesGTFine"] diff --git a/tensorbay/opendataset/Cityscapes/loader.py b/tensorbay/opendataset/Cityscapes/loader.py index b15ee983c..be56267f9 100644 --- a/tensorbay/opendataset/Cityscapes/loader.py +++ b/tensorbay/opendataset/Cityscapes/loader.py @@ -4,7 +4,7 @@ # # pylint: disable=invalid-name -"""Dataloader of the CityscapesGTFine dataset.""" +"""Dataloaders of the CityscapesGTCoarse dataset and the CityscapesGTFine dataset.""" import json import os @@ -14,8 +14,77 @@ from tensorbay.dataset import Data, Dataset from tensorbay.label import InstanceMask, LabeledPolygon, SemanticMask -DATASET_NAME = "CityscapesGTFine" -_SEGMENT_NAMES = {"test", "train", "val"} +DATASET_NAME_GTCOARSE = "CityscapesGTCoarse" +DATASET_NAME_GTFINE = "CityscapesGTFine" + +_SEGMENT_NAMES_GTCOARSE = ("train", "train_extra", "val") +_SEGMENT_NAMES_GTFINE = ("train", "test", "val") + + +def CityscapesGTCoarse(path: str) -> Dataset: + """`CityscapesGTCoarse `_ dataset. + + The file structure should be like:: + + + leftImg8bit/ + train/ + aachen/ + aachen_000000_000019_leftImg8bit.png + ... + ... + train_extra/ + augsburg/ + augsburg_000000_000019_leftImg8bit.png + ... + ... + val/ + frankfurt/ + frankfurt_000000_000019_leftImg8bit.png + ... + ... + ... + gtCoarse/ + train/ + aachen/ + aachen_000000_000019_gtCoarse_instanceIds.png + aachen_000000_000019_gtCoarse_labelIds.png + aachen_000000_000019_gtCoarse_polygons.json + ... + ... + train_extra/ + augsburg/ + augsburg_000000_000019_gtCoarse_instanceIds.png + augsburg_000000_000019_gtCoarse_labelIds.png + augsburg_000000_000019_gtCoarse_polygons.json + ... + ... + val/ + frankfurt/ + frankfurt_000000_000019_gtCoarse_instanceIds.png + frankfurt_000000_000019_gtCoarse_labelIds.png + frankfurt_000000_000019_gtCoarse_polygons.json + ... + ... + ... + + Arguments: + path: The root directory of the dataset. + + Returns: + Loaded :class:`~tensorbay.dataset.dataset.Dataset` instance. + + """ + root_path = os.path.join(os.path.abspath(os.path.expanduser(path))) + + dataset = Dataset(DATASET_NAME_GTCOARSE) + dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json")) + + for segment_name in _SEGMENT_NAMES_GTCOARSE: + segment = dataset.create_segment(segment_name) + for image_path in glob(os.path.join(root_path, "leftImg8bit", segment_name, "*", "*.png")): + segment.append(_get_data(image_path, root_path, segment_name, "gtCoarse")) + return dataset def CityscapesGTFine(path: str) -> Dataset: @@ -74,34 +143,39 @@ def CityscapesGTFine(path: str) -> Dataset: """ root_path = os.path.join(os.path.abspath(os.path.expanduser(path))) - dataset = Dataset(DATASET_NAME) + dataset = Dataset(DATASET_NAME_GTFINE) dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json")) - for segment_name in _SEGMENT_NAMES: + for segment_name in _SEGMENT_NAMES_GTFINE: segment = dataset.create_segment(segment_name) for image_path in glob(os.path.join(root_path, "leftImg8bit", segment_name, "*", "*.png")): - city = os.path.basename(image_path).split("_", 1)[0] - image_prefix = os.path.basename(image_path).rsplit("_", 1)[0] - label_dir = os.path.join(root_path, "gtFine", segment_name, city) - data = Data(image_path) - # get semantic mask and instance mask - label = data.label - label.semantic_mask = SemanticMask( - os.path.join(label_dir, f"{image_prefix}_gtFine_labelIds.png") - ) - label.instance_mask = InstanceMask( - os.path.join(label_dir, f"{image_prefix}_gtFine_instanceIds.png") - ) - # get polygons - polygons: List[LabeledPolygon] = [] - with open( - os.path.join(label_dir, f"{image_prefix}_gtFine_polygons.json"), - encoding="utf-8", - ) as fp: - objects = json.load(fp)["objects"] - for obj in objects: - polygons.append(LabeledPolygon(obj["polygon"], category=obj["label"])) - label.polygon = polygons - - segment.append(data) + segment.append(_get_data(image_path, root_path, segment_name, "gtFine")) return dataset + + +def _get_data(image_path: str, root_path: str, segment_name: str, folder_name: str) -> Data: + filename = os.path.basename(image_path) + city = filename.split("_", 1)[0] + image_prefix = filename.rsplit("_", 1)[0] + label_dir = os.path.join(root_path, folder_name, segment_name, city) + data = Data(image_path) + # get semantic mask and instance mask + label = data.label + label.semantic_mask = SemanticMask( + os.path.join(label_dir, f"{image_prefix}_{folder_name}_labelIds.png") + ) + label.instance_mask = InstanceMask( + os.path.join(label_dir, f"{image_prefix}_{folder_name}_instanceIds.png") + ) + # get polygons + polygons: List[LabeledPolygon] = [] + with open( + os.path.join(label_dir, f"{image_prefix}_{folder_name}_polygons.json"), + encoding="utf-8", + ) as fp: + objects = json.load(fp)["objects"] + for obj in objects: + polygons.append(LabeledPolygon(obj["polygon"], category=obj["label"])) + label.polygon = polygons + + return data diff --git a/tensorbay/opendataset/__init__.py b/tensorbay/opendataset/__init__.py index 8de67f7d5..02527f83c 100644 --- a/tensorbay/opendataset/__init__.py +++ b/tensorbay/opendataset/__init__.py @@ -17,7 +17,7 @@ from tensorbay.opendataset.CarConnection import CarConnection from tensorbay.opendataset.CCPD import CCPD, CCPDGreen from tensorbay.opendataset.CIHP import CIHP -from tensorbay.opendataset.Cityscapes import CityscapesGTFine +from tensorbay.opendataset.Cityscapes import CityscapesGTCoarse, CityscapesGTFine from tensorbay.opendataset.COCO2017 import COCO2017 from tensorbay.opendataset.CoinImage import CoinImage from tensorbay.opendataset.CompCars import CompCars @@ -81,6 +81,7 @@ "CCPD", "CCPDGreen", "CIHP", + "CityscapesGTCoarse", "CityscapesGTFine", "COCO2017", "CoinImage",