[Datumaro] Introduce image info (#1140)

* Employ transforms and item wrapper

* Add image class and tests

* Add image info support to formats

* Fix cli

* Fix merge and voc converte

* Update remote images extractor

* Codacy

* Remove item name, require path in Image

* Merge images of dataset items

* Update tests

* Add image dir converter

* Update Datumaro format

* Update COCO format with image info

* Update CVAT format with image info

* Update TFrecord format with image info

* Update VOC formar with image info

* Update YOLO format with image info

* Update dataset manager bindings with image info

* Add image name to id transform

* Fix coco export
main
zhiltsov-max 6 years ago committed by GitHub
parent 0db48afa9a
commit a376ee76fd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -14,7 +14,7 @@ from cvat.apps.engine.annotation import TaskAnnotation
from cvat.apps.engine.models import Task, ShapeType, AttributeType from cvat.apps.engine.models import Task, ShapeType, AttributeType
import datumaro.components.extractor as datumaro import datumaro.components.extractor as datumaro
from datumaro.util.image import lazy_image from datumaro.util.image import Image
class CvatImagesDirExtractor(datumaro.Extractor): class CvatImagesDirExtractor(datumaro.Extractor):
@ -29,8 +29,7 @@ class CvatImagesDirExtractor(datumaro.Extractor):
path = osp.join(dirpath, name) path = osp.join(dirpath, name)
if self._is_image(path): if self._is_image(path):
item_id = Task.get_image_frame(path) item_id = Task.get_image_frame(path)
item = datumaro.DatasetItem( item = datumaro.DatasetItem(id=item_id, image=path)
id=item_id, image=lazy_image(path))
items.append((item.id, item)) items.append((item.id, item))
items = sorted(items, key=lambda e: int(e[0])) items = sorted(items, key=lambda e: int(e[0]))
@ -49,11 +48,6 @@ class CvatImagesDirExtractor(datumaro.Extractor):
def subsets(self): def subsets(self):
return self._subsets return self._subsets
def get(self, item_id, subset=None, path=None):
if path or subset:
raise KeyError()
return self._items[item_id]
def _is_image(self, path): def _is_image(self, path):
for ext in self._SUPPORTED_FORMATS: for ext in self._SUPPORTED_FORMATS:
if osp.isfile(path) and path.endswith(ext): if osp.isfile(path) and path.endswith(ext):
@ -61,29 +55,24 @@ class CvatImagesDirExtractor(datumaro.Extractor):
return False return False
class CvatTaskExtractor(datumaro.Extractor): class CvatAnnotationsExtractor(datumaro.Extractor):
def __init__(self, url, db_task, user): def __init__(self, url, cvat_annotations):
self._db_task = db_task self._categories = self._load_categories(cvat_annotations)
self._categories = self._load_categories()
cvat_annotations = TaskAnnotation(db_task.id, user)
with transaction.atomic():
cvat_annotations.init_from_db()
cvat_annotations = Annotation(cvat_annotations.ir_data, db_task)
dm_annotations = [] dm_annotations = []
for cvat_anno in cvat_annotations.group_by_frame(): for cvat_frame_anno in cvat_annotations.group_by_frame():
dm_anno = self._read_cvat_anno(cvat_anno) dm_anno = self._read_cvat_anno(cvat_frame_anno, cvat_annotations)
dm_item = datumaro.DatasetItem( dm_image = Image(path=cvat_frame_anno.name, size=(
id=cvat_anno.frame, annotations=dm_anno) cvat_frame_anno.height, cvat_frame_anno.width)
)
dm_item = datumaro.DatasetItem(id=cvat_frame_anno.frame,
annotations=dm_anno, image=dm_image)
dm_annotations.append((dm_item.id, dm_item)) dm_annotations.append((dm_item.id, dm_item))
dm_annotations = sorted(dm_annotations, key=lambda e: int(e[0])) dm_annotations = sorted(dm_annotations, key=lambda e: int(e[0]))
self._items = OrderedDict(dm_annotations) self._items = OrderedDict(dm_annotations)
self._subsets = None
def __iter__(self): def __iter__(self):
for item in self._items.values(): for item in self._items.values():
yield item yield item
@ -91,70 +80,58 @@ class CvatTaskExtractor(datumaro.Extractor):
def __len__(self): def __len__(self):
return len(self._items) return len(self._items)
# pylint: disable=no-self-use
def subsets(self): def subsets(self):
return self._subsets return []
# pylint: enable=no-self-use
def get(self, item_id, subset=None, path=None): def categories(self):
if path or subset: return self._categories
raise KeyError()
return self._items[item_id]
def _load_categories(self): @staticmethod
def _load_categories(cvat_anno):
categories = {} categories = {}
label_categories = datumaro.LabelCategories() label_categories = datumaro.LabelCategories()
db_labels = self._db_task.label_set.all() for _, label in cvat_anno.meta['task']['labels']:
for db_label in db_labels: label_categories.add(label['name'])
db_attributes = db_label.attributespec_set.all() for _, attr in label['attributes']:
label_categories.add(db_label.name) label_categories.attributes.add(attr['name'])
for db_attr in db_attributes:
label_categories.attributes.add(db_attr.name)
categories[datumaro.AnnotationType.label] = label_categories categories[datumaro.AnnotationType.label] = label_categories
return categories return categories
def categories(self): def _read_cvat_anno(self, cvat_frame_anno, cvat_task_anno):
return self._categories
def _read_cvat_anno(self, cvat_anno):
item_anno = [] item_anno = []
categories = self.categories() categories = self.categories()
label_cat = categories[datumaro.AnnotationType.label] label_cat = categories[datumaro.AnnotationType.label]
map_label = lambda name: label_cat.find(name)[0]
label_map = {} label_attrs = {
label_attrs = {} label['name']: label['attributes']
db_labels = self._db_task.label_set.all() for _, label in cvat_task_anno.meta['task']['labels']
for db_label in db_labels: }
label_map[db_label.name] = label_cat.find(db_label.name)[0]
attrs = {}
db_attributes = db_label.attributespec_set.all()
for db_attr in db_attributes:
attrs[db_attr.name] = db_attr
label_attrs[db_label.name] = attrs
map_label = lambda label_db_name: label_map[label_db_name]
def convert_attrs(label, cvat_attrs): def convert_attrs(label, cvat_attrs):
cvat_attrs = {a.name: a.value for a in cvat_attrs} cvat_attrs = {a.name: a.value for a in cvat_attrs}
dm_attr = dict() dm_attr = dict()
for attr_name, attr_spec in label_attrs[label].items(): for _, a_desc in label_attrs[label]:
attr_value = cvat_attrs.get(attr_name, attr_spec.default_value) a_name = a_desc['name']
a_value = cvat_attrs.get(a_name, a_desc['default_value'])
try: try:
if attr_spec.input_type == AttributeType.NUMBER: if a_desc['input_type'] == AttributeType.NUMBER:
attr_value = float(attr_value) a_value = float(a_value)
elif attr_spec.input_type == AttributeType.CHECKBOX: elif a_desc['input_type'] == AttributeType.CHECKBOX:
attr_value = attr_value.lower() == 'true' a_value = (a_value.lower() == 'true')
dm_attr[attr_name] = attr_value dm_attr[a_name] = a_value
except Exception as e: except Exception as e:
slogger.task[self._db_task.id].error( raise Exception(
"Failed to convert attribute '%s'='%s': %s" % \ "Failed to convert attribute '%s'='%s': %s" %
(attr_name, attr_value, e)) (a_name, a_value, e))
return dm_attr return dm_attr
for tag_obj in cvat_anno.tags: for tag_obj in cvat_frame_anno.tags:
anno_group = tag_obj.group anno_group = tag_obj.group
anno_label = map_label(tag_obj.label) anno_label = map_label(tag_obj.label)
anno_attr = convert_attrs(tag_obj.label, tag_obj.attributes) anno_attr = convert_attrs(tag_obj.label, tag_obj.attributes)
@ -163,7 +140,7 @@ class CvatTaskExtractor(datumaro.Extractor):
attributes=anno_attr, group=anno_group) attributes=anno_attr, group=anno_group)
item_anno.append(anno) item_anno.append(anno)
for shape_obj in cvat_anno.labeled_shapes: for shape_obj in cvat_frame_anno.labeled_shapes:
anno_group = shape_obj.group anno_group = shape_obj.group
anno_label = map_label(shape_obj.label) anno_label = map_label(shape_obj.label)
anno_attr = convert_attrs(shape_obj.label, shape_obj.attributes) anno_attr = convert_attrs(shape_obj.label, shape_obj.attributes)
@ -183,8 +160,64 @@ class CvatTaskExtractor(datumaro.Extractor):
anno = datumaro.Bbox(x0, y0, x1 - x0, y1 - y0, anno = datumaro.Bbox(x0, y0, x1 - x0, y1 - y0,
label=anno_label, attributes=anno_attr, group=anno_group) label=anno_label, attributes=anno_attr, group=anno_group)
else: else:
raise Exception("Unknown shape type '%s'" % (shape_obj.type)) raise Exception("Unknown shape type '%s'" % shape_obj.type)
item_anno.append(anno) item_anno.append(anno)
return item_anno return item_anno
class CvatTaskExtractor(CvatAnnotationsExtractor):
def __init__(self, url, db_task, user):
cvat_annotations = TaskAnnotation(db_task.id, user)
with transaction.atomic():
cvat_annotations.init_from_db()
cvat_annotations = Annotation(cvat_annotations.ir_data, db_task)
super().__init__(url, cvat_annotations)
def match_frame(item, cvat_task_anno):
frame_number = None
if frame_number is None:
try:
frame_number = cvat_task_anno.match_frame(item.id)
except Exception:
pass
if frame_number is None and item.has_image:
try:
frame_number = cvat_task_anno.match_frame(item.image.filename)
except Exception:
pass
if frame_number is None:
try:
frame_number = int(item.id)
except Exception:
pass
if not frame_number in cvat_task_anno.frame_info:
raise Exception("Could not match item id: '%s' with any task frame" %
item.id)
return frame_number
def import_dm_annotations(dm_dataset, cvat_task_anno):
shapes = {
datumaro.AnnotationType.bbox: ShapeType.RECTANGLE,
datumaro.AnnotationType.polygon: ShapeType.POLYGON,
datumaro.AnnotationType.polyline: ShapeType.POLYLINE,
datumaro.AnnotationType.points: ShapeType.POINTS,
}
label_cat = dm_dataset.categories()[datumaro.AnnotationType.label]
for item in dm_dataset:
frame_number = match_frame(item, cvat_task_anno)
for ann in item.annotations:
if ann.type in shapes:
cvat_task_anno.add_shape(cvat_task_anno.LabeledShape(
type=shapes[ann.type],
frame=frame_number,
label=label_cat.items[ann.label].name,
points=ann.points,
occluded=False,
attributes=[],
))

@ -13,7 +13,7 @@ from datumaro.components.config import (Config,
SchemaBuilder as _SchemaBuilder, SchemaBuilder as _SchemaBuilder,
) )
import datumaro.components.extractor as datumaro import datumaro.components.extractor as datumaro
from datumaro.util.image import lazy_image, load_image from datumaro.util.image import lazy_image, load_image, Image
from cvat.utils.cli.core import CLI as CVAT_CLI, CVAT_API_V1 from cvat.utils.cli.core import CLI as CVAT_CLI, CVAT_API_V1
@ -103,8 +103,11 @@ class cvat_rest_api_task_images(datumaro.SourceExtractor):
items = [] items = []
for entry in image_list: for entry in image_list:
item_id = entry['id'] item_id = entry['id']
item = datumaro.DatasetItem( size = None
id=item_id, image=self._make_image_loader(item_id)) if entry.get('height') and entry.get('width'):
size = (entry['height'], entry['width'])
image = Image(data=self._make_image_loader(item_id), size=size)
item = datumaro.DatasetItem(id=item_id, image=image)
items.append((item.id, item)) items.append((item.id, item))
items = sorted(items, key=lambda e: int(e[0])) items = sorted(items, key=lambda e: int(e[0]))

@ -156,16 +156,17 @@ def import_command(args):
if project_name is None: if project_name is None:
project_name = osp.basename(project_dir) project_name = osp.basename(project_dir)
extra_args = {}
try: try:
env = Environment() env = Environment()
importer = env.make_importer(args.format) importer = env.make_importer(args.format)
if hasattr(importer, 'from_cmdline'):
extra_args = importer.from_cmdline(args.extra_args)
except KeyError: except KeyError:
raise CliException("Importer for format '%s' is not found" % \ raise CliException("Importer for format '%s' is not found" % \
args.format) args.format)
extra_args = {}
if hasattr(importer, 'from_cmdline'):
extra_args = importer.from_cmdline(args.extra_args)
log.info("Importing project from '%s' as '%s'" % \ log.info("Importing project from '%s' as '%s'" % \
(args.source, args.format)) (args.source, args.format))
@ -293,13 +294,14 @@ def export_command(args):
try: try:
converter = project.env.converters.get(args.format) converter = project.env.converters.get(args.format)
if hasattr(converter, 'from_cmdline'):
extra_args = converter.from_cmdline(args.extra_args)
converter = converter(**extra_args)
except KeyError: except KeyError:
raise CliException("Converter for format '%s' is not found" % \ raise CliException("Converter for format '%s' is not found" % \
args.format) args.format)
if hasattr(converter, 'from_cmdline'):
extra_args = converter.from_cmdline(args.extra_args)
converter = converter(**extra_args)
filter_args = FilterModes.make_filter_args(args.filter_mode) filter_args = FilterModes.make_filter_args(args.filter_mode)
log.info("Loading the project...") log.info("Loading the project...")
@ -559,14 +561,15 @@ def transform_command(args):
(project.config.project_name, make_file_name(args.transform))) (project.config.project_name, make_file_name(args.transform)))
dst_dir = osp.abspath(dst_dir) dst_dir = osp.abspath(dst_dir)
extra_args = {}
try: try:
transform = project.env.transforms.get(args.transform) transform = project.env.transforms.get(args.transform)
if hasattr(transform, 'from_cmdline'):
extra_args = transform.from_cmdline(args.extra_args)
except KeyError: except KeyError:
raise CliException("Transform '%s' is not found" % args.transform) raise CliException("Transform '%s' is not found" % args.transform)
extra_args = {}
if hasattr(transform, 'from_cmdline'):
extra_args = transform.from_cmdline(args.extra_args)
log.info("Loading the project...") log.info("Loading the project...")
dataset = project.make_dataset() dataset = project.make_dataset()

@ -4,7 +4,7 @@
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
from lxml import etree as ET # NOTE: lxml has proper XPath implementation from lxml import etree as ET # NOTE: lxml has proper XPath implementation
from datumaro.components.extractor import (DatasetItem, Extractor, from datumaro.components.extractor import (Transform,
Annotation, AnnotationType, Annotation, AnnotationType,
Label, Mask, Points, Polygon, PolyLine, Bbox, Caption, Label, Mask, Points, Polygon, PolyLine, Bbox, Caption,
) )
@ -31,11 +31,12 @@ class DatasetItemEncoder:
def encode_image(cls, image): def encode_image(cls, image):
image_elem = ET.Element('image') image_elem = ET.Element('image')
h, w = image.shape[:2] h, w = image.size
c = 1 if len(image.shape) == 2 else image.shape[2]
ET.SubElement(image_elem, 'width').text = str(w) ET.SubElement(image_elem, 'width').text = str(w)
ET.SubElement(image_elem, 'height').text = str(h) ET.SubElement(image_elem, 'height').text = str(h)
ET.SubElement(image_elem, 'depth').text = str(c)
ET.SubElement(image_elem, 'has_data').text = '%d' % int(image.has_data)
ET.SubElement(image_elem, 'path').text = image.path
return image_elem return image_elem
@ -82,10 +83,6 @@ class DatasetItemEncoder:
str(cls._get_label(obj.label, categories)) str(cls._get_label(obj.label, categories))
ET.SubElement(ann_elem, 'label_id').text = str(obj.label) ET.SubElement(ann_elem, 'label_id').text = str(obj.label)
mask = obj.image
if mask is not None:
ann_elem.append(cls.encode_image(mask))
return ann_elem return ann_elem
@classmethod @classmethod
@ -218,39 +215,9 @@ def XPathDatasetFilter(extractor, xpath=None):
DatasetItemEncoder.encode(item, extractor.categories()))) DatasetItemEncoder.encode(item, extractor.categories())))
return extractor.select(f) return extractor.select(f)
class XPathAnnotationsFilter(Extractor): # NOTE: essentially, a transform class XPathAnnotationsFilter(Transform):
class ItemWrapper(DatasetItem):
def __init__(self, item, annotations):
self._item = item
self._annotations = annotations
@DatasetItem.id.getter
def id(self):
return self._item.id
@DatasetItem.subset.getter
def subset(self):
return self._item.subset
@DatasetItem.path.getter
def path(self):
return self._item.path
@DatasetItem.annotations.getter
def annotations(self):
return self._annotations
@DatasetItem.has_image.getter
def has_image(self):
return self._item.has_image
@DatasetItem.image.getter
def image(self):
return self._item.image
def __init__(self, extractor, xpath=None, remove_empty=False): def __init__(self, extractor, xpath=None, remove_empty=False):
super().__init__() super().__init__(extractor)
self._extractor = extractor
if xpath is not None: if xpath is not None:
xpath = ET.XPath(xpath) xpath = ET.XPath(xpath)
@ -258,24 +225,16 @@ class XPathAnnotationsFilter(Extractor): # NOTE: essentially, a transform
self._remove_empty = remove_empty self._remove_empty = remove_empty
def __len__(self):
return len(self._extractor)
def __iter__(self): def __iter__(self):
for item in self._extractor: for item in self._extractor:
item = self._filter_item(item) item = self.transform_item(item)
if item is not None: if item is not None:
yield item yield item
def subsets(self): def transform_item(self, item):
return self._extractor.subsets()
def categories(self):
return self._extractor.categories()
def _filter_item(self, item):
if self._filter is None: if self._filter is None:
return item return item
encoded = DatasetItemEncoder.encode(item, self._extractor.categories()) encoded = DatasetItemEncoder.encode(item, self._extractor.categories())
filtered = self._filter(encoded) filtered = self._filter(encoded)
filtered = [elem for elem in filtered if elem.tag == 'annotation'] filtered = [elem for elem in filtered if elem.tag == 'annotation']
@ -285,4 +244,4 @@ class XPathAnnotationsFilter(Extractor): # NOTE: essentially, a transform
if self._remove_empty and len(annotations) == 0: if self._remove_empty and len(annotations) == 0:
return None return None
return self.ItemWrapper(item, annotations) return self.wrap_item(item, annotations=annotations)

@ -7,6 +7,7 @@ from collections import namedtuple
from enum import Enum from enum import Enum
import numpy as np import numpy as np
from datumaro.util.image import Image
AnnotationType = Enum('AnnotationType', AnnotationType = Enum('AnnotationType',
[ [
@ -418,8 +419,8 @@ class PolyLine(_Shape):
class Polygon(_Shape): class Polygon(_Shape):
# pylint: disable=redefined-builtin # pylint: disable=redefined-builtin
def __init__(self, points=None, z_order=None, def __init__(self, points=None, label=None,
label=None, id=None, attributes=None, group=None): z_order=None, id=None, attributes=None, group=None):
if points is not None: if points is not None:
# keep the message on the single line to produce # keep the message on the single line to produce
# informative output # informative output
@ -575,27 +576,34 @@ class Caption(Annotation):
class DatasetItem: class DatasetItem:
# pylint: disable=redefined-builtin # pylint: disable=redefined-builtin
def __init__(self, id, annotations=None, def __init__(self, id=None, annotations=None,
subset=None, path=None, image=None): subset=None, path=None, image=None):
assert id is not None assert id is not None
if not isinstance(id, str): self._id = str(id)
id = str(id)
assert len(id) != 0
self._id = id
if subset is None: if subset is None:
subset = '' subset = ''
assert isinstance(subset, str) else:
subset = str(subset)
self._subset = subset self._subset = subset
if path is None: if path is None:
path = [] path = []
else:
path = list(path)
self._path = path self._path = path
if annotations is None: if annotations is None:
annotations = [] annotations = []
else:
annotations = list(annotations)
self._annotations = annotations self._annotations = annotations
if callable(image) or isinstance(image, np.ndarray):
image = Image(data=image)
elif isinstance(image, str):
image = Image(path=image)
assert image is None or isinstance(image, Image)
self._image = image self._image = image
# pylint: enable=redefined-builtin # pylint: enable=redefined-builtin
@ -617,8 +625,6 @@ class DatasetItem:
@property @property
def image(self): def image(self):
if callable(self._image):
return self._image()
return self._image return self._image
@property @property
@ -631,11 +637,16 @@ class DatasetItem:
return \ return \
(self.id == other.id) and \ (self.id == other.id) and \
(self.subset == other.subset) and \ (self.subset == other.subset) and \
(self.annotations == other.annotations) and \
(self.path == other.path) and \ (self.path == other.path) and \
(self.has_image == other.has_image) and \ (self.annotations == other.annotations) and \
(self.has_image and np.array_equal(self.image, other.image) or \ (self.image == other.image)
not self.has_image)
def wrap(item, **kwargs):
expected_args = {'id', 'annotations', 'subset', 'path', 'image'}
for k in expected_args:
if k not in kwargs:
kwargs[k] = getattr(item, k)
return DatasetItem(**kwargs)
class IExtractor: class IExtractor:
def __iter__(self): def __iter__(self):
@ -656,9 +667,6 @@ class IExtractor:
def select(self, pred): def select(self, pred):
raise NotImplementedError() raise NotImplementedError()
def get(self, item_id, subset=None, path=None):
raise NotImplementedError()
class _DatasetFilter: class _DatasetFilter:
def __init__(self, iterable, predicate): def __init__(self, iterable, predicate):
self.iterable = iterable self.iterable = iterable
@ -741,16 +749,9 @@ class Importer:
raise NotImplementedError() raise NotImplementedError()
class Transform(Extractor): class Transform(Extractor):
@classmethod @staticmethod
def wrap_item(cls, item, **kwargs): def wrap_item(item, **kwargs):
expected_args = {'id', 'annotations', 'subset', 'path', 'image'} return item.wrap(**kwargs)
for k in expected_args:
if k not in kwargs:
if k == 'image' and item.has_image:
kwargs[k] = lambda: item.image
else:
kwargs[k] = getattr(item, k)
return DatasetItem(**kwargs)
def __init__(self, extractor): def __init__(self, extractor):
super().__init__() super().__init__()

@ -5,7 +5,7 @@
import numpy as np import numpy as np
from datumaro.components.extractor import DatasetItem, Extractor from datumaro.components.extractor import Transform
# pylint: disable=no-self-use # pylint: disable=no-self-use
@ -23,37 +23,9 @@ class Launcher:
return None return None
# pylint: enable=no-self-use # pylint: enable=no-self-use
class InferenceWrapper(Extractor): class InferenceWrapper(Transform):
class ItemWrapper(DatasetItem):
def __init__(self, item, annotations, path=None):
super().__init__(id=item.id)
self._annotations = annotations
self._item = item
self._path = path
@DatasetItem.id.getter
def id(self):
return self._item.id
@DatasetItem.subset.getter
def subset(self):
return self._item.subset
@DatasetItem.path.getter
def path(self):
return self._path
@DatasetItem.annotations.getter
def annotations(self):
return self._annotations
@DatasetItem.image.getter
def image(self):
return self._item.image
def __init__(self, extractor, launcher, batch_size=1): def __init__(self, extractor, launcher, batch_size=1):
super().__init__() super().__init__(extractor)
self._extractor = extractor
self._launcher = launcher self._launcher = launcher
self._batch_size = batch_size self._batch_size = batch_size
@ -71,25 +43,23 @@ class InferenceWrapper(Extractor):
if len(batch_items) == 0: if len(batch_items) == 0:
break break
inputs = np.array([item.image for item in batch_items]) inputs = np.array([item.image.data for item in batch_items])
inference = self._launcher.launch(inputs) inference = self._launcher.launch(inputs)
for item, annotations in zip(batch_items, inference): for item, annotations in zip(batch_items, inference):
yield self.ItemWrapper(item, annotations) yield self.wrap_item(item, annotations=annotations)
def __len__(self):
return len(self._extractor)
def subsets(self):
return self._extractor.subsets()
def get_subset(self, name): def get_subset(self, name):
subset = self._extractor.get_subset(name) subset = self._extractor.get_subset(name)
return InferenceWrapper(subset, return InferenceWrapper(subset, self._launcher, self._batch_size)
self._launcher, self._batch_size)
def categories(self): def categories(self):
launcher_override = self._launcher.get_categories() launcher_override = self._launcher.get_categories()
if launcher_override is not None: if launcher_override is not None:
return launcher_override return launcher_override
return self._extractor.categories() return self._extractor.categories()
def transform_item(self, item):
inputs = np.expand_dims(item.image, axis=0)
annotations = self._launcher.launch(inputs)[0]
return self.wrap_item(item, annotations=annotations)

@ -302,45 +302,6 @@ class Subset(Extractor):
def categories(self): def categories(self):
return self._parent.categories() return self._parent.categories()
class DatasetItemWrapper(DatasetItem):
def __init__(self, item, path, annotations, image=None):
self._item = item
if path is None:
path = []
self._path = path
self._annotations = annotations
self._image = image
@DatasetItem.id.getter
def id(self):
return self._item.id
@DatasetItem.subset.getter
def subset(self):
return self._item.subset
@DatasetItem.path.getter
def path(self):
return self._path
@DatasetItem.annotations.getter
def annotations(self):
return self._annotations
@DatasetItem.has_image.getter
def has_image(self):
if self._image is not None:
return True
return self._item.has_image
@DatasetItem.image.getter
def image(self):
if self._image is not None:
if callable(self._image):
return self._image()
return self._image
return self._item.image
class Dataset(Extractor): class Dataset(Extractor):
@classmethod @classmethod
def from_extractors(cls, *sources): def from_extractors(cls, *sources):
@ -364,17 +325,9 @@ class Dataset(Extractor):
existing_item = subsets[item.subset].items.get(item.id) existing_item = subsets[item.subset].items.get(item.id)
if existing_item is not None: if existing_item is not None:
image = None item = self._merge_items(existing_item, item, path=path)
if existing_item.has_image:
# TODO: think of image comparison
image = cls._lazy_image(existing_item)
item = DatasetItemWrapper(item=item, path=path,
image=image, annotations=self._merge_anno(
existing_item.annotations, item.annotations))
else: else:
item = DatasetItemWrapper(item=item, path=path, item = item.wrap(path=path, annotations=item.annotations)
annotations=item.annotations)
subsets[item.subset].items[item.id] = item subsets[item.subset].items[item.id] = item
@ -423,8 +376,7 @@ class Dataset(Extractor):
if subset is None: if subset is None:
subset = item.subset subset = item.subset
item = DatasetItemWrapper(item=item, path=None, item = item.wrap(path=None, annotations=item.annotations)
annotations=item.annotations)
if item.subset not in self._subsets: if item.subset not in self._subsets:
self._subsets[item.subset] = Subset(self) self._subsets[item.subset] = Subset(self)
self._subsets[subset].items[item_id] = item self._subsets[subset].items[item_id] = item
@ -453,6 +405,34 @@ class Dataset(Extractor):
# NOTE: avoid https://docs.python.org/3/faq/programming.html#why-do-lambdas-defined-in-a-loop-with-different-values-all-return-the-same-result # NOTE: avoid https://docs.python.org/3/faq/programming.html#why-do-lambdas-defined-in-a-loop-with-different-values-all-return-the-same-result
return lambda: item.image return lambda: item.image
@classmethod
def _merge_items(cls, existing_item, current_item, path=None):
image = None
if existing_item.has_image and current_item.has_image:
if existing_item.image.has_data:
image = existing_item.image
else:
image = current_item.image
if existing_item.image.path != current_item.image.path:
if not existing_item.image.path:
image._path = current_item.image.path
if all([existing_item.image._size, current_item.image._size]):
assert existing_item.image._size == current_item.image._size, "Image info differs for item '%s'" % item.id
elif existing_item.image._size:
image._size = existing_item.image._size
else:
image._size = current_item.image._size
elif existing_item.has_image:
image = existing_item.image
else:
image = current_item.image
return existing_item.wrap(path=path,
image=image, annotations=cls._merge_anno(
existing_item.annotations, current_item.annotations))
@staticmethod @staticmethod
def _merge_anno(a, b): def _merge_anno(a, b):
from itertools import chain from itertools import chain
@ -520,17 +500,10 @@ class ProjectDataset(Dataset):
for item in source: for item in source:
existing_item = subsets[item.subset].items.get(item.id) existing_item = subsets[item.subset].items.get(item.id)
if existing_item is not None: if existing_item is not None:
image = None
if existing_item.has_image:
# TODO: think of image comparison
image = self._lazy_image(existing_item)
path = existing_item.path path = existing_item.path
if item.path != path: if item.path != path:
path = None # NOTE: move to our own dataset path = None # NOTE: move to our own dataset
item = DatasetItemWrapper(item=item, path=path, item = self._merge_items(existing_item, item, path=path)
image=image, annotations=self._merge_anno(
existing_item.annotations, item.annotations))
else: else:
s_config = config.sources[source_name] s_config = config.sources[source_name]
if s_config and \ if s_config and \
@ -542,8 +515,7 @@ class ProjectDataset(Dataset):
if path is None: if path is None:
path = [] path = []
path = [source_name] + path path = [source_name] + path
item = DatasetItemWrapper(item=item, path=path, item = item.wrap(path=path, annotations=item.annotations)
annotations=item.annotations)
subsets[item.subset].items[item.id] = item subsets[item.subset].items[item.id] = item
@ -558,7 +530,7 @@ class ProjectDataset(Dataset):
if existing_item.has_image: if existing_item.has_image:
# TODO: think of image comparison # TODO: think of image comparison
image = self._lazy_image(existing_item) image = self._lazy_image(existing_item)
item = DatasetItemWrapper(item=item, path=None, item = item.wrap(path=None,
annotations=item.annotations, image=image) annotations=item.annotations, image=image)
subsets[item.subset].items[item.id] = item subsets[item.subset].items[item.id] = item
@ -597,8 +569,7 @@ class ProjectDataset(Dataset):
if subset is None: if subset is None:
subset = item.subset subset = item.subset
item = DatasetItemWrapper(item=item, path=path, item = item.wrap(path=path, annotations=item.annotations)
annotations=item.annotations)
if item.subset not in self._subsets: if item.subset not in self._subsets:
self._subsets[item.subset] = Subset(self) self._subsets[item.subset] = Subset(self)
self._subsets[subset].items[item_id] = item self._subsets[subset].items[item_id] = item

@ -67,15 +67,18 @@ class _TaskConverter:
def is_empty(self): def is_empty(self):
return len(self._data['annotations']) == 0 return len(self._data['annotations']) == 0
def _get_image_id(self, item):
return self._context._get_image_id(item)
def save_image_info(self, item, filename): def save_image_info(self, item, filename):
if item.has_image: if item.has_image:
h, w = item.image.shape[:2] h, w = item.image.size
else: else:
h = 0 h = 0
w = 0 w = 0
self._data['images'].append({ self._data['images'].append({
'id': _cast(item.id, int, 0), 'id': self._get_image_id(item),
'width': int(w), 'width': int(w),
'height': int(h), 'height': int(h),
'file_name': _cast(filename, str, ''), 'file_name': _cast(filename, str, ''),
@ -130,13 +133,13 @@ class _CaptionsConverter(_TaskConverter):
pass pass
def save_annotations(self, item): def save_annotations(self, item):
for ann in item.annotations: for ann_idx, ann in enumerate(item.annotations):
if ann.type != AnnotationType.caption: if ann.type != AnnotationType.caption:
continue continue
elem = { elem = {
'id': self._get_ann_id(ann), 'id': self._get_ann_id(ann),
'image_id': _cast(item.id, int, 0), 'image_id': self._get_image_id(item),
'category_id': 0, # NOTE: workaround for a bug in cocoapi 'category_id': 0, # NOTE: workaround for a bug in cocoapi
'caption': ann.caption, 'caption': ann.caption,
} }
@ -144,7 +147,8 @@ class _CaptionsConverter(_TaskConverter):
try: try:
elem['score'] = float(ann.attributes['score']) elem['score'] = float(ann.attributes['score'])
except Exception as e: except Exception as e:
log.warning("Failed to convert attribute 'score': %e" % e) log.warning("Item '%s', ann #%s: failed to convert "
"attribute 'score': %e" % (item.id, ann_idx, e))
self.annotations.append(elem) self.annotations.append(elem)
@ -293,10 +297,10 @@ class _InstancesConverter(_TaskConverter):
return return
if not item.has_image: if not item.has_image:
log.warn("Skipping writing instances for " log.warn("Item '%s': skipping writing instances "
"item '%s' as it has no image info" % item.id) "since no image info available" % item.id)
return return
h, w, _ = item.image.shape h, w = item.image.size
instances = [self.find_instance_parts(i, w, h) for i in instances] instances = [self.find_instance_parts(i, w, h) for i in instances]
if self._context._crop_covered: if self._context._crop_covered:
@ -319,7 +323,7 @@ class _InstancesConverter(_TaskConverter):
area = 0 area = 0
if segmentation: if segmentation:
if item.has_image: if item.has_image:
h, w, _ = item.image.shape h, w = item.image.size
else: else:
# NOTE: here we can guess the image size as # NOTE: here we can guess the image size as
# it is only needed for the area computation # it is only needed for the area computation
@ -339,7 +343,7 @@ class _InstancesConverter(_TaskConverter):
elem = { elem = {
'id': self._get_ann_id(ann), 'id': self._get_ann_id(ann),
'image_id': _cast(item.id, int, 0), '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, 'segmentation': segmentation,
'area': float(area), 'area': float(area),
@ -350,7 +354,8 @@ class _InstancesConverter(_TaskConverter):
try: try:
elem['score'] = float(ann.attributes['score']) elem['score'] = float(ann.attributes['score'])
except Exception as e: except Exception as e:
log.warning("Failed to convert attribute 'score': %e" % e) log.warning("Item '%s': failed to convert attribute "
"'score': %e" % (item.id, e))
return elem return elem
@ -452,14 +457,15 @@ class _LabelsConverter(_TaskConverter):
elem = { elem = {
'id': self._get_ann_id(ann), 'id': self._get_ann_id(ann),
'image_id': _cast(item.id, int, 0), 'image_id': self._get_image_id(item),
'category_id': int(ann.label) + 1, 'category_id': int(ann.label) + 1,
} }
if 'score' in ann.attributes: if 'score' in ann.attributes:
try: try:
elem['score'] = float(ann.attributes['score']) elem['score'] = float(ann.attributes['score'])
except Exception as e: except Exception as e:
log.warning("Failed to convert attribute 'score': %e" % e) log.warning("Item '%s': failed to convert attribute "
"'score': %e" % (item.id, e))
self.annotations.append(elem) self.annotations.append(elem)
@ -501,31 +507,50 @@ class _Converter:
self._crop_covered = crop_covered self._crop_covered = crop_covered
def make_dirs(self): self._image_ids = {}
def _make_dirs(self):
self._images_dir = osp.join(self._save_dir, CocoPath.IMAGES_DIR) self._images_dir = osp.join(self._save_dir, CocoPath.IMAGES_DIR)
os.makedirs(self._images_dir, exist_ok=True) os.makedirs(self._images_dir, exist_ok=True)
self._ann_dir = osp.join(self._save_dir, CocoPath.ANNOTATIONS_DIR) self._ann_dir = osp.join(self._save_dir, CocoPath.ANNOTATIONS_DIR)
os.makedirs(self._ann_dir, exist_ok=True) os.makedirs(self._ann_dir, exist_ok=True)
def make_task_converter(self, task): def _make_task_converter(self, task):
if task not in self._TASK_CONVERTER: if task not in self._TASK_CONVERTER:
raise NotImplementedError() raise NotImplementedError()
return self._TASK_CONVERTER[task](self) return self._TASK_CONVERTER[task](self)
def make_task_converters(self): def _make_task_converters(self):
return { return {
task: self.make_task_converter(task) for task in self._tasks task: self._make_task_converter(task) for task in self._tasks
} }
def save_image(self, item, filename): 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)
self._image_ids[item.id] = image_id
return image_id
def _save_image(self, item):
image = item.image.data
if image is None:
log.warning("Item '%s' has no image" % item.id)
return ''
filename = item.image.filename
if filename:
filename = osp.splitext(filename)[0]
else:
filename = item.id
filename += CocoPath.IMAGE_EXT
path = osp.join(self._images_dir, filename) path = osp.join(self._images_dir, filename)
save_image(path, item.image) save_image(path, image)
return filename
return path
def convert(self): def convert(self):
self.make_dirs() self._make_dirs()
subsets = self._extractor.subsets() subsets = self._extractor.subsets()
if len(subsets) == 0: if len(subsets) == 0:
@ -538,18 +563,18 @@ class _Converter:
subset_name = DEFAULT_SUBSET_NAME subset_name = DEFAULT_SUBSET_NAME
subset = self._extractor subset = self._extractor
task_converters = self.make_task_converters() task_converters = self._make_task_converters()
for task_conv in task_converters.values(): for task_conv in task_converters.values():
task_conv.save_categories(subset) task_conv.save_categories(subset)
for item in subset: for item in subset:
filename = '' filename = ''
if item.has_image: if item.has_image:
filename = str(item.id) + CocoPath.IMAGE_EXT filename = item.image.filename
if self._save_images: if self._save_images:
if item.has_image: if item.has_image:
self.save_image(item, filename) filename = self._save_image(item)
else: else:
log.debug("Item '%s' has no image" % item.id) log.debug("Item '%s' has no image info" % item.id)
for task_conv in task_converters.values(): for task_conv in task_converters.values():
task_conv.save_image_info(item, filename) task_conv.save_image_info(item, filename)
task_conv.save_annotations(item) task_conv.save_annotations(item)

@ -15,7 +15,7 @@ from datumaro.components.extractor import (SourceExtractor,
AnnotationType, Label, RleMask, Points, Polygon, Bbox, Caption, AnnotationType, Label, RleMask, Points, Polygon, Bbox, Caption,
LabelCategories, PointsCategories LabelCategories, PointsCategories
) )
from datumaro.util.image import lazy_image from datumaro.util.image import Image
from .format import CocoTask, CocoPath from .format import CocoTask, CocoPath
@ -117,7 +117,15 @@ class _CocoExtractor(SourceExtractor):
for img_id in loader.getImgIds(): for img_id in loader.getImgIds():
image_info = loader.loadImgs(img_id)[0] image_info = loader.loadImgs(img_id)[0]
image = self._find_image(image_info['file_name']) image_path = self._find_image(image_info['file_name'])
if not image_path:
image_path = image_info['file_name']
image_size = (image_info.get('height'), image_info.get('width'))
if all(image_size):
image_size = (int(image_size[0]), int(image_size[1]))
else:
image_size = None
image = Image(path=image_path, size=image_size)
anns = loader.getAnnIds(imgIds=img_id) anns = loader.getAnnIds(imgIds=img_id)
anns = loader.loadAnns(anns) anns = loader.loadAnns(anns)
@ -232,7 +240,8 @@ class _CocoExtractor(SourceExtractor):
] ]
for image_path in search_paths: for image_path in search_paths:
if osp.exists(image_path): if osp.exists(image_path):
return lazy_image(image_path) return image_path
return None
class CocoImageInfoExtractor(_CocoExtractor): class CocoImageInfoExtractor(_CocoExtractor):
def __init__(self, path, **kwargs): def __init__(self, path, **kwargs):

@ -157,35 +157,45 @@ class _SubsetWriter:
self._writer.open_root() self._writer.open_root()
self._write_meta() self._write_meta()
for item in self._extractor: for index, item in enumerate(self._extractor):
if self._context._save_images: self._write_item(item, index)
if item.has_image:
self._save_image(item)
else:
log.debug("Item '%s' has no image" % item.id)
self._write_item(item)
self._writer.close_root() self._writer.close_root()
def _save_image(self, item): def _save_image(self, item):
image = item.image image = item.image.data
if image is None: if image is None:
return log.warning("Item '%s' has no image" % item.id)
return ''
image_path = osp.join(self._context._images_dir, filename = item.image.filename
str(item.id) + CvatPath.IMAGE_EXT) if filename:
filename = osp.splitext(filename)[0]
else:
filename = item.id
filename += CvatPath.IMAGE_EXT
image_path = osp.join(self._context._images_dir, filename)
save_image(image_path, image) save_image(image_path, image)
return filename
def _write_item(self, item): def _write_item(self, item, index):
h, w = 0, 0 image_info = OrderedDict([
("id", str(_cast(item.id, int, index))),
])
if item.has_image: if item.has_image:
h, w = item.image.shape[:2] size = item.image.size
self._writer.open_image(OrderedDict([ if size:
("id", str(item.id)), h, w = size
("name", str(item.id)), image_info["width"] = str(w)
("width", str(w)), image_info["height"] = str(h)
("height", str(h))
])) filename = item.image.filename
if self._context._save_images:
filename = self._save_image(item)
image_info["name"] = filename
else:
log.debug("Item '%s' has no image info" % item.id)
self._writer.open_image(image_info)
for ann in item.annotations: for ann in item.annotations:
if ann.type in {AnnotationType.points, AnnotationType.polyline, if ann.type in {AnnotationType.points, AnnotationType.polyline,

@ -12,7 +12,7 @@ from datumaro.components.extractor import (SourceExtractor,
AnnotationType, Points, Polygon, PolyLine, Bbox, AnnotationType, Points, Polygon, PolyLine, Bbox,
LabelCategories LabelCategories
) )
from datumaro.util.image import lazy_image from datumaro.util.image import Image
from .format import CvatPath from .format import CvatPath
@ -67,7 +67,7 @@ class CvatExtractor(SourceExtractor):
context = ET.ElementTree.iterparse(path, events=("start", "end")) context = ET.ElementTree.iterparse(path, events=("start", "end"))
context = iter(context) context = iter(context)
categories = cls._parse_meta(context) categories, frame_size = cls._parse_meta(context)
items = OrderedDict() items = OrderedDict()
@ -81,11 +81,15 @@ class CvatExtractor(SourceExtractor):
'id': el.attrib.get('id'), 'id': el.attrib.get('id'),
'label': el.attrib.get('label'), 'label': el.attrib.get('label'),
'group': int(el.attrib.get('group_id', 0)), 'group': int(el.attrib.get('group_id', 0)),
'height': frame_size[0],
'width': frame_size[1],
} }
elif el.tag == 'image': elif el.tag == 'image':
image = { image = {
'name': el.attrib.get('name'), 'name': el.attrib.get('name'),
'frame': el.attrib['id'], 'frame': el.attrib['id'],
'width': el.attrib.get('width'),
'height': el.attrib.get('height'),
} }
elif el.tag in cls._SUPPORTED_SHAPES and (track or image): elif el.tag in cls._SUPPORTED_SHAPES and (track or image):
shape = { shape = {
@ -130,10 +134,7 @@ class CvatExtractor(SourceExtractor):
for pair in el.attrib['points'].split(';'): for pair in el.attrib['points'].split(';'):
shape['points'].extend(map(float, pair.split(','))) shape['points'].extend(map(float, pair.split(',')))
frame_desc = items.get(shape['frame'], { frame_desc = items.get(shape['frame'], {'annotations': []})
'name': shape.get('name'),
'annotations': [],
})
frame_desc['annotations'].append( frame_desc['annotations'].append(
cls._parse_ann(shape, categories)) cls._parse_ann(shape, categories))
items[shape['frame']] = frame_desc items[shape['frame']] = frame_desc
@ -142,6 +143,13 @@ class CvatExtractor(SourceExtractor):
elif el.tag == 'track': elif el.tag == 'track':
track = None track = None
elif el.tag == 'image': elif el.tag == 'image':
frame_desc = items.get(image['frame'], {'annotations': []})
frame_desc.update({
'name': image.get('name'),
'height': image.get('height'),
'width': image.get('width'),
})
items[image['frame']] = frame_desc
image = None image = None
el.clear() el.clear()
@ -155,6 +163,7 @@ class CvatExtractor(SourceExtractor):
categories = {} categories = {}
frame_size = None
has_z_order = False has_z_order = False
mode = 'annotation' mode = 'annotation'
labels = OrderedDict() labels = OrderedDict()
@ -183,6 +192,10 @@ class CvatExtractor(SourceExtractor):
if accepted('annotations', 'meta'): pass if accepted('annotations', 'meta'): pass
elif accepted('meta', 'task'): pass elif accepted('meta', 'task'): pass
elif accepted('task', 'z_order'): pass elif accepted('task', 'z_order'): pass
elif accepted('task', 'original_size'):
frame_size = [None, None]
elif accepted('original_size', 'height', next_state='frame_height'): pass
elif accepted('original_size', 'width', next_state='frame_width'): pass
elif accepted('task', 'labels'): pass elif accepted('task', 'labels'): pass
elif accepted('labels', 'label'): elif accepted('labels', 'label'):
label = { 'name': None, 'attributes': set() } label = { 'name': None, 'attributes': set() }
@ -202,6 +215,11 @@ class CvatExtractor(SourceExtractor):
elif consumed('task', 'task'): pass elif consumed('task', 'task'): pass
elif consumed('z_order', 'z_order'): elif consumed('z_order', 'z_order'):
has_z_order = (el.text == 'True') has_z_order = (el.text == 'True')
elif consumed('original_size', 'original_size'): pass
elif consumed('frame_height', 'height'):
frame_size[0] = int(el.text)
elif consumed('frame_width', 'width'):
frame_size[1] = int(el.text)
elif consumed('label_name', 'name'): elif consumed('label_name', 'name'):
label['name'] = el.text label['name'] = el.text
elif consumed('attr_name', 'name'): elif consumed('attr_name', 'name'):
@ -231,7 +249,7 @@ class CvatExtractor(SourceExtractor):
categories[AnnotationType.label] = label_cat categories[AnnotationType.label] = label_cat
return categories return categories, frame_size
@classmethod @classmethod
def _parse_ann(cls, ann, categories): def _parse_ann(cls, ann, categories):
@ -277,27 +295,37 @@ class CvatExtractor(SourceExtractor):
raise NotImplementedError("Unknown annotation type '%s'" % ann_type) raise NotImplementedError("Unknown annotation type '%s'" % ann_type)
def _load_items(self, parsed): def _load_items(self, parsed):
for item_id, item_desc in parsed.items(): for frame_id, item_desc in parsed.items():
file_name = item_desc.get('name') filename = item_desc.get('name')
if not file_name: if filename:
file_name = item_id filename = self._find_image(filename)
image = self._find_image(file_name) if not filename:
filename = item_desc.get('name')
parsed[item_id] = DatasetItem(id=item_id, subset=self._subset, image_size = (item_desc.get('height'), item_desc.get('width'))
image=image, annotations=item_desc.get('annotations', None)) if all(image_size):
image_size = (int(image_size[0]), int(image_size[1]))
else:
image_size = None
image = None
if filename:
image = Image(path=filename, size=image_size)
parsed[frame_id] = DatasetItem(id=frame_id, subset=self._subset,
image=image, annotations=item_desc.get('annotations'))
return parsed return parsed
def _find_image(self, file_name): def _find_image(self, file_name):
search_paths = [ search_paths = []
osp.join(osp.dirname(self._path), file_name)
]
if self._images_dir: if self._images_dir:
search_paths += [ search_paths += [
osp.join(self._images_dir, file_name), osp.join(self._images_dir, file_name),
osp.join(self._images_dir, self._subset or DEFAULT_SUBSET_NAME, osp.join(self._images_dir, self._subset or DEFAULT_SUBSET_NAME,
file_name), file_name),
] ]
search_paths += [
osp.join(osp.dirname(self._path), file_name)
]
for image_path in search_paths: for image_path in search_paths:
if osp.isfile(image_path): if osp.isfile(image_path):
return lazy_image(image_path) return image_path
return None return None

@ -30,9 +30,9 @@ def _cast(value, type_conv, default=None):
return default return default
class _SubsetWriter: class _SubsetWriter:
def __init__(self, name, converter): def __init__(self, name, context):
self._name = name self._name = name
self._converter = converter self._context = context
self._data = { self._data = {
'info': {}, 'info': {},
@ -58,6 +58,15 @@ class _SubsetWriter:
} }
if item.path: if item.path:
item_desc['path'] = item.path item_desc['path'] = item.path
if item.has_image:
path = item.image.path
if self._context._save_images:
path = self._context._save_image(item)
item_desc['image'] = {
'size': item.image.size,
'path': path,
}
self.items.append(item_desc) self.items.append(item_desc)
for ann in item.annotations: for ann in item.annotations:
@ -123,7 +132,7 @@ class _SubsetWriter:
self._next_mask_id += 1 self._next_mask_id += 1
filename = '%d%s' % (mask_id, DatumaroPath.MASK_EXT) filename = '%d%s' % (mask_id, DatumaroPath.MASK_EXT)
masks_dir = osp.join(self._converter._annotations_dir, masks_dir = osp.join(self._context._annotations_dir,
DatumaroPath.MASKS_DIR) DatumaroPath.MASKS_DIR)
os.makedirs(masks_dir, exist_ok=True) os.makedirs(masks_dir, exist_ok=True)
path = osp.join(masks_dir, filename) path = osp.join(masks_dir, filename)
@ -226,7 +235,7 @@ class _SubsetWriter:
return converted return converted
class _Converter: class _Converter:
def __init__(self, extractor, save_dir, save_images=False,): def __init__(self, extractor, save_dir, save_images=False):
self._extractor = extractor self._extractor = extractor
self._save_dir = save_dir self._save_dir = save_dir
self._save_images = save_images self._save_images = save_images
@ -257,21 +266,25 @@ class _Converter:
subset = DEFAULT_SUBSET_NAME subset = DEFAULT_SUBSET_NAME
writer = subsets[subset] writer = subsets[subset]
if self._save_images:
self._save_image(item)
writer.write_item(item) writer.write_item(item)
for subset, writer in subsets.items(): for subset, writer in subsets.items():
writer.write(annotations_dir) writer.write(annotations_dir)
def _save_image(self, item): def _save_image(self, item):
image = item.image image = item.image.data
if image is None: if image is None:
return return ''
image_path = osp.join(self._images_dir, filename = item.image.filename
str(item.id) + DatumaroPath.IMAGE_EXT) if filename:
filename = osp.splitext(filename)[0]
else:
filename = item.id
filename += DatumaroPath.IMAGE_EXT
image_path = osp.join(self._images_dir, filename)
save_image(image_path, image) save_image(image_path, image)
return filename
class DatumaroConverter(Converter, CliPlugin): class DatumaroConverter(Converter, CliPlugin):
@classmethod @classmethod

@ -12,7 +12,7 @@ from datumaro.components.extractor import (SourceExtractor,
AnnotationType, Label, Mask, Points, Polygon, PolyLine, Bbox, Caption, AnnotationType, Label, Mask, Points, Polygon, PolyLine, Bbox, Caption,
LabelCategories, MaskCategories, PointsCategories LabelCategories, MaskCategories, PointsCategories
) )
from datumaro.util.image import lazy_image from datumaro.util.image import Image
from datumaro.util.mask_tools import lazy_mask from datumaro.util.mask_tools import lazy_mask
from .format import DatumaroPath from .format import DatumaroPath
@ -93,11 +93,13 @@ class DatumaroExtractor(SourceExtractor):
items = [] items = []
for item_desc in parsed['items']: for item_desc in parsed['items']:
item_id = item_desc['id'] item_id = item_desc['id']
image = None image = None
image_path = osp.join(self._path, DatumaroPath.IMAGES_DIR, image_info = item_desc.get('image', {})
item_id + DatumaroPath.IMAGE_EXT) if image_info:
if osp.exists(image_path): image_path = osp.join(self._path, DatumaroPath.IMAGES_DIR,
image = lazy_image(image_path) image_info.get('path', '')) # relative or absolute fits
image = Image(path=image_path, size=image_info.get('size'))
annotations = self._load_annotations(item_desc) annotations = self._load_annotations(item_desc)

@ -8,7 +8,8 @@ import os
import os.path as osp import os.path as osp
from datumaro.components.extractor import DatasetItem, SourceExtractor, Importer from datumaro.components.extractor import DatasetItem, SourceExtractor, Importer
from datumaro.util.image import lazy_image from datumaro.components.converter import Converter
from datumaro.util.image import save_image
class ImageDirImporter(Importer): class ImageDirImporter(Importer):
@ -44,7 +45,7 @@ class ImageDirExtractor(SourceExtractor):
path = osp.join(url, name) path = osp.join(url, name)
if self._is_image(path): if self._is_image(path):
item_id = osp.splitext(name)[0] item_id = osp.splitext(name)[0]
item = DatasetItem(id=item_id, image=lazy_image(path)) item = DatasetItem(id=item_id, image=path)
items.append((item.id, item)) items.append((item.id, item))
items = sorted(items, key=lambda e: e[0]) items = sorted(items, key=lambda e: e[0])
@ -73,3 +74,18 @@ class ImageDirExtractor(SourceExtractor):
if osp.isfile(path) and path.endswith(ext): if osp.isfile(path) and path.endswith(ext):
return True return True
return False return False
class ImageDirConverter(Converter):
def __call__(self, extractor, save_dir):
os.makedirs(save_dir, exist_ok=True)
for item in extractor:
if item.has_image and item.image.has_data:
filename = item.image.filename
if filename:
filename = osp.splitext(filename)[0]
else:
filename = item.id
filename += '.jpg'
save_image(osp.join(save_dir, filename), item.image.data)

@ -10,7 +10,9 @@ import os
import os.path as osp import os.path as osp
import string import string
from datumaro.components.extractor import AnnotationType, DEFAULT_SUBSET_NAME from datumaro.components.extractor import (AnnotationType, DEFAULT_SUBSET_NAME,
LabelCategories
)
from datumaro.components.converter import Converter from datumaro.components.converter import Converter
from datumaro.components.cli_plugin import CliPlugin from datumaro.components.cli_plugin import CliPlugin
from datumaro.util.image import encode_image from datumaro.util.image import encode_image
@ -49,26 +51,30 @@ def _make_tf_example(item, get_label_id, get_label, save_images=False):
} }
if not item.has_image: if not item.has_image:
raise Exception( raise Exception("Failed to export dataset item '%s': "
"Failed to export dataset item '%s': item has no image" % item.id) "item has no image info" % item.id)
height, width = item.image.shape[:2] height, width = item.image.size
features.update({ features.update({
'image/height': int64_feature(height), 'image/height': int64_feature(height),
'image/width': int64_feature(width), 'image/width': int64_feature(width),
}) })
features.update({
'image/encoded': bytes_feature(b''),
'image/format': bytes_feature(b'')
})
if save_images: if save_images:
if item.has_image: if item.has_image and item.image.has_data:
fmt = DetectionApiPath.IMAGE_FORMAT fmt = DetectionApiPath.IMAGE_FORMAT
buffer = encode_image(item.image, DetectionApiPath.IMAGE_EXT) buffer = encode_image(item.image.data, DetectionApiPath.IMAGE_EXT)
features.update({ features.update({
'image/encoded': bytes_feature(buffer), 'image/encoded': bytes_feature(buffer),
'image/format': bytes_feature(fmt.encode('utf-8')), 'image/format': bytes_feature(fmt.encode('utf-8')),
}) })
else: else:
log.debug("Item '%s' has no image" % item.id) log.warning("Item '%s' has no image" % item.id)
xmins = [] # List of normalized left x coordinates in bounding box (1 per box) xmins = [] # List of normalized left x coordinates in bounding box (1 per box)
xmaxs = [] # List of normalized right x coordinates in bounding box (1 per box) xmaxs = [] # List of normalized right x coordinates in bounding box (1 per box)
@ -130,7 +136,8 @@ class TfDetectionApiConverter(Converter, CliPlugin):
subset_name = DEFAULT_SUBSET_NAME subset_name = DEFAULT_SUBSET_NAME
subset = extractor subset = extractor
label_categories = subset.categories()[AnnotationType.label] label_categories = subset.categories().get(AnnotationType.label,
LabelCategories())
get_label = lambda label_id: label_categories.items[label_id].name \ get_label = lambda label_id: label_categories.items[label_id].name \
if label_id is not None else '' if label_id is not None else ''
label_ids = OrderedDict((label.name, 1 + idx) label_ids = OrderedDict((label.name, 1 + idx)

@ -12,7 +12,7 @@ from datumaro.components.extractor import (SourceExtractor,
DEFAULT_SUBSET_NAME, DatasetItem, DEFAULT_SUBSET_NAME, DatasetItem,
AnnotationType, Bbox, LabelCategories AnnotationType, Bbox, LabelCategories
) )
from datumaro.util.image import lazy_image, decode_image from datumaro.util.image import Image, decode_image, lazy_image
from datumaro.util.tf_util import import_tf as _import_tf from datumaro.util.tf_util import import_tf as _import_tf
from .format import DetectionApiPath from .format import DetectionApiPath
@ -163,24 +163,29 @@ class TfDetectionApiExtractor(SourceExtractor):
item_id = osp.splitext(frame_filename)[0] item_id = osp.splitext(frame_filename)[0]
annotations = [] annotations = []
for index, shape in enumerate( for shape in np.dstack((labels, xmins, ymins, xmaxs, ymaxs))[0]:
np.dstack((labels, xmins, ymins, xmaxs, ymaxs))[0]):
label = shape[0].decode('utf-8') label = shape[0].decode('utf-8')
x = clamp(shape[1] * frame_width, 0, frame_width) x = clamp(shape[1] * frame_width, 0, frame_width)
y = clamp(shape[2] * frame_height, 0, frame_height) y = clamp(shape[2] * frame_height, 0, frame_height)
w = clamp(shape[3] * frame_width, 0, frame_width) - x w = clamp(shape[3] * frame_width, 0, frame_width) - x
h = clamp(shape[4] * frame_height, 0, frame_height) - y h = clamp(shape[4] * frame_height, 0, frame_height) - y
annotations.append(Bbox(x, y, w, h, annotations.append(Bbox(x, y, w, h,
label=dataset_labels.get(label, None), id=index label=dataset_labels.get(label)
)) ))
image_size = None
if frame_height and frame_width:
image_size = (frame_height, frame_width)
image_params = {}
if frame_image and frame_format:
image_params['data'] = lazy_image(frame_image, decode_image)
if frame_filename and images_dir:
image_params['path'] = osp.join(images_dir, frame_filename)
image = None image = None
if image is None and frame_image and frame_format: if image_params:
image = lazy_image(frame_image, loader=decode_image) image = Image(**image_params, size=image_size)
if image is None and frame_filename and images_dir:
image_path = osp.join(images_dir, frame_filename)
if osp.exists(image_path):
image = lazy_image(image_path)
dataset_items.append(DatasetItem(id=item_id, subset=subset_name, dataset_items.append(DatasetItem(id=item_id, subset=subset_name,
image=image, annotations=annotations)) image=image, annotations=annotations))

@ -5,6 +5,7 @@
from itertools import groupby from itertools import groupby
import logging as log import logging as log
import os.path as osp
import pycocotools.mask as mask_utils import pycocotools.mask as mask_utils
@ -28,7 +29,7 @@ class CropCoveredSegments(Transform, CliPlugin):
if not item.has_image: if not item.has_image:
raise Exception("Image info is required for this transform") raise Exception("Image info is required for this transform")
h, w = item.image.shape[:2] h, w = item.image.size
segments = self.crop_segments(segments, w, h) segments = self.crop_segments(segments, w, h)
annotations += segments annotations += segments
@ -107,7 +108,7 @@ class MergeInstanceSegments(Transform, CliPlugin):
if not item.has_image: if not item.has_image:
raise Exception("Image info is required for this transform") raise Exception("Image info is required for this transform")
h, w = item.image.shape[:2] h, w = item.image.size
instances = self.find_instances(segments) instances = self.find_instances(segments)
segments = [self.merge_segments(i, w, h, self._include_polygons) segments = [self.merge_segments(i, w, h, self._include_polygons)
for i in instances] for i in instances]
@ -196,7 +197,7 @@ class PolygonsToMasks(Transform, CliPlugin):
if ann.type == AnnotationType.polygon: if ann.type == AnnotationType.polygon:
if not item.has_image: if not item.has_image:
raise Exception("Image info is required for this transform") raise Exception("Image info is required for this transform")
h, w = item.image.shape[:2] h, w = item.image.size
annotations.append(self.convert_polygon(ann, h, w)) annotations.append(self.convert_polygon(ann, h, w))
else: else:
annotations.append(ann) annotations.append(ann)
@ -273,7 +274,6 @@ class Reindex(Transform, CliPlugin):
for i, item in enumerate(self._extractor): for i, item in enumerate(self._extractor):
yield self.wrap_item(item, id=i + self._start) yield self.wrap_item(item, id=i + self._start)
class MapSubsets(Transform, CliPlugin): class MapSubsets(Transform, CliPlugin):
@staticmethod @staticmethod
def _mapping_arg(s): def _mapping_arg(s):
@ -303,3 +303,10 @@ class MapSubsets(Transform, CliPlugin):
def transform_item(self, item): def transform_item(self, item):
return self.wrap_item(item, return self.wrap_item(item,
subset=self._mapping.get(item.subset, item.subset)) subset=self._mapping.get(item.subset, item.subset))
class IdFromImageName(Transform, CliPlugin):
def transform_item(self, item):
name = item.id
if item.has_image and item.image.filename:
name = osp.splitext(item.image.filename)[0]
return self.wrap_item(item, id=name)

@ -135,11 +135,18 @@ class _Converter:
for item in subset: for item in subset:
log.debug("Converting item '%s'", item.id) log.debug("Converting item '%s'", item.id)
image_filename = ''
if item.has_image:
image_filename = item.image.filename
if self._save_images: if self._save_images:
if item.has_image: if item.has_image and item.image.has_data:
save_image(osp.join(self._images_dir, if image_filename:
item.id + VocPath.IMAGE_EXT), image_filename = osp.splitext(image_filename)[0]
item.image) else:
image_filename = item.id
image_filename += VocPath.IMAGE_EXT
save_image(osp.join(self._images_dir, image_filename),
item.image.data)
else: else:
log.debug("Item '%s' has no image" % item.id) log.debug("Item '%s' has no image" % item.id)
@ -161,8 +168,7 @@ class _Converter:
else: else:
folder = '' folder = ''
ET.SubElement(root_elem, 'folder').text = folder ET.SubElement(root_elem, 'folder').text = folder
ET.SubElement(root_elem, 'filename').text = \ ET.SubElement(root_elem, 'filename').text = image_filename
item.id + VocPath.IMAGE_EXT
source_elem = ET.SubElement(root_elem, 'source') source_elem = ET.SubElement(root_elem, 'source')
ET.SubElement(source_elem, 'database').text = 'Unknown' ET.SubElement(source_elem, 'database').text = 'Unknown'
@ -170,9 +176,12 @@ class _Converter:
ET.SubElement(source_elem, 'image').text = 'Unknown' ET.SubElement(source_elem, 'image').text = 'Unknown'
if item.has_image: if item.has_image:
image_shape = item.image.shape h, w = item.image.size
h, w = image_shape[:2] if item.image.has_data:
c = 1 if len(image_shape) == 2 else image_shape[2] image_shape = item.image.data.shape
c = 1 if len(image_shape) == 2 else image_shape[2]
else:
c = 3
size_elem = ET.SubElement(root_elem, 'size') size_elem = ET.SubElement(root_elem, 'size')
ET.SubElement(size_elem, 'width').text = str(w) ET.SubElement(size_elem, 'width').text = str(w)
ET.SubElement(size_elem, 'height').text = str(h) ET.SubElement(size_elem, 'height').text = str(h)
@ -467,7 +476,8 @@ class _Converter:
def _make_label_id_map(self): def _make_label_id_map(self):
source_labels = { source_labels = {
id: label.name for id, label in id: label.name for id, label in
enumerate(self._extractor.categories()[AnnotationType.label].items) enumerate(self._extractor.categories().get(
AnnotationType.label, LabelCategories()).items)
} }
target_labels = { target_labels = {
label.name: id for id, label in label.name: id for id, label in
@ -510,7 +520,11 @@ class VocConverter(Converter, CliPlugin):
def _get_labelmap(s): def _get_labelmap(s):
if osp.isfile(s): if osp.isfile(s):
return s return s
return LabelmapType[s].name try:
return LabelmapType[s].name
except KeyError:
import argparse
raise argparse.ArgumentTypeError()
@classmethod @classmethod
def build_cmdline_parser(cls, **kwargs): def build_cmdline_parser(cls, **kwargs):

@ -86,6 +86,8 @@ class VocExtractor(SourceExtractor):
ann_file_data = f.read() ann_file_data = f.read()
ann_file_root = ET.fromstring(ann_file_data) ann_file_root = ET.fromstring(ann_file_data)
item = ann_file_root.find('filename').text item = ann_file_root.find('filename').text
if not item:
item = ann_item
item = osp.splitext(item)[0] item = osp.splitext(item)[0]
det_annotations[item] = ann_file_data det_annotations[item] = ann_file_data
@ -130,11 +132,8 @@ class VocExtractor(SourceExtractor):
yield item yield item
def _get(self, item_id, subset_name): def _get(self, item_id, subset_name):
image = None image = osp.join(self._path, VocPath.IMAGES_DIR,
image_path = osp.join(self._path, VocPath.IMAGES_DIR,
item_id + VocPath.IMAGE_EXT) item_id + VocPath.IMAGE_EXT)
if osp.isfile(image_path):
image = lazy_image(image_path)
annotations = self._get_annotations(item_id) annotations = self._get_annotations(item_id)

@ -75,17 +75,24 @@ class YoloConverter(Converter, CliPlugin):
image_paths = OrderedDict() image_paths = OrderedDict()
for item in subset: for item in subset:
image_name = '%s.jpg' % item.id if not item.has_image:
image_paths[item.id] = osp.join('data', raise Exception("Failed to export item '%s': "
osp.basename(subset_dir), image_name) "item has no image info" % item.id)
height, width = item.image.size
image_name = item.image.filename
item_name = osp.splitext(item.image.filename)[0]
if self._save_images: if self._save_images:
if item.has_image: if item.has_image and item.image.has_data:
save_image(osp.join(subset_dir, image_name), item.image) if not item_name:
item_name = item.id
image_name = item_name + '.jpg'
save_image(osp.join(subset_dir, image_name),
item.image.data)
else: else:
log.debug("Item '%s' has no images" % item.id) log.warning("Item '%s' has no image" % item.id)
image_paths[item.id] = osp.join('data',
height, width = item.image.shape[:2] osp.basename(subset_dir), image_name)
yolo_annotation = '' yolo_annotation = ''
for bbox in item.annotations: for bbox in item.annotations:
@ -98,7 +105,7 @@ class YoloConverter(Converter, CliPlugin):
yolo_bb = ' '.join('%.6f' % p for p in yolo_bb) yolo_bb = ' '.join('%.6f' % p for p in yolo_bb)
yolo_annotation += '%s %s\n' % (bbox.label, yolo_bb) yolo_annotation += '%s %s\n' % (bbox.label, yolo_bb)
annotation_path = osp.join(subset_dir, '%s.txt' % item.id) annotation_path = osp.join(subset_dir, '%s.txt' % item_name)
with open(annotation_path, 'w') as f: with open(annotation_path, 'w') as f:
f.write(yolo_annotation) f.write(yolo_annotation)

@ -10,7 +10,7 @@ import re
from datumaro.components.extractor import (SourceExtractor, Extractor, from datumaro.components.extractor import (SourceExtractor, Extractor,
DatasetItem, AnnotationType, Bbox, LabelCategories DatasetItem, AnnotationType, Bbox, LabelCategories
) )
from datumaro.util.image import lazy_image from datumaro.util.image import Image
from .format import YoloPath from .format import YoloPath
@ -33,16 +33,31 @@ class YoloExtractor(SourceExtractor):
def categories(self): def categories(self):
return self._parent.categories() return self._parent.categories()
def __init__(self, config_path): def __init__(self, config_path, image_info=None):
super().__init__() super().__init__()
if not osp.isfile(config_path): if not osp.isfile(config_path):
raise Exception("Can't read dataset descriptor file '%s'" % \ raise Exception("Can't read dataset descriptor file '%s'" %
config_path) config_path)
rootpath = osp.dirname(config_path) rootpath = osp.dirname(config_path)
self._path = rootpath self._path = rootpath
assert image_info is None or isinstance(image_info, (str, dict))
if image_info is None:
image_info = osp.join(rootpath, YoloPath.IMAGE_META_FILE)
if not osp.isfile(image_info):
image_info = {}
if isinstance(image_info, str):
if not osp.isfile(image_info):
raise Exception("Can't read image meta file '%s'" % image_info)
with open(image_info) as f:
image_info = {}
for line in f:
image_name, h, w = line.strip().split()
image_info[image_name] = (int(h), int(w))
self._image_info = image_info
with open(config_path, 'r') as f: with open(config_path, 'r') as f:
config_lines = f.readlines() config_lines = f.readlines()
@ -77,10 +92,10 @@ class YoloExtractor(SourceExtractor):
subset.items = OrderedDict( subset.items = OrderedDict(
(osp.splitext(osp.basename(p))[0], p.strip()) for p in f) (osp.splitext(osp.basename(p))[0], p.strip()) for p in f)
for image_path in subset.items.values(): for item_id, image_path in subset.items.items():
image_path = self._make_local_path(image_path) image_path = self._make_local_path(image_path)
if not osp.isfile(image_path): if not osp.isfile(image_path) and item_id not in image_info:
raise Exception("Can't find image '%s'" % image_path) raise Exception("Can't find image '%s'" % item_id)
subsets[subset_name] = subset subsets[subset_name] = subset
@ -103,8 +118,10 @@ class YoloExtractor(SourceExtractor):
if isinstance(item, str): if isinstance(item, str):
image_path = self._make_local_path(item) image_path = self._make_local_path(item)
image = lazy_image(image_path) image_size = self._image_info.get(item_id)
h, w, _ = image().shape image = Image(path=image_path, size=image_size)
h, w = image.size
anno_path = osp.splitext(image_path)[0] + '.txt' anno_path = osp.splitext(image_path)[0] + '.txt'
annotations = self._parse_annotations(anno_path, w, h) annotations = self._parse_annotations(anno_path, w, h)

@ -7,3 +7,5 @@
class YoloPath: class YoloPath:
DEFAULT_SUBSET_NAME = 'train' DEFAULT_SUBSET_NAME = 'train'
SUBSET_NAMES = ['train', 'valid'] SUBSET_NAMES = ['train', 'valid']
IMAGE_META_FILE = 'images.meta'

@ -15,18 +15,20 @@ class YoloImporter(Importer):
from datumaro.components.project import Project # cyclic import from datumaro.components.project import Project # cyclic import
project = Project() project = Project()
if not osp.exists(path):
raise Exception("Failed to find 'yolo' dataset at '%s'" % path)
if path.endswith('.data') and osp.isfile(path): if path.endswith('.data') and osp.isfile(path):
config_paths = [path] config_paths = [path]
else: else:
config_paths = glob(osp.join(path, '*.data')) config_paths = glob(osp.join(path, '*.data'))
if not osp.exists(path) or not config_paths:
raise Exception("Failed to find 'yolo' dataset at '%s'" % path)
for config_path in config_paths: for config_path in config_paths:
log.info("Found a dataset at '%s'" % config_path) log.info("Found a dataset at '%s'" % config_path)
source_name = osp.splitext(osp.basename(config_path))[0] source_name = '%s_%s' % (
osp.basename(osp.dirname(config_path)),
osp.splitext(osp.basename(config_path))[0])
project.add_source(source_name, { project.add_source(source_name, {
'url': config_path, 'url': config_path,
'format': 'yolo', 'format': 'yolo',

@ -7,6 +7,7 @@
from io import BytesIO from io import BytesIO
import numpy as np import numpy as np
import os.path as osp
from enum import Enum from enum import Enum
_IMAGE_BACKENDS = Enum('_IMAGE_BACKENDS', ['cv2', 'PIL']) _IMAGE_BACKENDS = Enum('_IMAGE_BACKENDS', ['cv2', 'PIL'])
@ -123,14 +124,16 @@ def decode_image(image_bytes):
class lazy_image: class lazy_image:
def __init__(self, path, loader=load_image, cache=None): def __init__(self, path, loader=None, cache=None):
if loader is None:
loader = load_image
self.path = path self.path = path
self.loader = loader self.loader = loader
# Cache: # Cache:
# - False: do not cache # - False: do not cache
# - None: use default (don't store in a class variable) # - None: use the global cache
# - object: use this object as a cache # - object: an object to be used as cache
assert cache in {None, False} or isinstance(cache, object) assert cache in {None, False} or isinstance(cache, object)
self.cache = cache self.cache = cache
@ -138,9 +141,9 @@ class lazy_image:
image = None image = None
image_id = hash(self) # path is not necessary hashable or a file path image_id = hash(self) # path is not necessary hashable or a file path
cache = self._get_cache() cache = self._get_cache(self.cache)
if cache is not None: if cache is not None:
image = self._get_cache().get(image_id) image = cache.get(image_id)
if image is None: if image is None:
image = self.loader(self.path) image = self.loader(self.path)
@ -148,8 +151,8 @@ class lazy_image:
cache.push(image_id, image) cache.push(image_id, image)
return image return image
def _get_cache(self): @staticmethod
cache = self.cache def _get_cache(cache):
if cache is None: if cache is None:
cache = _ImageCache.get_instance() cache = _ImageCache.get_instance()
elif cache == False: elif cache == False:
@ -158,3 +161,63 @@ class lazy_image:
def __hash__(self): def __hash__(self):
return hash((id(self), self.path, self.loader)) return hash((id(self), self.path, self.loader))
class Image:
def __init__(self, data=None, path=None, loader=None, cache=None,
size=None):
assert size is None or len(size) == 2
if size is not None:
assert len(size) == 2 and 0 < size[0] and 0 < size[1], size
size = tuple(size)
else:
size = None
self._size = size # (H, W)
assert path is None or isinstance(path, str)
if path is None:
path = ''
self._path = path
assert data is not None or path, "Image can not be empty"
if data is None and path:
if osp.isfile(path):
data = lazy_image(path, loader=loader, cache=cache)
self._data = data
@property
def path(self):
return self._path
@property
def filename(self):
return osp.basename(self._path)
@property
def data(self):
if callable(self._data):
return self._data()
return self._data
@property
def has_data(self):
return self._data is not None
@property
def size(self):
if self._size is None:
data = self.data
if data is not None:
self._size = data.shape[:2]
return self._size
def __eq__(self, other):
if isinstance(other, np.ndarray):
return self.has_data and np.array_equal(self.data, other)
if not isinstance(other, __class__):
return False
return \
(np.array_equal(self.size, other.size)) and \
(self.has_data == other.has_data) and \
(self.has_data and np.array_equal(self.data, other.data) or \
not self.has_data)

@ -19,8 +19,7 @@ from datumaro.plugins.coco_format.converter import (
CocoLabelsConverter, CocoLabelsConverter,
) )
from datumaro.plugins.coco_format.importer import CocoImporter from datumaro.plugins.coco_format.importer import CocoImporter
from datumaro.util.image import save_image from datumaro.util.image import save_image, Image
from datumaro.util import find
from datumaro.util.test_utils import TestDir, compare_datasets from datumaro.util.test_utils import TestDir, compare_datasets
@ -100,7 +99,7 @@ class CocoImporterTest(TestCase):
os.makedirs(img_dir) os.makedirs(img_dir)
os.makedirs(ann_dir) os.makedirs(ann_dir)
image = np.ones((10, 5, 3), dtype=np.uint8) image = np.ones((10, 5, 3))
save_image(osp.join(img_dir, '000000000001.jpg'), image) save_image(osp.join(img_dir, '000000000001.jpg'), image)
annotation = self.generate_annotation() annotation = self.generate_annotation()
@ -109,30 +108,33 @@ class CocoImporterTest(TestCase):
json.dump(annotation, outfile) json.dump(annotation, outfile)
def test_can_import(self): def test_can_import(self):
with TestDir() as test_dir: class DstExtractor(Extractor):
self.COCO_dataset_generate(test_dir) def __iter__(self):
project = Project.import_from(test_dir, 'coco') return iter([
dataset = project.make_dataset() DatasetItem(id=1, image=np.ones((10, 5, 3)), subset='val',
annotations=[
Polygon([0, 0, 1, 0, 1, 2, 0, 2], label=0,
id=1, group=1, attributes={'is_crowd': False}),
Mask(np.array(
[[1, 0, 0, 1, 0]] * 5 +
[[1, 1, 1, 1, 0]] * 5
), label=0,
id=2, group=2, attributes={'is_crowd': True}),
]
),
])
self.assertListEqual(['val'], sorted(dataset.subsets())) def categories(self):
self.assertEqual(1, len(dataset)) label_cat = LabelCategories()
label_cat.add('TEST')
return { AnnotationType.label: label_cat }
item = next(iter(dataset)) with TestDir() as test_dir:
self.assertTrue(item.has_image) self.COCO_dataset_generate(test_dir)
self.assertEqual(np.sum(item.image), np.prod(item.image.shape))
self.assertEqual(2, len(item.annotations))
ann_1 = find(item.annotations, lambda x: x.id == 1) dataset = Project.import_from(test_dir, 'coco').make_dataset()
ann_1_poly = find(item.annotations, lambda x: \
x.group == ann_1.id and x.type == AnnotationType.polygon)
self.assertFalse(ann_1 is None)
self.assertFalse(ann_1_poly is None)
ann_2 = find(item.annotations, lambda x: x.id == 2) compare_datasets(self, DstExtractor(), dataset)
ann_2_mask = find(item.annotations, lambda x: \
x.group == ann_2.id and x.type == AnnotationType.mask)
self.assertFalse(ann_2 is None)
self.assertFalse(ann_2_mask is None)
class CocoConverterTest(TestCase): class CocoConverterTest(TestCase):
def _test_save_and_load(self, source_dataset, converter, test_dir, def _test_save_and_load(self, source_dataset, converter, test_dir,
@ -633,3 +635,14 @@ class CocoConverterTest(TestCase):
with TestDir() as test_dir: with TestDir() as test_dir:
self._test_save_and_load(TestExtractor(), self._test_save_and_load(TestExtractor(),
CocoConverter(), test_dir) CocoConverter(), test_dir)
def test_can_save_dataset_with_image_info(self):
class TestExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=1, image=Image(path='1.jpg', size=(10, 15))),
])
with TestDir() as test_dir:
self._test_save_and_load(TestExtractor(),
CocoConverter(), test_dir)

@ -12,7 +12,7 @@ from datumaro.components.extractor import (Extractor, DatasetItem,
from datumaro.plugins.cvat_format.importer import CvatImporter from datumaro.plugins.cvat_format.importer import CvatImporter
from datumaro.plugins.cvat_format.converter import CvatConverter from datumaro.plugins.cvat_format.converter import CvatConverter
from datumaro.plugins.cvat_format.format import CvatPath from datumaro.plugins.cvat_format.format import CvatPath
from datumaro.util.image import save_image from datumaro.util.image import save_image, Image
from datumaro.util.test_utils import TestDir, compare_datasets from datumaro.util.test_utils import TestDir, compare_datasets
@ -82,7 +82,7 @@ class CvatExtractorTest(TestCase):
save_image(osp.join(images_dir, 'img1.jpg'), np.ones((10, 10, 3))) save_image(osp.join(images_dir, 'img1.jpg'), np.ones((10, 10, 3)))
item2_elem = ET.SubElement(root_elem, 'image') item2_elem = ET.SubElement(root_elem, 'image')
item2_elem.attrib.update({ item2_elem.attrib.update({
'id': '1', 'name': 'img1', 'width': '8', 'height': '8' 'id': '1', 'name': 'img1', 'width': '10', 'height': '10'
}) })
item2_ann1_elem = ET.SubElement(item2_elem, 'polygon') item2_ann1_elem = ET.SubElement(item2_elem, 'polygon')
@ -159,7 +159,7 @@ class CvatConverterTest(TestCase):
label_categories.items[2].attributes.update(['a1', 'a2']) label_categories.items[2].attributes.update(['a1', 'a2'])
label_categories.attributes.update(['z_order', 'occluded']) label_categories.attributes.update(['z_order', 'occluded'])
class SrcTestExtractor(Extractor): class SrcExtractor(Extractor):
def __iter__(self): def __iter__(self):
return iter([ return iter([
DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)), DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)),
@ -192,12 +192,15 @@ class CvatConverterTest(TestCase):
PolyLine([5, 0, 9, 0, 5, 5]), # will be skipped as no label PolyLine([5, 0, 9, 0, 5, 5]), # will be skipped as no label
] ]
), ),
DatasetItem(id=3, subset='s3', image=Image(
path='3.jpg', size=(2, 4))),
]) ])
def categories(self): def categories(self):
return { AnnotationType.label: label_categories } return { AnnotationType.label: label_categories }
class DstTestExtractor(Extractor): class DstExtractor(Extractor):
def __iter__(self): def __iter__(self):
return iter([ return iter([
DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)), DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)),
@ -232,12 +235,15 @@ class CvatConverterTest(TestCase):
attributes={ 'z_order': 1, 'occluded': False }), attributes={ 'z_order': 1, 'occluded': False }),
] ]
), ),
DatasetItem(id=3, subset='s3', image=Image(
path='3.jpg', size=(2, 4))),
]) ])
def categories(self): def categories(self):
return { AnnotationType.label: label_categories } return { AnnotationType.label: label_categories }
with TestDir() as test_dir: with TestDir() as test_dir:
self._test_save_and_load(SrcTestExtractor(), self._test_save_and_load(SrcExtractor(),
CvatConverter(save_images=True), test_dir, CvatConverter(save_images=True), test_dir,
target_dataset=DstTestExtractor()) target_dataset=DstExtractor())

@ -9,8 +9,9 @@ from datumaro.components.extractor import (Extractor, DatasetItem,
LabelCategories, MaskCategories, PointsCategories LabelCategories, MaskCategories, PointsCategories
) )
from datumaro.plugins.datumaro_format.converter import DatumaroConverter from datumaro.plugins.datumaro_format.converter import DatumaroConverter
from datumaro.util.test_utils import TestDir, item_to_str
from datumaro.util.mask_tools import generate_colormap from datumaro.util.mask_tools import generate_colormap
from datumaro.util.image import Image
from datumaro.util.test_utils import TestDir, item_to_str
class DatumaroConverterTest(TestCase): class DatumaroConverterTest(TestCase):
@ -48,7 +49,7 @@ class DatumaroConverterTest(TestCase):
DatasetItem(id=42, subset='test'), DatasetItem(id=42, subset='test'),
DatasetItem(id=42), DatasetItem(id=42),
DatasetItem(id=43), DatasetItem(id=43, image=Image(path='1/b/c.qq', size=(2, 4))),
]) ])
def categories(self): def categories(self):

@ -8,7 +8,7 @@ import datumaro.util.image as image_module
from datumaro.util.test_utils import TestDir from datumaro.util.test_utils import TestDir
class ImageTest(TestCase): class ImageOperationsTest(TestCase):
def setUp(self): def setUp(self):
self.default_backend = image_module._IMAGE_BACKEND self.default_backend = image_module._IMAGE_BACKEND

@ -1,12 +1,11 @@
import numpy as np import numpy as np
import os.path as osp
from unittest import TestCase from unittest import TestCase
from datumaro.components.project import Project from datumaro.components.project import Project
from datumaro.components.extractor import Extractor, DatasetItem from datumaro.components.extractor import Extractor, DatasetItem
from datumaro.plugins.image_dir import ImageDirConverter
from datumaro.util.test_utils import TestDir, compare_datasets from datumaro.util.test_utils import TestDir, compare_datasets
from datumaro.util.image import save_image
class ImageDirFormatTest(TestCase): class ImageDirFormatTest(TestCase):
@ -21,8 +20,7 @@ class ImageDirFormatTest(TestCase):
with TestDir() as test_dir: with TestDir() as test_dir:
source_dataset = self.TestExtractor() source_dataset = self.TestExtractor()
for item in source_dataset: ImageDirConverter()(source_dataset, save_dir=test_dir)
save_image(osp.join(test_dir, '%s.jpg' % item.id), item.image)
project = Project.import_from(test_dir, 'image_dir') project = Project.import_from(test_dir, 'image_dir')
parsed_dataset = project.make_dataset() parsed_dataset = project.make_dataset()

@ -1,11 +1,10 @@
import numpy as np import numpy as np
import os.path as osp import os.path as osp
from PIL import Image
from unittest import TestCase from unittest import TestCase
from datumaro.util.test_utils import TestDir from datumaro.util.test_utils import TestDir
from datumaro.util.image import lazy_image from datumaro.util.image import lazy_image, load_image, save_image, Image
from datumaro.util.image_cache import ImageCache from datumaro.util.image_cache import ImageCache
@ -13,10 +12,8 @@ class LazyImageTest(TestCase):
def test_cache_works(self): def test_cache_works(self):
with TestDir() as test_dir: with TestDir() as test_dir:
image = np.ones((100, 100, 3), dtype=np.uint8) image = np.ones((100, 100, 3), dtype=np.uint8)
image = Image.fromarray(image).convert('RGB')
image_path = osp.join(test_dir, 'image.jpg') image_path = osp.join(test_dir, 'image.jpg')
image.save(image_path) save_image(image_path, image)
caching_loader = lazy_image(image_path, cache=None) caching_loader = lazy_image(image_path, cache=None)
self.assertTrue(caching_loader() is caching_loader()) self.assertTrue(caching_loader() is caching_loader())
@ -47,3 +44,36 @@ class ImageCacheTest(TestCase):
ImageCache.get_instance().clear() ImageCache.get_instance().clear()
self.assertTrue(loader() is loader()) self.assertTrue(loader() is loader())
self.assertEqual(ImageCache.get_instance().size(), 1) self.assertEqual(ImageCache.get_instance().size(), 1)
class ImageTest(TestCase):
def test_lazy_image_shape(self):
data = np.ones((5, 6, 7))
image_lazy = Image(data=data, size=(2, 4))
image_eager = Image(data=data)
self.assertEqual((2, 4), image_lazy.size)
self.assertEqual((5, 6), image_eager.size)
@staticmethod
def test_ctors():
with TestDir() as test_dir:
path = osp.join(test_dir, 'path.png')
image = np.ones([2, 4, 3])
save_image(path, image)
for args in [
{ 'data': image },
{ 'data': image, 'path': path },
{ 'data': image, 'path': path, 'size': (2, 4) },
{ 'data': image, 'path': path, 'loader': load_image, 'size': (2, 4) },
{ 'path': path },
{ 'path': path, 'loader': load_image },
{ 'path': path, 'size': (2, 4) },
]:
img = Image(**args)
# pylint: disable=pointless-statement
if img.has_data:
img.data
img.size
# pylint: enable=pointless-statement

@ -68,6 +68,25 @@ class PolygonConversionsTest(TestCase):
self.assertTrue(np.array_equal(e_mask, c_mask), self.assertTrue(np.array_equal(e_mask, c_mask),
'#%s: %s\n%s\n' % (i, e_mask, c_mask)) '#%s: %s\n%s\n' % (i, e_mask, c_mask))
def test_mask_to_rle(self):
source_mask = np.array([
[0, 1, 1, 1, 0, 1, 1, 1, 1, 0],
[0, 0, 1, 1, 0, 1, 0, 1, 0, 0],
[0, 0, 0, 1, 0, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
])
rle_uncompressed = mask_tools.mask_to_rle(source_mask)
from pycocotools import mask as mask_utils
resulting_mask = mask_utils.frPyObjects(
rle_uncompressed, *rle_uncompressed['size'])
resulting_mask = mask_utils.decode(resulting_mask)
self.assertTrue(np.array_equal(source_mask, resulting_mask),
'%s\n%s\n' % (source_mask, resulting_mask))
class ColormapOperationsTest(TestCase): class ColormapOperationsTest(TestCase):
def test_can_paint_mask(self): def test_can_paint_mask(self):
mask = np.zeros((1, 3), dtype=np.uint8) mask = np.zeros((1, 3), dtype=np.uint8)

@ -11,10 +11,11 @@ from datumaro.components.converter import Converter
from datumaro.components.extractor import (Extractor, DatasetItem, from datumaro.components.extractor import (Extractor, DatasetItem,
Label, Mask, Points, Polygon, PolyLine, Bbox, Caption, Label, Mask, Points, Polygon, PolyLine, Bbox, Caption,
) )
from datumaro.util.image import Image
from datumaro.components.config import Config, DefaultConfig, SchemaBuilder from datumaro.components.config import Config, DefaultConfig, SchemaBuilder
from datumaro.components.dataset_filter import \ from datumaro.components.dataset_filter import \
XPathDatasetFilter, XPathAnnotationsFilter, DatasetItemEncoder XPathDatasetFilter, XPathAnnotationsFilter, DatasetItemEncoder
from datumaro.util.test_utils import TestDir from datumaro.util.test_utils import TestDir, compare_datasets
class ProjectTest(TestCase): class ProjectTest(TestCase):
@ -135,12 +136,12 @@ class ProjectTest(TestCase):
class TestExtractor(Extractor): class TestExtractor(Extractor):
def __iter__(self): def __iter__(self):
for i in range(5): for i in range(5):
yield DatasetItem(id=i, subset='train', image=i) yield DatasetItem(id=i, subset='train', image=np.array([i]))
class TestLauncher(Launcher): class TestLauncher(Launcher):
def launch(self, inputs): def launch(self, inputs):
for i, inp in enumerate(inputs): for i, inp in enumerate(inputs):
yield [ Label(attributes={'idx': i, 'data': inp}) ] yield [ Label(attributes={'idx': i, 'data': inp.item()}) ]
model_name = 'model' model_name = 'model'
launcher_name = 'custom_launcher' launcher_name = 'custom_launcher'
@ -165,19 +166,18 @@ class ProjectTest(TestCase):
class TestExtractorSrc(Extractor): class TestExtractorSrc(Extractor):
def __iter__(self): def __iter__(self):
for i in range(2): for i in range(2):
yield DatasetItem(id=i, subset='train', image=i, yield DatasetItem(id=i, image=np.ones([2, 2, 3]) * i,
annotations=[ Label(i) ]) annotations=[Label(i)])
class TestLauncher(Launcher): class TestLauncher(Launcher):
def launch(self, inputs): def launch(self, inputs):
for inp in inputs: for inp in inputs:
yield [ Label(inp) ] yield [ Label(inp[0, 0, 0]) ]
class TestConverter(Converter): class TestConverter(Converter):
def __call__(self, extractor, save_dir): def __call__(self, extractor, save_dir):
for item in extractor: for item in extractor:
with open(osp.join(save_dir, '%s.txt' % item.id), 'w') as f: with open(osp.join(save_dir, '%s.txt' % item.id), 'w') as f:
f.write(str(item.subset) + '\n')
f.write(str(item.annotations[0].label) + '\n') f.write(str(item.annotations[0].label) + '\n')
class TestExtractorDst(Extractor): class TestExtractorDst(Extractor):
@ -189,11 +189,8 @@ class ProjectTest(TestCase):
for path in self.items: for path in self.items:
with open(path, 'r') as f: with open(path, 'r') as f:
index = osp.splitext(osp.basename(path))[0] index = osp.splitext(osp.basename(path))[0]
subset = f.readline().strip()
label = int(f.readline().strip()) label = int(f.readline().strip())
assert subset == 'train' yield DatasetItem(id=index, annotations=[Label(label)])
yield DatasetItem(id=index, subset=subset,
annotations=[ Label(label) ])
model_name = 'model' model_name = 'model'
launcher_name = 'custom_launcher' launcher_name = 'custom_launcher'
@ -476,6 +473,11 @@ class ExtractorTest(TestCase):
DatasetItem(id=2, subset='train'), DatasetItem(id=2, subset='train'),
DatasetItem(id=3, subset='test'), DatasetItem(id=3, subset='test'),
DatasetItem(id=4, subset='test'),
DatasetItem(id=1),
DatasetItem(id=2),
DatasetItem(id=3),
]) ])
extractor_name = 'ext1' extractor_name = 'ext1'
@ -485,8 +487,30 @@ class ExtractorTest(TestCase):
'url': 'path', 'url': 'path',
'format': extractor_name, 'format': extractor_name,
}) })
project.set_subsets(['train'])
dataset = project.make_dataset() dataset = project.make_dataset()
self.assertEqual(3, len(dataset)) compare_datasets(self, CustomExtractor(), dataset)
class DatasetItemTest(TestCase):
def test_ctor_requires_id(self):
has_error = False
try:
# pylint: disable=no-value-for-parameter
DatasetItem()
# pylint: enable=no-value-for-parameter
except AssertionError:
has_error = True
self.assertTrue(has_error)
@staticmethod
def test_ctors_with_image():
for args in [
{ 'id': 0, 'image': None },
{ 'id': 0, 'image': 'path.jpg' },
{ 'id': 0, 'image': np.array([1, 2, 3]) },
{ 'id': 0, 'image': lambda f: np.array([1, 2, 3]) },
{ 'id': 0, 'image': Image(data=np.array([1, 2, 3])) },
]:
DatasetItem(**args)

@ -8,6 +8,7 @@ from datumaro.components.extractor import (Extractor, DatasetItem,
from datumaro.plugins.tf_detection_api_format.importer import TfDetectionApiImporter from datumaro.plugins.tf_detection_api_format.importer import TfDetectionApiImporter
from datumaro.plugins.tf_detection_api_format.extractor import TfDetectionApiExtractor from datumaro.plugins.tf_detection_api_format.extractor import TfDetectionApiExtractor
from datumaro.plugins.tf_detection_api_format.converter import TfDetectionApiConverter from datumaro.plugins.tf_detection_api_format.converter import TfDetectionApiConverter
from datumaro.util.image import Image
from datumaro.util.test_utils import TestDir, compare_datasets from datumaro.util.test_utils import TestDir, compare_datasets
@ -33,16 +34,16 @@ class TfrecordConverterTest(TestCase):
DatasetItem(id=1, subset='train', DatasetItem(id=1, subset='train',
image=np.ones((16, 16, 3)), image=np.ones((16, 16, 3)),
annotations=[ annotations=[
Bbox(0, 4, 4, 8, label=2, id=0), Bbox(0, 4, 4, 8, label=2),
Bbox(0, 4, 4, 4, label=3, id=1), Bbox(0, 4, 4, 4, label=3),
Bbox(2, 4, 4, 4, id=2), Bbox(2, 4, 4, 4),
] ]
), ),
DatasetItem(id=2, subset='val', DatasetItem(id=2, subset='val',
image=np.ones((8, 8, 3)), image=np.ones((8, 8, 3)),
annotations=[ annotations=[
Bbox(1, 2, 4, 2, label=3, id=0), Bbox(1, 2, 4, 2, label=3),
] ]
), ),
@ -71,15 +72,15 @@ class TfrecordConverterTest(TestCase):
DatasetItem(id=1, DatasetItem(id=1,
image=np.ones((16, 16, 3)), image=np.ones((16, 16, 3)),
annotations=[ annotations=[
Bbox(2, 1, 4, 4, label=2, id=0), Bbox(2, 1, 4, 4, label=2),
Bbox(4, 2, 8, 4, label=3, id=1), Bbox(4, 2, 8, 4, label=3),
] ]
), ),
DatasetItem(id=2, DatasetItem(id=2,
image=np.ones((8, 8, 3)) * 2, image=np.ones((8, 8, 3)) * 2,
annotations=[ annotations=[
Bbox(4, 4, 4, 4, label=3, id=0), Bbox(4, 4, 4, 4, label=3),
] ]
), ),
@ -101,6 +102,20 @@ class TfrecordConverterTest(TestCase):
TestExtractor(), TfDetectionApiConverter(save_images=True), TestExtractor(), TfDetectionApiConverter(save_images=True),
test_dir) test_dir)
def test_can_save_dataset_with_image_info(self):
class TestExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=1, image=Image(path='1/q.e', size=(10, 15))),
])
def categories(self):
return { AnnotationType.label: LabelCategories() }
with TestDir() as test_dir:
self._test_save_and_load(TestExtractor(),
TfDetectionApiConverter(), test_dir)
def test_labelmap_parsing(self): def test_labelmap_parsing(self):
text = """ text = """
{ {

@ -244,3 +244,20 @@ class TransformsTest(TestCase):
actual = transforms.ShapesToBoxes(SrcExtractor()) actual = transforms.ShapesToBoxes(SrcExtractor())
compare_datasets(self, DstExtractor(), actual) compare_datasets(self, DstExtractor(), actual)
def test_id_from_image(self):
class SrcExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=1, image='path.jpg'),
DatasetItem(id=2),
])
class DstExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id='path', image='path.jpg'),
DatasetItem(id=2),
])
actual = transforms.IdFromImageName(SrcExtractor())
compare_datasets(self, DstExtractor(), actual)

@ -27,7 +27,7 @@ from datumaro.plugins.voc_format.converter import (
) )
from datumaro.plugins.voc_format.importer import VocImporter from datumaro.plugins.voc_format.importer import VocImporter
from datumaro.components.project import Project from datumaro.components.project import Project
from datumaro.util.image import save_image from datumaro.util.image import save_image, Image
from datumaro.util.test_utils import TestDir, compare_datasets from datumaro.util.test_utils import TestDir, compare_datasets
@ -171,13 +171,11 @@ def generate_dummy_voc(path):
return subsets return subsets
class TestExtractorBase(Extractor): class TestExtractorBase(Extractor):
_categories = VOC.make_voc_categories()
def _label(self, voc_label): def _label(self, voc_label):
return self.categories()[AnnotationType.label].find(voc_label)[0] return self.categories()[AnnotationType.label].find(voc_label)[0]
def categories(self): def categories(self):
return self._categories return VOC.make_voc_categories()
class VocExtractorTest(TestCase): class VocExtractorTest(TestCase):
def test_can_load_voc_cls(self): def test_can_load_voc_cls(self):
@ -694,6 +692,17 @@ class VocConverterTest(TestCase):
SrcExtractor(), VocConverter(label_map=label_map), SrcExtractor(), VocConverter(label_map=label_map),
test_dir, target_dataset=DstExtractor()) test_dir, target_dataset=DstExtractor())
def test_can_save_dataset_with_image_info(self):
class TestExtractor(TestExtractorBase):
def __iter__(self):
return iter([
DatasetItem(id=1, image=Image(path='1.jpg', size=(10, 15))),
])
with TestDir() as test_dir:
self._test_save_and_load(TestExtractor(),
VocConverter(label_map='voc'), test_dir)
class VocImportTest(TestCase): class VocImportTest(TestCase):
def test_can_import(self): def test_can_import(self):
with TestDir() as test_dir: with TestDir() as test_dir:

@ -1,4 +1,5 @@
import numpy as np import numpy as np
import os.path as osp
from unittest import TestCase from unittest import TestCase
@ -7,6 +8,7 @@ from datumaro.components.extractor import (Extractor, DatasetItem,
) )
from datumaro.plugins.yolo_format.importer import YoloImporter from datumaro.plugins.yolo_format.importer import YoloImporter
from datumaro.plugins.yolo_format.converter import YoloConverter from datumaro.plugins.yolo_format.converter import YoloConverter
from datumaro.util.image import Image, save_image
from datumaro.util.test_utils import TestDir, compare_datasets from datumaro.util.test_utils import TestDir, compare_datasets
@ -51,3 +53,64 @@ class YoloFormatTest(TestCase):
parsed_dataset = YoloImporter()(test_dir).make_dataset() parsed_dataset = YoloImporter()(test_dir).make_dataset()
compare_datasets(self, source_dataset, parsed_dataset) compare_datasets(self, source_dataset, parsed_dataset)
def test_can_save_dataset_with_image_info(self):
class TestExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=1, subset='train',
image=Image(path='1.jpg', size=(10, 15)),
annotations=[
Bbox(0, 2, 4, 2, label=2),
Bbox(3, 3, 2, 3, label=4),
]),
])
def categories(self):
label_categories = LabelCategories()
for i in range(10):
label_categories.add('label_' + str(i))
return {
AnnotationType.label: label_categories,
}
with TestDir() as test_dir:
source_dataset = TestExtractor()
YoloConverter()(source_dataset, test_dir)
save_image(osp.join(test_dir, 'obj_train_data', '1.jpg'),
np.ones((10, 15, 3))) # put the image for dataset
parsed_dataset = YoloImporter()(test_dir).make_dataset()
compare_datasets(self, source_dataset, parsed_dataset)
def test_can_load_dataset_with_exact_image_info(self):
class TestExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=1, subset='train',
image=Image(path='1.jpg', size=(10, 15)),
annotations=[
Bbox(0, 2, 4, 2, label=2),
Bbox(3, 3, 2, 3, label=4),
]),
])
def categories(self):
label_categories = LabelCategories()
for i in range(10):
label_categories.add('label_' + str(i))
return {
AnnotationType.label: label_categories,
}
with TestDir() as test_dir:
source_dataset = TestExtractor()
YoloConverter()(source_dataset, test_dir)
parsed_dataset = YoloImporter()(test_dir,
image_info={'1': (10, 15)}).make_dataset()
compare_datasets(self, source_dataset, parsed_dataset)

Loading…
Cancel
Save