# Copyright (C) 2020 Intel Corporation
#
# SPDX-License-Identifier: MIT
import os
import os.path as osp
from datumaro.components.annotation import AnnotationType, CompiledMask
from datumaro.components.errors import DatasetExportError, MediaTypeError
from datumaro.components.exporter import Exporter
from datumaro.components.media import Image
from datumaro.util.image import save_image
from datumaro.util.mask_tools import generate_colormap, paint_mask
from .format import IcdarPath
[docs]
class IcdarWordRecognitionExporter(Exporter):
DEFAULT_IMAGE_EXT = IcdarPath.IMAGE_EXT
def _apply_impl(self):
if self._extractor.media_type() and not issubclass(self._extractor.media_type(), Image):
raise MediaTypeError("Media type is not an image")
for subset_name, subset in self._extractor.subsets().items():
annotation = ""
for item in subset:
image_filename = self._make_image_filename(item)
if self._save_media and item.media:
self._save_image(
item,
osp.join(self._save_dir, subset_name, IcdarPath.IMAGES_DIR, image_filename),
)
annotation += "%s, " % image_filename
for ann in item.annotations:
if ann.type != AnnotationType.caption:
continue
annotation += '"%s"' % ann.caption
annotation += "\n"
if len(annotation):
anno_file = osp.join(self._save_dir, subset_name, "gt.txt")
os.makedirs(osp.dirname(anno_file), exist_ok=True)
with open(anno_file, "w", encoding="utf-8") as f:
f.write(annotation)
[docs]
class IcdarTextLocalizationExporter(Exporter):
DEFAULT_IMAGE_EXT = IcdarPath.IMAGE_EXT
def _apply_impl(self):
if self._extractor.media_type() and not issubclass(self._extractor.media_type(), Image):
raise MediaTypeError("Media type is not an image")
for subset_name, subset in self._extractor.subsets().items():
for item in subset:
if self._save_media and item.media:
self._save_image(item, subdir=osp.join(subset_name, IcdarPath.IMAGES_DIR))
annotation = ""
for ann in item.annotations:
if ann.type == AnnotationType.bbox:
annotation += " ".join(str(p) for p in ann.points)
if ann.attributes and "text" in ann.attributes:
annotation += ' "%s"' % ann.attributes["text"]
elif ann.type == AnnotationType.polygon:
annotation += ",".join(str(p) for p in ann.points)
if ann.attributes and "text" in ann.attributes:
annotation += ',"%s"' % ann.attributes["text"]
annotation += "\n"
anno_file = osp.join(
self._save_dir,
subset_name,
osp.dirname(item.id),
"gt_" + osp.basename(item.id) + ".txt",
)
os.makedirs(osp.dirname(anno_file), exist_ok=True)
with open(anno_file, "w", encoding="utf-8") as f:
f.write(annotation)
[docs]
class IcdarTextSegmentationExporter(Exporter):
DEFAULT_IMAGE_EXT = IcdarPath.IMAGE_EXT
def _apply_impl(self):
if self._extractor.media_type() and not issubclass(self._extractor.media_type(), Image):
raise MediaTypeError("Media type is not an image")
for subset_name, subset in self._extractor.subsets().items():
for item in subset:
self._save_item(subset_name, subset, item)
def _save_item(self, subset_name, subset, item):
if self._save_media and item.media:
self._save_image(item, subdir=osp.join(subset_name, IcdarPath.IMAGES_DIR))
annotation = ""
anns = [a for a in item.annotations if a.type == AnnotationType.mask]
color_bank = iter(generate_colormap(len(anns), include_background=False).values())
colormap = [(255, 255, 255)]
used_colors = set(colormap)
if anns:
anns = sorted(anns, key=lambda a: int(a.attributes.get("index", 0)))
group = anns[0].group
for i, ann in enumerate(anns):
# Assign new color if it is not defined
color = ann.attributes.get("color", "")
if color:
color = color.split()
if len(color) != 3:
raise DatasetExportError(
"Item %s: mask #%s has invalid color" % (item.id, i)
)
color = tuple(map(int, color))
else:
color = next(color_bank)
while color in used_colors:
color = next(color_bank)
colormap.append(color)
used_colors.add(color)
text = ann.attributes.get("text", "")
bbox = ann.get_bbox()
if ann.group != group or (not ann.group and anns[0].group != 0):
annotation += "\n"
if text == " ":
annotation += "#"
annotation += " ".join(str(p) for p in color)
annotation += " %s" % ann.attributes.get("center", "- -")
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)
