-
Notifications
You must be signed in to change notification settings - Fork 3.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support for Market-1501 dataset format (#2869)
* Add support for Market-1501 dataset format * fix data access * Update Datumaro version * Add transforms * Update Changelog
- Loading branch information
1 parent
d62e176
commit ce1666f
Showing
7 changed files
with
387 additions
and
219 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,77 @@ | ||
# Copyright (C) 2021 Intel Corporation | ||
# | ||
# SPDX-License-Identifier: MIT | ||
|
||
import zipfile | ||
from tempfile import TemporaryDirectory | ||
|
||
from datumaro.components.dataset import Dataset | ||
from datumaro.components.extractor import (AnnotationType, Label, | ||
LabelCategories, Transform) | ||
|
||
from cvat.apps.dataset_manager.bindings import (CvatTaskDataExtractor, | ||
import_dm_annotations) | ||
from cvat.apps.dataset_manager.util import make_zip_archive | ||
|
||
from .registry import dm_env, exporter, importer | ||
|
||
class AttrToLabelAttr(Transform): | ||
def __init__(self, extractor, label): | ||
super().__init__(extractor) | ||
|
||
assert isinstance(label, str) | ||
self._categories = {} | ||
label_cat = self._extractor.categories().get(AnnotationType.label) | ||
if not label_cat: | ||
label_cat = LabelCategories() | ||
self._label = label_cat.add(label) | ||
self._categories[AnnotationType.label] = label_cat | ||
|
||
def categories(self): | ||
return self._categories | ||
|
||
def transform_item(self, item): | ||
annotations = item.annotations | ||
if item.attributes: | ||
annotations.append(Label(self._label, attributes=item.attributes)) | ||
item.attributes = {} | ||
return item.wrap(annotations=annotations) | ||
|
||
class LabelAttrToAttr(Transform): | ||
def __init__(self, extractor, label): | ||
super().__init__(extractor) | ||
|
||
assert isinstance(label, str) | ||
label_cat = self._extractor.categories().get(AnnotationType.label) | ||
self._label = label_cat.find(label)[0] | ||
|
||
def transform_item(self, item): | ||
annotations = item.annotations | ||
attributes = item.attributes | ||
if self._label != None: | ||
labels = [ann for ann in annotations | ||
if ann.type == AnnotationType.label \ | ||
and ann.label == self._label] | ||
if len(labels) == 1: | ||
attributes.update(labels[0].attributes) | ||
annotations.remove(labels[0]) | ||
return item.wrap(annotations=annotations, attributes=attributes) | ||
|
||
|
||
@exporter(name='Market-1501', ext='ZIP', version='1.0') | ||
def _export(dst_file, task_data, save_images=False): | ||
dataset = Dataset.from_extractors(CvatTaskDataExtractor( | ||
task_data, include_images=save_images), env=dm_env) | ||
with TemporaryDirectory() as temp_dir: | ||
dataset.transform(LabelAttrToAttr, 'market-1501') | ||
dataset.export(temp_dir, 'market1501', save_images=save_images) | ||
make_zip_archive(temp_dir, dst_file) | ||
|
||
@importer(name='Market-1501', ext='ZIP', version='1.0') | ||
def _import(src_file, task_data): | ||
with TemporaryDirectory() as tmp_dir: | ||
zipfile.ZipFile(src_file).extractall(tmp_dir) | ||
|
||
dataset = Dataset.import_from(tmp_dir, 'market1501', env=dm_env) | ||
dataset.transform(AttrToLabelAttr, 'market-1501') | ||
import_dm_annotations(dataset, task_data) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.