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 _save_image(self,
item,
path=None,
*,
name=None,
subdir=None,
basedir=None):
assert not ((subdir or name or basedir) and path), \
"Can't use both subdir or name or basedir and path arguments"
if not item.image.has_data:
log.warning("Item '%s' has no image", item.id)
return
basedir = basedir or self._save_dir
path = path or osp.join(
basedir, self._make_image_filename(item, name=name, subdir=subdir))
path = osp.abspath(path)
src_ext = item.image.ext.lower()
dst_ext = osp.splitext(osp.basename(path))[1].lower()
os.makedirs(osp.dirname(path), exist_ok=True)
if src_ext == dst_ext and osp.isfile(item.image.path):
if item.image.path != path:
shutil.copyfile(item.image.path, path)
elif src_ext == dst_ext and isinstance(item.image, ByteImage):
with open(path, 'wb') as f:
f.write(item.image.get_bytes())
else:
save_image(path, item.image.data)
def test_ctors(self):
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, 'ext': 'png' },
{ 'data': image, 'ext': 'png', 'size': (2, 4) },
{ 'data': lambda p: image },
{ 'data': lambda p: image, 'path': 'somepath' },
{ 'data': lambda p: image, 'ext': 'jpg' },
{ 'path': path },
{ 'path': path, 'data': load_image },
{ 'path': path, 'data': load_image, 'size': (2, 4) },
{ 'path': path, 'size': (2, 4) },
]:
with self.subTest(**args):
img = Image(**args)
self.assertTrue(img.has_data)
np.testing.assert_array_equal(img.data, image)
self.assertEqual(img.size, tuple(image.shape[:2]))
with self.subTest():
img = Image(size=(2, 4))
self.assertEqual(img.size, (2, 4))
def _save_item(self, subset_name, subset, item):
if self._save_images and item.has_image:
self._save_image(item,
subdir=osp.join(subset_name,
IcdarPath.IMAGES_DIR))
annotation = ''
colormap = [(255, 255, 255)]
anns = [a for a in item.annotations if a.type == AnnotationType.mask]
if anns:
anns = sorted(anns,
key=lambda a: int(a.attributes.get('index', 0)))
group = anns[0].group
for ann in anns:
if ann.group != group or (not ann.group
and anns[0].group != 0):
annotation += '\n'
text = ''
if ann.attributes:
if 'text' in ann.attributes:
text = ann.attributes['text']
if text == ' ':
annotation += '#'
if 'color' in ann.attributes and \
len(ann.attributes['color'].split()) == 3:
color = ann.attributes['color'].split()
colormap.append(
(int(color[0]), int(color[1]), int(color[2])))
annotation += ' '.join(p for p in color)
else:
raise Exception(
"Item %s: a mask must have "
"an RGB color attribute, e.g. '10 7 50'" % item.id)
if 'center' in ann.attributes:
annotation += ' %s' % ann.attributes['center']
else:
annotation += ' - -'
bbox = ann.get_bbox()
annotation += ' %s %s %s %s' % (bbox[0], bbox[1], bbox[0] +
bbox[2], bbox[1] + bbox[3])
annotation += ' \"%s\"' % text
annotation += '\n'
group = ann.group
mask = CompiledMask.from_instance_masks(
anns,
instance_labels=[m.attributes['index'] + 1 for m in anns])
mask = paint_mask(mask.class_mask,
{i: colormap[i]
for i in range(len(colormap))})
save_image(osp.join(self._save_dir, subset_name,
item.id + '_GT' + IcdarPath.GT_EXT),
mask,
create_dir=True)
anno_file = osp.join(self._save_dir, subset_name,
item.id + '_GT' + '.txt')
os.makedirs(osp.dirname(anno_file), exist_ok=True)
with open(anno_file, 'w', encoding='utf-8') as f:
f.write(annotation)
def save_segm(self, path, annotation, colormap):
data = annotation.image
if self._apply_colormap:
if colormap is None:
colormap = VocColormap
data = apply_colormap(data, colormap)
save_image(path, data)
def test_can_save_dataset_with_image_info(self):
source_dataset = Dataset.from_iterable(
[
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),
]),
],
categories={
AnnotationType.label:
LabelCategories.from_iterable('label_' + str(i)
for i in range(10)),
})
with TestDir() as test_dir:
YoloConverter.convert(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 = Dataset.import_from(test_dir, 'yolo')
compare_datasets(self, source_dataset, parsed_dataset)
def save_item_bbox_diff(self, item_a, item_b, diff):
_, mispred, a_unmatched, b_unmatched = diff
if 0 < len(a_unmatched) + len(b_unmatched) + len(mispred):
if not item_a.has_image or not item_a.image.has_data:
log.warning("Item %s: item has no image data, "
"it will be skipped" % (item_a.id))
return
img_a = item_a.image.data.copy()
img_b = img_a.copy()
for a_bbox, b_bbox in mispred:
self.draw_bbox(img_a, a_bbox, self.get_a_label(a_bbox.label),
(0, 255, 0))
self.draw_bbox(img_b, b_bbox, self.get_b_label(b_bbox.label),
(0, 0, 255))
for a_bbox in a_unmatched:
self.draw_bbox(img_a, a_bbox, self.get_a_label(a_bbox.label),
(255, 255, 0))
for b_bbox in b_unmatched:
self.draw_bbox(img_b, b_bbox, self.get_b_label(b_bbox.label),
(255, 255, 0))
img = np.hstack([img_a, img_b])
path = osp.join(self.save_dir, item_a.subset, item_a.id)
if self.output_format is OutputFormat.simple:
save_image(path + '.png', img, create_dir=True)
elif self.output_format is OutputFormat.tensorboard:
self.save_as_tensorboard(img, path)
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:
save_image(osp.join(save_dir, item.id + '.jpg'),
item.image.data,
create_dir=True)
def _save_image(self, item):
image = item.image
if image is None:
return
image_path = osp.join(self._images_dir,
str(item.id) + DatumaroPath.IMAGE_EXT)
save_image(image_path, image)
def _save_image(self, item):
image = item.image.data
if image is None:
return ''
filename = item.id + DatumaroPath.IMAGE_EXT
image_path = osp.join(self._images_dir, filename)
save_image(image_path, image, create_dir=True)
return filename
def _save_image(self, item, filename):
image = item.image.data
if image is None:
log.warning("Item '%s' has no image" % item.id)
return ''
save_image(osp.join(self._context._images_dir, filename),
image,
create_dir=True)
def _save_image(self, item, index):
if item.image.filename:
frame_id = osp.splitext(item.image.filename)[0]
else:
frame_id = item.id
frame_id = cast(frame_id, int, index)
image_filename = '%06d%s' % (frame_id, MotPath.IMAGE_EXT)
save_image(osp.join(self._images_dir, image_filename),
item.image.data)
def __call__(self, extractor, save_dir):
images_dir = osp.join(save_dir, MotPath.IMAGE_DIR)
os.makedirs(images_dir, exist_ok=True)
self._images_dir = images_dir
anno_dir = osp.join(save_dir, 'gt')
os.makedirs(anno_dir, exist_ok=True)
anno_file = osp.join(anno_dir, MotPath.GT_FILENAME)
with open(anno_file, 'w', encoding="utf-8") as csv_file:
writer = csv.DictWriter(csv_file, fieldnames=MotPath.FIELDS)
track_id_mapping = {-1: -1}
for idx, item in enumerate(extractor):
log.debug("Converting item '%s'", item.id)
frame_id = cast(item.id, int, 1 + idx)
for anno in item.annotations:
if anno.type != AnnotationType.bbox:
continue
track_id = int(anno.attributes.get('track_id', -1))
if track_id not in track_id_mapping:
track_id_mapping[track_id] = len(track_id_mapping)
track_id = track_id_mapping[track_id]
writer.writerow({
'frame_id': frame_id,
'track_id': track_id,
'x': anno.x,
'y': anno.y,
'w': anno.w,
'h': anno.h,
'confidence': int(anno.attributes.get('ignored') != True),
'class_id': 1 + cast(anno.label, int, -2),
'visibility': float(
anno.attributes.get('visibility',
1 - float(
anno.attributes.get('occluded', False)
)
)
)
})
if self._save_images:
if item.has_image and item.image.has_data:
save_image(osp.join(self._images_dir,
'%06d%s' % (frame_id, MotPath.IMAGE_EXT)),
item.image.data)
else:
log.debug("Item '%s' has no image" % item.id)
labels_file = osp.join(anno_dir, MotPath.LABELS_FILE)
with open(labels_file, 'w', encoding='utf-8') as f:
f.write('\n'.join(l.name
for l in extractor.categories()[AnnotationType.label].items)
)
def test_image_true_when_true(self):
with TestDir() as test_dir:
path = osp.join(test_dir, 'test.jpg')
save_image(path, np.ones([10, 7, 3]))
target = ImageTarget()
status = target.test(path)
self.assertTrue(status)
def test_cache_works(self):
with TestDir() as test_dir:
image = np.ones((100, 100, 3), dtype=np.uint8)
image_path = osp.join(test_dir, 'image.jpg')
save_image(image_path, image)
caching_loader = lazy_image(image_path, cache=True)
self.assertTrue(caching_loader() is caching_loader())
non_caching_loader = lazy_image(image_path, cache=False)
self.assertFalse(non_caching_loader() is non_caching_loader())
def save_mask(self,
path,
mask,
colormap=None,
apply_colormap=True,
dtype=np.uint8):
if self._apply_colormap and apply_colormap:
if colormap is None:
colormap = self._categories[AnnotationType.mask].colormap
mask = paint_mask(mask, colormap)
save_image(path, mask, create_dir=True, dtype=dtype)
def save_annotations(self, item, path):
annotation = ''
colormap = [(255, 255, 255)]
anns = [a for a in item.annotations if a.type == AnnotationType.mask]
if anns:
is_not_index = len(
[p for p in anns if 'index' not in p.attributes])
if is_not_index:
raise Exception("Item %s: a mask must have"
"'index' attribute" % item.id)
anns = sorted(anns, key=lambda a: a.attributes['index'])
group = anns[0].group
for ann in anns:
if ann.group != group or (not ann.group
and anns[0].group != 0):
annotation += '\n'
text = ''
if ann.attributes:
if 'text' in ann.attributes:
text = ann.attributes['text']
if text == ' ':
annotation += '#'
if 'color' in ann.attributes and \
len(ann.attributes['color'].split()) == 3:
color = ann.attributes['color'].split()
colormap.append(
(int(color[0]), int(color[1]), int(color[2])))
annotation += ' '.join(p for p in color)
else:
raise Exception(
"Item %s: a mask must have "
"an RGB color attribute, e. g. '10 7 50'" %
item.id)
if 'center' in ann.attributes:
annotation += ' %s' % ann.attributes['center']
else:
annotation += ' - -'
bbox = ann.get_bbox()
annotation += ' %s %s %s %s' % (bbox[0], bbox[1], bbox[0] +
bbox[2], bbox[1] + bbox[3])
annotation += ' \"%s\"' % text
annotation += '\n'
group = ann.group
mask = CompiledMask.from_instance_masks(
anns,
instance_labels=[m.attributes['index'] + 1 for m in anns])
mask = paint_mask(mask.class_mask,
{i: colormap[i]
for i in range(len(colormap))})
save_image(osp.join(path, item.id + '_GT' + IcdarPath.GT_EXT),
mask,
create_dir=True)
self.annotations[item.id] = annotation
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)
def COCO_dataset_generate(self, path):
img_dir = osp.join(path, 'images', 'val')
ann_dir = osp.join(path, 'annotations')
os.makedirs(img_dir)
os.makedirs(ann_dir)
image = np.ones((10, 5, 3))
save_image(osp.join(img_dir, '000000000001.jpg'), image)
annotation = self.generate_annotation()
with open(osp.join(ann_dir, 'instances_val.json'), 'w') as outfile:
json.dump(annotation, outfile)
def save(self, path):
cur_path = osp.abspath(self.path)
path = osp.abspath(path)
cur_ext = self.ext.lower()
new_ext = osp.splitext(osp.basename(path))[1].lower()
os.makedirs(osp.dirname(path), exist_ok=True)
if cur_ext == new_ext and osp.isfile(cur_path):
if cur_path != path:
shutil.copyfile(cur_path, path)
else:
save_image(path, self.data)
def _save_image(self, item):
image = item.image.data
if image is None:
return ''
filename = item.image.filename
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)
return filename
def test_ctors(self):
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': 'somepath',
'loader': lambda p: image
},
{
'loader': lambda p: image
},
{
'path': path,
'size': (2, 4)
},
]:
with self.subTest(**args):
img = Image(**args)
# pylint: disable=pointless-statement
self.assertTrue(img.has_data)
self.assertEqual(img, image)
self.assertEqual(img.size, tuple(image.shape[:2]))
def _save_annotations(self, item, anno_dir):
masks = [a for a in item.annotations if a.type == AnnotationType.mask]
if not masks:
return
instance_ids = [int(a.attributes['track_id']) for a in masks]
masks = sorted(zip(masks, instance_ids), key=lambda e: e[0].z_order)
mask = merge_masks(
(m.image, MotsPath.MAX_INSTANCES * (1 + m.label) + id)
for m, id in masks)
save_image(osp.join(anno_dir, item.id + '.png'),
mask,
create_dir=True,
dtype=np.uint16)
def test_can_save_and_load_image_with_custom_extension(self):
expected = Dataset.from_iterable([
DatasetItem(id='a/3', image=Image(path='a/3.qq',
data=np.zeros((3, 4, 3)))),
])
with TestDir() as test_dir:
image_path = osp.join(test_dir, 'a', '3.jpg')
save_image(image_path, expected.get('a/3').image.data,
create_dir=True)
os.rename(image_path, osp.join(test_dir, 'a', '3.qq'))
actual = Dataset.import_from(test_dir, 'image_dir', exts='qq')
compare_datasets(self, expected, actual, require_images=True)
def _test_can_save_and_load(self,
src_image,
path,
save_backend=None,
load_backend=None):
if save_backend:
image_module._IMAGE_BACKEND = save_backend
image_module.save_image(path, src_image)
if load_backend:
image_module._IMAGE_BACKEND = load_backend
dst_image = image_module.load_image(path)
self.assertTrue(np.all(src_image == dst_image), 'save: %s, load: %s' % \
(save_backend, load_backend))
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)
save_image(path, image)
return path
def _save_image(self, item, path=None):
image = item.image.data
if image is None:
log.warning("Item '%s' has no image", item.id)
return item.image.path
path = path or self._make_image_filename(item)
src_ext = osp.splitext(osp.basename(item.image.path))[1]
dst_ext = osp.splitext(osp.basename(path))[1]
os.makedirs(osp.dirname(path), exist_ok=True)
if src_ext == dst_ext and osp.isfile(item.image.path):
shutil.copyfile(item.image.path, path)
else:
save_image(path, image)
def _save_image(self, item, path=None):
if not item.image.has_data:
log.warning("Item '%s' has no image", item.id)
return
path = path or self._make_image_filename(item)
src_ext = item.image.ext.lower()
dst_ext = osp.splitext(osp.basename(path))[1].lower()
os.makedirs(osp.dirname(path), exist_ok=True)
if src_ext == dst_ext and osp.isfile(item.image.path):
shutil.copyfile(item.image.path, path)
elif src_ext == dst_ext and isinstance(item.image, ByteImage):
with open(path, 'wb') as f:
f.write(item.image.get_bytes())
else:
save_image(path, item.image.data)
def test_save_and_load_backends(self):
backends = image_module._IMAGE_BACKENDS
for save_backend, load_backend, c in product(backends, backends, [1, 3]):
with TestDir() as test_dir:
if c == 1:
src_image = np.random.randint(0, 255 + 1, (2, 4))
else:
src_image = np.random.randint(0, 255 + 1, (2, 4, c))
path = osp.join(test_dir, 'img.png') # lossless
image_module._IMAGE_BACKEND = save_backend
image_module.save_image(path, src_image, jpeg_quality=100)
image_module._IMAGE_BACKEND = load_backend
dst_image = image_module.load_image(path)
self.assertTrue(np.array_equal(src_image, dst_image),
'save: %s, load: %s' % (save_backend, load_backend))
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
