yiqun/PascalPart
This PACO dataset is designed to load coco-stuff only & coco stuff thing.
027
1# Copyright 2020 The HuggingFace Datasets Authors and the current dataset2# script contributor.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15# TODO: Address all TODOs and remove all explanatory comments16"""TODO: Add a description here."""17 18 19import os20import copy21import logging22import os.path as osp23from scipy.io import loadmat24from PIL import Image25 26import datasets27from datasets import Value28 29logger = logging.getLogger(__name__)30 31 32def build_class_to_parts_dict():33 pimap = {}34 35 # [aeroplane]36 pimap[1] = {}37 pimap[1]["body"] = 138 pimap[1]["stern"] = 239 pimap[1]["lwing"] = 3 # left wing40 pimap[1]["rwing"] = 4 # right wing41 pimap[1]["tail"] = 542 for ii in range(1, 10 + 1):43 pimap[1][('engine_%d' % ii)] = 10+ii # multiple engines44 for ii in range(1, 10 + 1):45 pimap[1][('wheel_%d' % ii)] = 20+ii # multiple wheels46 47 # [bicycle]48 pimap[2] = {}49 pimap[2]['fwheel'] = 1 # front wheel50 pimap[2]['bwheel'] = 2 # back wheel51 pimap[2]['saddle'] = 352 pimap[2]['handlebar'] = 4 # handle bar53 pimap[2]['chainwheel'] = 5 # chain wheel54 for ii in range(1, 10 + 1):55 pimap[2][('headlight_%d' % ii)] = 10 + ii56 57 # [bird]58 pimap[3] = {}59 pimap[3]['head'] = 160 pimap[3]['leye'] = 2 # left eye61 pimap[3]['reye'] = 3 # right eye62 pimap[3]['beak'] = 463 pimap[3]['torso'] = 564 pimap[3]['neck'] = 665 pimap[3]['lwing'] = 7 # left wing66 pimap[3]['rwing'] = 8 # right wing67 pimap[3]['lleg'] = 9 # left leg68 pimap[3]['lfoot'] = 10 # left foot69 pimap[3]['rleg'] = 11 # right leg70 pimap[3]['rfoot'] = 12 # right foot71 pimap[3]['tail'] = 1372 73 # [boat]74 # only has silhouette mask75 76 # [bottle]77 pimap[5] = {}78 pimap[5]['cap'] = 179 pimap[5]['body'] = 280 81 # [bus]82 pimap[6] = {}83 pimap[6]['frontside'] = 184 pimap[6]['leftside'] = 285 pimap[6]['rightside'] = 386 pimap[6]['backside'] = 487 pimap[6]['roofside'] = 588 pimap[6]['leftmirror'] = 689 pimap[6]['rightmirror'] = 790 pimap[6]['fliplate'] = 8 # front license plate91 pimap[6]['bliplate'] = 9 # back license plate92 for ii in range(1, 10 + 1):93 pimap[6][('door_%d' % ii)] = 10 + ii94 for ii in range(1, 10 + 1):95 pimap[6][('wheel_%d' % ii)] = 20 + ii96 for ii in range(1, 10 + 1):97 pimap[6][('headlight_%d' % ii)] = 30 + ii98 for ii in range(1, 20 + 1):99 pimap[6][('window_%d' % ii)] = 40 + ii100 101 # [car]102 pimap[7] = pimap[6].copy() # car has the same set of parts with bus103 104 # [cat]105 pimap[8] = {}106 pimap[8]['head'] = 1107 pimap[8]['leye'] = 2 # left eye108 pimap[8]['reye'] = 3 # right eye109 pimap[8]['lear'] = 4 # left ear110 pimap[8]['rear'] = 5 # right ear111 pimap[8]['nose'] = 6112 pimap[8]['torso'] = 7113 pimap[8]['neck'] = 8114 pimap[8]['lfleg'] = 9 # left front leg115 pimap[8]['lfpa'] = 10 # left front paw116 pimap[8]['rfleg'] = 11 # right front leg117 pimap[8]['rfpa'] = 12 # right front paw118 pimap[8]['lbleg'] = 13 # left back leg119 pimap[8]['lbpa'] = 14 # left back paw120 pimap[8]['rbleg'] = 15 # right back leg121 pimap[8]['rbpa'] = 16 # right back paw122 pimap[8]['tail'] = 17123 124 # [chair]125 # only has sihouette mask126 127 # [cow]128 pimap[10] = {}129 pimap[10]['head'] = 1130 pimap[10]['leye'] = 2 # left eye131 pimap[10]['reye'] = 3 # right eye132 pimap[10]['lear'] = 4 # left ear133 pimap[10]['rear'] = 5 # right ear134 pimap[10]['muzzle'] = 6135 pimap[10]['lhorn'] = 7 # left horn136 pimap[10]['rhorn'] = 8 # right horn137 pimap[10]['torso'] = 9138 pimap[10]['neck'] = 10139 pimap[10]['lfuleg'] = 11 # left front upper leg140 pimap[10]['lflleg'] = 12 # left front lower leg141 pimap[10]['rfuleg'] = 13 # right front upper leg142 pimap[10]['rflleg'] = 14 # right front lower leg143 pimap[10]['lbuleg'] = 15 # left back upper leg144 pimap[10]['lblleg'] = 16 # left back lower leg145 pimap[10]['rbuleg'] = 17 # right back upper leg146 pimap[10]['rblleg'] = 18 # right back lower leg147 pimap[10]['tail'] = 19148 149 # [table]150 # only has silhouette mask151 152 # [dog]153 # dog has the same set of parts with cat,154 pimap[12] = pimap[8].copy()155 # except for the additional156 # muzzle157 pimap[12]['muzzle'] = 20158 159 # [horse]160 # horse has the same set of parts with cow,161 pimap[13] = pimap[10].copy()162 # except it has hoof instead of horn163 del pimap[13]['lhorn']164 del pimap[13]['rhorn']165 pimap[13]['lfho'] = 30166 pimap[13]['rfho'] = 31167 pimap[13]['lbho'] = 32168 pimap[13]['rbho'] = 33169 170 # [motorbike]171 pimap[14] = {}172 pimap[14]['fwheel'] = 1173 pimap[14]['bwheel'] = 2174 pimap[14]['handlebar'] = 3175 pimap[14]['saddle'] = 4176 for ii in range(1, 10 + 1):177 pimap[14][('headlight_%d' % ii)] = 10 + ii178 179 # [person]180 pimap[15] = {}181 pimap[15]['head'] = 1182 pimap[15]['leye'] = 2 # left eye183 pimap[15]['reye'] = 3 # right eye184 pimap[15]['lear'] = 4 # left ear185 pimap[15]['rear'] = 5 # right ear186 pimap[15]['lebrow'] = 6 # left eyebrow187 pimap[15]['rebrow'] = 7 # right eyebrow188 pimap[15]['nose'] = 8189 pimap[15]['mouth'] = 9190 pimap[15]['hair'] = 10191 192 pimap[15]['torso'] = 11193 pimap[15]['neck'] = 12194 pimap[15]['llarm'] = 13 # left lower arm195 pimap[15]['luarm'] = 14 # left upper arm196 pimap[15]['lhand'] = 15 # left hand197 pimap[15]['rlarm'] = 16 # right lower arm198 pimap[15]['ruarm'] = 17 # right upper arm199 pimap[15]['rhand'] = 18 # right hand200 201 pimap[15]['llleg'] = 19 # left lower leg202 pimap[15]['luleg'] = 20 # left upper leg203 pimap[15]['lfoot'] = 21 # left foot204 pimap[15]['rlleg'] = 22 # right lower leg205 pimap[15]['ruleg'] = 23 # right upper leg206 pimap[15]['rfoot'] = 24 # right foot207 208 # [pottedplant]209 pimap[16] = {}210 pimap[16]['pot'] = 1211 pimap[16]['plant'] = 2212 213 # [sheep]214 # sheep has the same set of parts with cow215 pimap[17] = pimap[10].copy()216 217 # [sofa]218 # only has sihouette mask219 220 # [train]221 pimap[19] = {}222 pimap[19]['head'] = 1223 pimap[19]['hfrontside'] = 2 # head front side224 pimap[19]['hleftside'] = 3 # head left side225 pimap[19]['hrightside'] = 4 # head right side226 pimap[19]['hbackside'] = 5 # head back side227 pimap[19]['hroofside'] = 6 # head roof side228 229 for ii in range(1, 10 + 1):230 pimap[19][('headlight_%d' % ii)] = 10 + ii231 232 for ii in range(1, 10 + 1):233 pimap[19][('coach_%d' % ii)] = 20 + ii234 235 for ii in range(1, 10 + 1):236 pimap[19][('cfrontside_%d' % ii)] = 30 + ii # coach front side237 238 for ii in range(1, 10 + 1):239 pimap[19][('cleftside_%d' % ii)] = 40 + ii # coach left side240 241 for ii in range(1, 10 + 1):242 pimap[19][('crightside_%d' % ii)] = 50 + ii # coach right side243 244 for ii in range(1, 10 + 1):245 pimap[19][('cbackside_%d' % ii)] = 60 + ii # coach back side246 247 for ii in range(1, 10 + 1):248 pimap[19][('croofside_%d' % ii)] = 70 + ii # coach roof side249 250 # [tvmonitor]251 pimap[20] = {}252 pimap[20]['screen'] = 1253 254 return pimap255 256 257# TODO: Add BibTeX citation258# Find for instance the citation on arxiv or on the dataset repo/website259_CITATION = """\260@InProceedings{huggingface:dataset,261title = {A great new dataset},262author={huggingface, Inc.263},264year={2020}265}266"""267 268# TODO: Add description of the dataset here269# You can copy an official description270_DESCRIPTION = """\271This PACO dataset is designed to load coco-stuff only & coco stuff thing.272"""273 274# TODO: Add a link to an official homepage for the dataset here275_HOMEPAGE = ""276 277# TODO: Add the licence for the dataset here if you can find it278_LICENSE = ""279 280# TODO: Add link to the official dataset URLs here281# The HuggingFace Datasets library doesn't host the datasets but only points282# to the original files.283# This can be an arbitrary nested dict/list of URLs284# (see below in `_split_generators` method)285_URLS = {}286 287 288VALID_SPLIT_NAMES = ("trainval",)289 290 291CATEGORY_ID_TO_CATEGORY = {292 1: 'aeroplane', 2: 'bicycle', 3: 'bird', 4: 'boat', 5: 'bottle',293 6: 'bus', 7: 'car', 8: 'cat', 9: 'chair', 10: 'cow', 11: 'table',294 12: 'dog', 13: 'horse', 14: 'motorbike', 15: 'person', 16: 'pottedplant',295 17: 'sheep', 18: 'sofa', 19: 'train', 20: 'tvmonitor',296}297CLASS_TO_PARTS = build_class_to_parts_dict()298 299 300# TODO: Name of the dataset usually matches the script name with CamelCase301# instead of snake_case302class PascalPartDataset(datasets.GeneratorBasedBuilder):303 """TODO: Short description of my dataset."""304 305 VERSION = datasets.Version("0.0.1")306 307 # This is an example of a dataset with multiple configurations.308 # If you don't want/need to define several sub-sets in your dataset,309 # just remove the BUILDER_CONFIG_CLASS and the BUILDER_CONFIGS attributes.310 311 # If you need to make complex sub-parts in the datasets with configurable312 # options313 # You can create your own builder configuration class to store attribute,314 # inheriting from datasets.BuilderConfig315 BUILDER_CONFIG_CLASS = datasets.BuilderConfig316 317 # You will be able to load one or the other configurations318 # in the following list with319 # data = datasets.load_dataset('my_dataset', 'first_domain')320 # data = datasets.load_dataset('my_dataset', 'second_domain')321 BUILDER_CONFIGS = [322 datasets.BuilderConfig(323 version=VERSION, description="paco-lvis."),324 ]325 326 # It's not mandatory to have a default configuration.327 # Just use one if it make sense.328 329 def _info(self):330 self.config: datasets.BuilderConfig331 mask_type = datasets.Image()332 features = {333 "image": {334 "file_name": Value("string"),335 "image_path": Value("string"),336 "width": Value("int32"),337 "height": Value("int32"),338 },339 "annotations": [340 {341 # category id in this object.342 'category_id': Value("int32"),343 'category': Value("string"),344 'segmentation': mask_type,345 'parts': [346 {347 "category_id": Value("int32"),348 "category": Value("string"),349 "segmentation": mask_type,350 }351 ],352 }353 ],354 }355 features = datasets.Features(features)356 return datasets.DatasetInfo(357 # This is the description that will appear on the datasets page.358 description=_DESCRIPTION,359 # This defines the different columns of the dataset and their types360 features=features,361 # If there's a common (input, target) tuple from the features,362 # uncomment supervised_keys line below and specify them.363 # They'll be used if as_supervised=True in builder.as_dataset.364 # supervised_keys=("sentence", "label"),365 # Homepage of the dataset for documentation366 homepage=_HOMEPAGE,367 # License for the dataset if available368 license=_LICENSE,369 # Citation for the dataset370 citation=_CITATION,371 )372 373 def _split_generators(self, dl_manager):374 # TODO: This method is tasked with downloading/extracting the data and375 # defining the splits depending on the configuration376 # If several configurations are possible (listed in BUILDER_CONFIGS),377 # the configuration selected by the user is in self.config.name378 379 # dl_manager is a datasets.download.DownloadManager that can be used380 # to download and extract URLS. It can accept any type381 # or nested list/dict and will give back the same structure with382 # the url replaced with path to local files.383 # By default the archives will be extracted and a path to a cached384 # folder where they are extracted is returned instead of the archive385 # urls = _URLS[self.config.name]386 # data_dir = dl_manager.download_and_extract(urls)387 splits = []388 split_names = ("trainval",)389 for split in split_names:390 splits.append(datasets.SplitGenerator(391 name=datasets.NamedSplit(split),392 gen_kwargs={"split": split},393 ))394 return splits395 396 # method parameters are unpacked from `gen_kwargs` as given in397 # `_split_generators`398 def _generate_examples(self, split: str):399 # TODO: This method handles input defined in _split_generators to400 # yield (key, example) tuples from the dataset.401 # The `key` is for legacy reasons (tfds) and is not important402 # in itself, but must be unique for each example.403 data_root = self.config.data_dir404 path = osp.expanduser(data_root)405 ann_dir = osp.join(path, "Annotations_Part")406 img_dir = osp.join(path, "VOCdevkit", "VOC2010", "JPEGImages")407 files = os.listdir(ann_dir)408 anno_path = [osp.join(ann_dir, f) for f in files]409 410 for path in anno_path:411 data = self._load_from_mat_data(path)412 file_name = data["image"]["file_name"]413 file_path = osp.join(img_dir, f"{file_name}.jpg")414 data["image"]["image_path"] = file_path415 data = copy.deepcopy(data)416 yield file_name, data417 418 def _load_from_mat_data(self, path: str):419 data = loadmat(path)["anno"][0, 0]420 image_name: str = data["imname"][0]421 objects = data["objects"][0]422 classes: list[str] = [obj["class"][0] for obj in objects]423 class_index: list[int] = [obj["class_ind"][0, 0] for obj in objects]424 masks = [obj["mask"] for obj in objects]425 annotations = []426 # part_info = []427 # part_name = []428 # part_mask = []429 for idx, obj in enumerate(objects):430 anno = {431 "category_id": class_index[idx],432 "category": classes[idx],433 "segmentation": Image.fromarray(masks[idx], 'L'),434 "parts": []435 }436 if len(obj["parts"]) == 0:437 anno["parts"] = []438 else:439 part_info = obj["parts"][0]440 part_name = [part["part_name"][0] for part in obj["parts"][0]]441 part_mask = [part["mask"] for part in obj["parts"][0]]442 anno["parts"] = [443 {444 "category_id": (CLASS_TO_PARTS445 [class_index[idx]]446 [part_name[j]]),447 "category": part_name[j],448 "segmentation": Image.fromarray(449 part_mask[j], 'L'),450 }451 for j in range(len(part_info))452 ]453 annotations.append(anno)454 feature = {455 "image": {456 "file_name": image_name,457 "height": masks[0].shape[0],458 "width": masks[0].shape[1],459 },460 "annotations": annotations,461 }462 return feature463