Skip to content

Commit

Permalink
Support arbitrary image extensions in formats (CamVid, CVAT, COCO, IC…
Browse files Browse the repository at this point in the history 
…DAR, Datumaro) (#170)

* support etxs in camvid

* update coco format

* update icdar format

* update cvat format

* add datumaro format tests
  • Loading branch information
Maxim Zhiltsov authored Mar 20, 2021
1 parent 755068f commit e3db419
Show file tree
Hide file tree
Showing 11 changed files with 308 additions and 142 deletions.
38 changes: 22 additions & 16 deletions datumaro/plugins/camvid_format.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,11 +135,14 @@ def make_camvid_categories(label_map=None):


class CamvidExtractor(SourceExtractor):
def __init__(self, path):
def __init__(self, path, subset=None):
assert osp.isfile(path), path
self._path = path
self._dataset_dir = osp.dirname(path)
super().__init__(subset=osp.splitext(osp.basename(path))[0])

if not subset:
subset = osp.splitext(osp.basename(path))[0]
super().__init__(subset=subset)

self._categories = self._load_categories(self._dataset_dir)
self._items = list(self._load_items(path).values())
Expand All @@ -156,6 +159,11 @@ def _load_categories(self, path):

def _load_items(self, path):
items = {}

labels = self._categories[AnnotationType.label]._indices
labels = { labels[label_name]: label_name
for label_name in labels }

with open(path, encoding='utf-8') as f:
for line in f:
line = line.strip()
Expand All @@ -170,29 +178,27 @@ def _load_items(self, path):
else:
objects = line.split()
image = objects[0]
item_id = ('/'.join(image.split('/')[2:]))[:-len(CamvidPath.IMAGE_EXT)]
image_path = osp.join(self._dataset_dir,
(image, image[1:])[image[0] == '/'])
item_id = osp.splitext(osp.join(*image.split('/')[2:]))[0]
image_path = osp.join(self._dataset_dir, image.lstrip('/'))

item_annotations = []
if 1 < len(objects):
gt = objects[1]
gt_path = osp.join(self._dataset_dir,
(gt, gt[1:]) [gt[0] == '/'])
inverse_cls_colormap = \
self._categories[AnnotationType.mask].inverse_colormap
mask = lazy_mask(gt_path, inverse_cls_colormap)
# loading mask through cache
mask = mask()
gt_path = osp.join(self._dataset_dir, gt.lstrip('/'))
mask = lazy_mask(gt_path,
self._categories[AnnotationType.mask].inverse_colormap)
mask = mask() # loading mask through cache

classes = np.unique(mask)
labels = self._categories[AnnotationType.label]._indices
labels = { labels[label_name]: label_name
for label_name in labels }
for label_id in classes:
if labels[label_id] in self._labels:
image = self._lazy_extract_mask(mask, label_id)
item_annotations.append(Mask(image=image, label=label_id))
item_annotations.append(
Mask(image=image, label=label_id))

items[item_id] = DatasetItem(id=item_id, subset=self._subset,
image=image_path, annotations=item_annotations)

return items

@staticmethod
Expand Down
7 changes: 4 additions & 3 deletions datumaro/plugins/coco_format/extractor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,12 @@


class _CocoExtractor(SourceExtractor):
def __init__(self, path, task, merge_instance_polygons=False):
def __init__(self, path, task, merge_instance_polygons=False, subset=None):
assert osp.isfile(path), path

subset = osp.splitext(osp.basename(path))[0].rsplit('_', maxsplit=1)
subset = subset[1] if len(subset) == 2 else None
if not subset:
subset = osp.splitext(osp.basename(path))[0].rsplit('_', maxsplit=1)
subset = subset[1] if len(subset) == 2 else None
super().__init__(subset=subset)

rootpath = ''
Expand Down
8 changes: 6 additions & 2 deletions datumaro/plugins/cvat_format/converter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@

from datumaro.components.converter import Converter
from datumaro.components.dataset import ItemStatus
from datumaro.components.extractor import AnnotationType, DatasetItem
from datumaro.components.extractor import (AnnotationType, DatasetItem,
LabelCategories)
from datumaro.util import cast, pairs

from .format import CvatPath
Expand Down Expand Up @@ -190,7 +191,8 @@ def _write_item(self, item, index):
self._writer.close_image()

def _write_meta(self):
label_cat = self._extractor.categories()[AnnotationType.label]
label_cat = self._extractor.categories().get(
AnnotationType.label, LabelCategories())
meta = OrderedDict([
("task", OrderedDict([
("id", ""),
Expand Down Expand Up @@ -222,6 +224,8 @@ def _write_meta(self):
self._writer.write_meta(meta)

def _get_label(self, label_id):
if label_id is None:
return ""
label_cat = self._extractor.categories()[AnnotationType.label]
return label_cat.items[label_id]

Expand Down
6 changes: 4 additions & 2 deletions datumaro/plugins/cvat_format/extractor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
class CvatExtractor(SourceExtractor):
_SUPPORTED_SHAPES = ('box', 'polygon', 'polyline', 'points')

def __init__(self, path):
def __init__(self, path, subset=None):
assert osp.isfile(path), path
rootpath = osp.dirname(path)
images_dir = ''
Expand All @@ -28,7 +28,9 @@ def __init__(self, path):
self._images_dir = images_dir
self._path = path

super().__init__(subset=osp.splitext(osp.basename(path))[0])
if not subset:
subset = osp.splitext(osp.basename(path))[0]
super().__init__(subset=subset)

items, categories = self._parse(path)
self._items = list(self._load_items(items).values())
Expand Down
127 changes: 64 additions & 63 deletions datumaro/plugins/icdar_format/converter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,16 +20,18 @@ def apply(self):
for subset_name, subset in self._extractor.subsets().items():
annotation = ''
for item in subset:
if item.has_image and self._save_images:
self._save_image(item, osp.join(self._save_dir, subset_name,
IcdarPath.IMAGES_DIR, item.id + IcdarPath.IMAGE_EXT))
image_filename = self._make_image_filename(item)
if self._save_images and item.has_image:
self._save_image(item, osp.join(self._save_dir,
subset_name, IcdarPath.IMAGES_DIR, image_filename))

annotation += '%s, ' % (item.id + IcdarPath.IMAGE_EXT)
annotation += '%s, ' % image_filename
for ann in item.annotations:
if ann.type != AnnotationType.caption:
continue
annotation += '\"%s\"' % ann.caption
annotation += '\n'

if len(annotation):
anno_file = osp.join(self._save_dir, subset_name, 'gt.txt')
os.makedirs(osp.dirname(anno_file), exist_ok=True)
Expand All @@ -42,9 +44,9 @@ class IcdarTextLocalizationConverter(Converter):
def apply(self):
for subset_name, subset in self._extractor.subsets().items():
for item in subset:
if item.has_image and self._save_images:
self._save_image(item, osp.join(self._save_dir, subset_name,
IcdarPath.IMAGES_DIR, item.id + IcdarPath.IMAGE_EXT))
if self._save_images and item.has_image:
self._save_image(item,
subdir=osp.join(subset_name, IcdarPath.IMAGES_DIR))

annotation = ''
for ann in item.annotations:
Expand All @@ -57,8 +59,9 @@ def apply(self):
if ann.attributes and 'text' in ann.attributes:
annotation += ',\"%s\"' % ann.attributes['text']
annotation += '\n'
anno_file = osp.join(self._save_dir, subset_name, osp.dirname(item.id),
'gt_' + osp.basename(item.id) + '.txt')

anno_file = osp.join(self._save_dir, subset_name,
osp.dirname(item.id), 'gt_' + osp.basename(item.id) + '.txt')
os.makedirs(osp.dirname(anno_file), exist_ok=True)
with open(anno_file, 'w', encoding='utf-8') as f:
f.write(annotation)
Expand All @@ -69,59 +72,57 @@ class IcdarTextSegmentationConverter(Converter):
def apply(self):
for subset_name, subset in self._extractor.subsets().items():
for item in subset:
if item.has_image and self._save_images:
self._save_image(item, osp.join(self._save_dir, subset_name,
IcdarPath.IMAGES_DIR, item.id + IcdarPath.IMAGE_EXT))
self._save_item(subset_name, subset, item)

annotation = ''
colormap = [(255, 255, 255)]
anns = [a for a in item.annotations
if a.type == AnnotationType.mask]
if anns:
is_not_index = len([p for p in anns if 'index' not in p.attributes])
if is_not_index:
raise Exception("Item %s: a mask must have"
"'index' attribute" % item.id)
anns = sorted(anns, key=lambda a: a.attributes['index'])
group = anns[0].group
for ann in anns:
if ann.group != group or (not ann.group and anns[0].group != 0):
annotation += '\n'
text = ''
if ann.attributes:
if 'text' in ann.attributes:
text = ann.attributes['text']
if text == ' ':
annotation += '#'
if 'color' in ann.attributes and \
len(ann.attributes['color'].split()) == 3:
color = ann.attributes['color'].split()
colormap.append(
(int(color[0]), int(color[1]), int(color[2])))
annotation += ' '.join(p for p in color)
else:
raise Exception("Item %s: a mask must have "
"an RGB color attribute, e. g. '10 7 50'" % item.id)
if 'center' in ann.attributes:
annotation += ' %s' % ann.attributes['center']
else:
annotation += ' - -'
bbox = ann.get_bbox()
annotation += ' %s %s %s %s' % (bbox[0], bbox[1],
bbox[0] + bbox[2], bbox[1] + bbox[3])
annotation += ' \"%s\"' % text
annotation += '\n'
group = ann.group

mask = CompiledMask.from_instance_masks(anns,
instance_labels=[m.attributes['index'] + 1 for m in anns])
mask = paint_mask(mask.class_mask,
{ i: colormap[i] for i in range(len(colormap)) })
save_image(osp.join(self._save_dir, subset_name,
item.id + '_GT' + IcdarPath.GT_EXT), mask, create_dir=True)
def _save_item(self, subset_name, subset, item):
if self._save_images and item.has_image:
self._save_image(item,
subdir=osp.join(subset_name, IcdarPath.IMAGES_DIR))

anno_file = osp.join(self._save_dir, subset_name,
item.id + '_GT' + '.txt')
os.makedirs(osp.dirname(anno_file), exist_ok=True)
with open(anno_file, 'w', encoding='utf-8') as f:
f.write(annotation)
annotation = ''
colormap = [(255, 255, 255)]
anns = [a for a in item.annotations if a.type == AnnotationType.mask]
if anns:
anns = sorted(anns, key=lambda a: int(a.attributes.get('index', 0)))
group = anns[0].group
for ann in anns:
if ann.group != group or (not ann.group and anns[0].group != 0):
annotation += '\n'
text = ''
if ann.attributes:
if 'text' in ann.attributes:
text = ann.attributes['text']
if text == ' ':
annotation += '#'
if 'color' in ann.attributes and \
len(ann.attributes['color'].split()) == 3:
color = ann.attributes['color'].split()
colormap.append(
(int(color[0]), int(color[1]), int(color[2])))
annotation += ' '.join(p for p in color)
else:
raise Exception("Item %s: a mask must have "
"an RGB color attribute, e.g. '10 7 50'" % item.id)
if 'center' in ann.attributes:
annotation += ' %s' % ann.attributes['center']
else:
annotation += ' - -'
bbox = ann.get_bbox()
annotation += ' %s %s %s %s' % (bbox[0], bbox[1],
bbox[0] + bbox[2], bbox[1] + bbox[3])
annotation += ' \"%s\"' % text
annotation += '\n'
group = ann.group

mask = CompiledMask.from_instance_masks(anns,
instance_labels=[m.attributes['index'] + 1 for m in anns])
mask = paint_mask(mask.class_mask,
{ i: colormap[i] for i in range(len(colormap)) })
save_image(osp.join(self._save_dir, subset_name,
item.id + '_GT' + IcdarPath.GT_EXT), mask, create_dir=True)

anno_file = osp.join(self._save_dir, subset_name,
item.id + '_GT' + '.txt')
os.makedirs(osp.dirname(anno_file), exist_ok=True)
with open(anno_file, 'w', encoding='utf-8') as f:
f.write(annotation)
Loading

0 comments on commit e3db419

Please sign in to comment.