Skip to content

Commit

Permalink
[Datumaro] MOT format (#1289)
Browse files Browse the repository at this point in the history
* Add mot format base

* Add mot format

* Extract common code
  • Loading branch information
zhiltsov-max authored Mar 20, 2020
1 parent 8efaf58 commit 0744c6a
Show file tree
Hide file tree
Showing 6 changed files with 540 additions and 50 deletions.
30 changes: 11 additions & 19 deletions datumaro/datumaro/plugins/coco_format/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,23 +17,14 @@
AnnotationType, Points
)
from datumaro.components.cli_plugin import CliPlugin
from datumaro.util import find
from datumaro.util import find, cast
from datumaro.util.image import save_image
import datumaro.util.mask_tools as mask_tools
import datumaro.util.annotation_tools as anno_tools

from .format import CocoTask, CocoPath


def _cast(value, type_conv, default=None):
if value is None:
return default
try:
return type_conv(value)
except Exception:
return default


SegmentationMode = Enum('SegmentationMode', ['guess', 'polygons', 'mask'])

class _TaskConverter:
Expand Down Expand Up @@ -82,7 +73,7 @@ def save_image_info(self, item, filename):
'id': self._get_image_id(item),
'width': int(w),
'height': int(h),
'file_name': _cast(filename, str, ''),
'file_name': cast(filename, str, ''),
'license': 0,
'flickr_url': '',
'coco_url': '',
Expand Down Expand Up @@ -162,8 +153,8 @@ def save_categories(self, dataset):
for idx, cat in enumerate(label_categories.items):
self.categories.append({
'id': 1 + idx,
'name': _cast(cat.name, str, ''),
'supercategory': _cast(cat.parent, str, ''),
'name': cast(cat.name, str, ''),
'supercategory': cast(cat.parent, str, ''),
})

@classmethod
Expand Down Expand Up @@ -309,7 +300,7 @@ def convert_instance(self, instance, item):
elem = {
'id': self._get_ann_id(ann),
'image_id': self._get_image_id(item),
'category_id': _cast(ann.label, int, -1) + 1,
'category_id': cast(ann.label, int, -1) + 1,
'segmentation': segmentation,
'area': float(area),
'bbox': list(map(float, bbox)),
Expand All @@ -334,10 +325,11 @@ def save_categories(self, dataset):
for idx, label_cat in enumerate(label_categories.items):
cat = {
'id': 1 + idx,
'name': _cast(label_cat.name, str, ''),
'supercategory': _cast(label_cat.parent, str, ''),
'name': cast(label_cat.name, str, ''),
'supercategory': cast(label_cat.parent, str, ''),
'keypoints': [],
'skeleton': [],

}

if point_categories is not None:
Expand Down Expand Up @@ -416,8 +408,8 @@ def save_categories(self, dataset):
for idx, cat in enumerate(label_categories.items):
self.categories.append({
'id': 1 + idx,
'name': _cast(cat.name, str, ''),
'supercategory': _cast(cat.parent, str, ''),
'name': cast(cat.name, str, ''),
'supercategory': cast(cat.parent, str, ''),
})

def save_annotations(self, item):
Expand Down Expand Up @@ -504,7 +496,7 @@ def _make_task_converters(self):
def _get_image_id(self, item):
image_id = self._image_ids.get(item.id)
if image_id is None:
image_id = _cast(item.id, int, len(self._image_ids) + 1)
image_id = cast(item.id, int, len(self._image_ids) + 1)
self._image_ids[item.id] = image_id
return image_id

Expand Down
11 changes: 2 additions & 9 deletions datumaro/datumaro/plugins/cvat_format/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,19 +12,12 @@
from datumaro.components.cli_plugin import CliPlugin
from datumaro.components.converter import Converter
from datumaro.components.extractor import DEFAULT_SUBSET_NAME, AnnotationType
from datumaro.util import cast
from datumaro.util.image import save_image

from .format import CvatPath


def _cast(value, type_conv, default=None):
if value is None:
return default
try:
return type_conv(value)
except Exception:
return default

def pairwise(iterable):
a = iter(iterable)
return zip(a, a)
Expand Down Expand Up @@ -188,7 +181,7 @@ def _save_image(self, item):

def _write_item(self, item, index):
image_info = OrderedDict([
("id", str(_cast(item.id, int, index))),
("id", str(cast(item.id, int, index))),
])
if item.has_image:
size = item.image.size
Expand Down
35 changes: 14 additions & 21 deletions datumaro/datumaro/plugins/datumaro_format/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,21 +16,14 @@
Label, Mask, RleMask, Points, Polygon, PolyLine, Bbox, Caption,
LabelCategories, MaskCategories, PointsCategories
)
from datumaro.util import cast
from datumaro.util.image import save_image
import pycocotools.mask as mask_utils
from datumaro.components.cli_plugin import CliPlugin

from .format import DatumaroPath


def _cast(value, type_conv, default=None):
if value is None:
return default
try:
return type_conv(value)
except Exception:
return default

class _SubsetWriter:
def __init__(self, name, context):
self._name = name
Expand Down Expand Up @@ -108,18 +101,18 @@ def _convert_annotation(self, obj):
assert isinstance(obj, Annotation)

ann_json = {
'id': _cast(obj.id, int),
'type': _cast(obj.type.name, str),
'id': cast(obj.id, int),
'type': cast(obj.type.name, str),
'attributes': obj.attributes,
'group': _cast(obj.group, int, 0),
'group': cast(obj.group, int, 0),
}
return ann_json

def _convert_label_object(self, obj):
converted = self._convert_annotation(obj)

converted.update({
'label_id': _cast(obj.label, int),
'label_id': cast(obj.label, int),
})
return converted

Expand All @@ -133,7 +126,7 @@ def _convert_mask_object(self, obj):
np.require(obj.image, dtype=np.uint8, requirements='F'))

converted.update({
'label_id': _cast(obj.label, int),
'label_id': cast(obj.label, int),
'rle': {
# serialize as compressed COCO mask
'counts': rle['counts'].decode('ascii'),
Expand All @@ -146,7 +139,7 @@ def _convert_polyline_object(self, obj):
converted = self._convert_annotation(obj)

converted.update({
'label_id': _cast(obj.label, int),
'label_id': cast(obj.label, int),
'points': [float(p) for p in obj.points],
})
return converted
Expand All @@ -155,7 +148,7 @@ def _convert_polygon_object(self, obj):
converted = self._convert_annotation(obj)

converted.update({
'label_id': _cast(obj.label, int),
'label_id': cast(obj.label, int),
'points': [float(p) for p in obj.points],
})
return converted
Expand All @@ -164,7 +157,7 @@ def _convert_bbox_object(self, obj):
converted = self._convert_annotation(obj)

converted.update({
'label_id': _cast(obj.label, int),
'label_id': cast(obj.label, int),
'bbox': [float(p) for p in obj.get_bbox()],
})
return converted
Expand All @@ -173,7 +166,7 @@ def _convert_points_object(self, obj):
converted = self._convert_annotation(obj)

converted.update({
'label_id': _cast(obj.label, int),
'label_id': cast(obj.label, int),
'points': [float(p) for p in obj.points],
'visibility': [int(v.value) for v in obj.visibility],
})
Expand All @@ -183,7 +176,7 @@ def _convert_caption_object(self, obj):
converted = self._convert_annotation(obj)

converted.update({
'caption': _cast(obj.caption, str),
'caption': cast(obj.caption, str),
})
return converted

Expand All @@ -193,8 +186,8 @@ def _convert_label_categories(self, obj):
}
for label in obj.items:
converted['labels'].append({
'name': _cast(label.name, str),
'parent': _cast(label.parent, str),
'name': cast(label.name, str),
'parent': cast(label.parent, str),
})
return converted

Expand All @@ -218,7 +211,7 @@ def _convert_points_categories(self, obj):
for label_id, item in obj.items.items():
converted['items'].append({
'label_id': int(label_id),
'labels': [_cast(label, str) for label in item.labels],
'labels': [cast(label, str) for label in item.labels],
'adjacent': [int(v) for v in item.adjacent],
})
return converted
Expand Down
Loading

0 comments on commit 0744c6a

Please sign in to comment.