Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixes in ICDAR and Market-1501 dataset formats #114

Merged
merged 4 commits into from
Feb 26, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 49 additions & 45 deletions datumaro/plugins/icdar_format/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@
import os.path as osp

from datumaro.components.converter import Converter
from datumaro.components.extractor import AnnotationType
from datumaro.components.extractor import AnnotationType, CompiledMask
from datumaro.util.image import save_image
from datumaro.util.mask_tools import paint_mask, merge_masks
from datumaro.util.mask_tools import paint_mask

from .format import IcdarPath, IcdarTask

Expand All @@ -17,7 +17,7 @@ class _WordRecognitionConverter:
def __init__(self):
self.annotations = ''

def save_annotations(self, item):
def save_annotations(self, item, path):
self.annotations += '%s, ' % (item.id + IcdarPath.IMAGE_EXT)
for ann in item.annotations:
if ann.type != AnnotationType.caption:
Expand All @@ -38,7 +38,7 @@ class _TextLocalizationConverter:
def __init__(self):
self.annotations = {}

def save_annotations(self, item):
def save_annotations(self, item, path):
annotation = ''
for ann in item.annotations:
if ann.type == AnnotationType.bbox:
Expand All @@ -65,59 +65,62 @@ def is_empty(self):
class _TextSegmentationConverter:
def __init__(self):
self.annotations = {}
self.masks = {}

def save_annotations(self, item):
masks = []
def save_annotations(self, item, path):
annotation = ''
colormap = [(255, 255, 255)]
anns = [a for a in item.annotations
if a.type == AnnotationType.mask]
anns = sorted(anns, key=lambda a: a.id)
group = anns[0].group
for ann in anns:
if ann.group != group or ann.group == 0:
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']
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure if we need to store color and center information as text.

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'
text = ''
if ann.attributes:
if 'text' in ann.attributes:
text = ann.attributes['text']
if text == ' ':
annotation += '#'
if 'color' in ann.attributes:
colormap.append(ann.attributes['color'])
annotation += ' '.join(str(p)
for p in ann.attributes['color'])
else:
annotation += '- - -'
if 'center' in ann.attributes:
annotation += ' '
annotation += ' '.join(str(p)
for p in 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
masks.append(ann.as_class_mask(ann.id))

mask = merge_masks(masks)
mask = paint_mask(mask,
{ i: colormap[i] for i in range(len(colormap)) })
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)) })
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This can easily consume lots of memory (for instance - imagine an uncompressed 4k video), try not to store masks.

save_image(osp.join(path, item.id + '_GT' + IcdarPath.GT_EXT),
mask, create_dir=True)
self.annotations[item.id] = annotation
self.masks[item.id] = mask

def write(self, path):
os.makedirs(path, exist_ok=True)
for item in self.annotations:
file = osp.join(path, item + '_GT' + '.txt')
with open(file, 'w') as f:
f.write(self.annotations[item])
save_image(osp.join(path, item + '_GT' + IcdarPath.GT_EXT),
self.masks[item], create_dir=True)

def is_empty(self):
return len(self.annotations) == 0
Expand Down Expand Up @@ -161,12 +164,13 @@ def apply(self):
for subset_name, subset in self._extractor.subsets().items():
task_converters = self._make_task_converters()
for item in subset:
for task_conv in task_converters.values():
for task, task_conv in task_converters.items():
if item.has_image and self._save_images:
self._save_image(item, osp.join(
self._save_dir, subset_name, IcdarPath.IMAGES_DIR,
item.id + IcdarPath.IMAGE_EXT))
task_conv.save_annotations(item)
task_conv.save_annotations(item, osp.join(self._save_dir,
IcdarPath.TASK_DIR[task], subset_name))

for task, task_conv in task_converters.items():
if task_conv.is_empty() and not self._tasks:
Expand Down
10 changes: 5 additions & 5 deletions datumaro/plugins/icdar_format/extractor.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ def _load_segmentation_items(self):
if len(objects) != 10:
continue

centers.append([float(objects[3]), float(objects[4])])
centers.append(objects[3] + ' ' + objects[4])
groups.append(group)
colors.append(tuple(int(o) for o in objects[:3]))
char = objects[9]
Expand All @@ -177,10 +177,10 @@ def _load_segmentation_items(self):
if label_id == 0:
continue
i = int(label_id)
annotations.append(Mask(id=i, group=groups[i],
annotations.append(Mask(group=groups[i],
image=self._lazy_extract_mask(mask, label_id),
attributes={ 'color': colors[i], 'text': chars[i],
'center': centers[i] }
attributes={ 'index': i - 1, 'color': ' '.join(str(p) for p in colors[i]),
'text': chars[i], 'center': centers[i] }
))
return items

Expand Down Expand Up @@ -229,4 +229,4 @@ def find_sources(cls, path):
sources += cls._find_sources_recursive(path, ext,
extractor_type, file_filter=lambda p:
osp.basename(p) != IcdarPath.VOCABULARY_FILE)
return sources
return sources
69 changes: 55 additions & 14 deletions datumaro/plugins/market1501_format.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import os.path as osp
import re
from distutils.util import strtobool
from glob import glob

from datumaro.components.converter import Converter
Expand All @@ -16,17 +17,24 @@ class Market1501Path:
BBOX_DIR = 'bounding_box_'
IMAGE_EXT = '.jpg'
PATTERN = re.compile(r'([-\d]+)_c(\d)')
IMAGE_NAMES = 'images_'

class Market1501Extractor(SourceExtractor):
def __init__(self, path):
if not osp.isdir(path):
raise NotADirectoryError("Can't open folder with annotation files '%s'" % path)
raise NotADirectoryError(
"Can't open folder with annotation files '%s'" % path)

subset = ''
for dirname in glob(osp.join(path, '*')):
if osp.basename(dirname).startswith(Market1501Path.BBOX_DIR):
subset = osp.basename(dirname).replace(Market1501Path.BBOX_DIR, '')
if osp.basename(dirname).startswith(Market1501Path.IMAGE_NAMES):
subset = osp.basename(dirname).replace(Market1501Path.IMAGE_NAMES, '')
subset = osp.splitext(subset)[0]
break
super().__init__(subset=subset)

self._path = path
self._items = list(self._load_items(path).values())

Expand All @@ -36,26 +44,38 @@ def _load_items(self, path):
paths = glob(osp.join(path, Market1501Path.QUERY_DIR, '*'))
paths += glob(osp.join(path, Market1501Path.BBOX_DIR + self._subset, '*'))

anno_file = osp.join(path,
Market1501Path.IMAGE_NAMES + self._subset + '.txt')
if len(paths) == 0 and osp.isfile(anno_file):
with open(anno_file, encoding='utf-8') as f:
for line in f:
paths.append(line.strip())

for image_path in paths:
if not osp.isfile(image_path) or \
osp.splitext(image_path)[-1] != Market1501Path.IMAGE_EXT:
if osp.splitext(image_path)[-1] != Market1501Path.IMAGE_EXT:
continue

item_id = osp.splitext(osp.basename(image_path))[0]
attributes = {}
pid, camid = map(int, Market1501Path.PATTERN.search(image_path).groups())
pid, camid = -1, -1
search = Market1501Path.PATTERN.search(image_path)
if search:
pid, camid = map(int, search.groups())
if 19 < len(item_id):
item_id = item_id[19:]
items[item_id] = DatasetItem(id=item_id, subset=self._subset,
image=image_path)

if pid == -1:
continue

attributes = items[item_id].attributes
camid -= 1
attributes['person_id'] = pid
attributes['camera_id'] = camid
if osp.basename(osp.dirname(image_path)) == Market1501Path.QUERY_DIR:
attributes['query'] = True
else:
attributes['query'] = False
items[item_id] = DatasetItem(id=item_id, subset=self._subset,
image=image_path, attributes=attributes)
return items

class Market1501Importer(Importer):
Expand All @@ -70,12 +90,33 @@ class Market1501Converter(Converter):

def apply(self):
for subset_name, subset in self._extractor.subsets().items():
annotation = ''
for item in subset:
image_name = item.id
if Market1501Path.PATTERN.search(image_name) == None:
if 'person_id' in item.attributes and \
'camera_id' in item.attributes:
image_pattern = '{:04d}_c{}s1_000000_00{}'
pid = int(item.attributes.get('person_id'))
camid = int(item.attributes.get('camera_id')) + 1
image_name = image_pattern.format(pid, camid, item.id)

dirname = Market1501Path.BBOX_DIR + subset_name
if 'query' in item.attributes:
query = item.attributes.get('query')
if isinstance(query, str):
query = strtobool(query)
if query:
dirname = Market1501Path.QUERY_DIR
image_path = osp.join(self._save_dir, dirname,
image_name + Market1501Path.IMAGE_EXT)
if item.has_image and self._save_images:
if item.attributes and 'query' in item.attributes:
if item.attributes.get('query'):
dirname = Market1501Path.QUERY_DIR
else:
dirname = Market1501Path.BBOX_DIR + subset_name
self._save_image(item, osp.join(self._save_dir,
dirname, item.id + Market1501Path.IMAGE_EXT))
self._save_image(item, image_path)
else:
annotation += '%s\n' % image_path

if 0 < len(annotation):
annotation_file = osp.join(self._save_dir,
Market1501Path.IMAGE_NAMES + subset_name + '.txt')
with open(annotation_file, 'w') as f:
f.write(annotation)
54 changes: 27 additions & 27 deletions tests/test_icdar_format.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,20 +69,20 @@ def test_can_import_masks(self):
DatasetItem(id='1', subset='train',
image=np.ones((2, 5, 3)),
annotations=[
Mask(id=1, group=0,
Mask(group=0,
image=np.array([[0, 1, 1, 0, 0], [0, 0, 0, 0, 0]]),
attributes={ 'color': (108, 225, 132),
'text': 'F', 'center': [0, 1]
attributes={ 'index': 0, 'color': '108 225 132',
'text': 'F', 'center': '0 1'
}),
Mask(id=2, group=1,
Mask(group=1,
image=np.array([[0, 0, 0, 1, 0], [0, 0, 0, 1, 0]]),
attributes = { 'color': (82, 174, 214),
'text': 'T', 'center': [1, 3]
attributes={ 'index': 1, 'color': '82 174 214',
'text': 'T', 'center': '1 3'
}),
Mask(id=3, group=1,
Mask(group=1,
image=np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 1]]),
attributes = { 'color': (241, 73, 144),
'text': 'h', 'center': [1, 4]
attributes={ 'index': 2, 'color': '241 73 144',
'text': 'h', 'center': '1 4'
}),
]
),
Expand Down Expand Up @@ -145,27 +145,27 @@ def test_can_save_and_load_masks(self):
expected_dataset = Dataset.from_iterable([
DatasetItem(id=1, subset='train',
annotations=[
Mask(id=2, image=np.array([[0, 0, 0, 1, 1]]), group=1,
attributes={ 'color': (82, 174, 214), 'text': 'j',
'center': [0, 3] }),
Mask(id=1, image=np.array([[0, 1, 1, 0, 0]]), group=1,
attributes={ 'color': (108, 225, 132), 'text': 'F',
'center': [0, 1] }),
Mask(image=np.array([[0, 0, 0, 1, 1]]), group=1,
attributes={ 'index': 1, 'color': '82 174 214', 'text': 'j',
'center': '0 3' }),
Mask(image=np.array([[0, 1, 1, 0, 0]]), group=1,
attributes={ 'index': 0, 'color': '108 225 132', 'text': 'F',
'center': '0 1' }),
]),
DatasetItem(id=2, subset='train',
annotations=[
Mask(id=4, image=np.array([[0, 0, 0, 0, 0, 1]]), group=0,
attributes={ 'color': (183, 6, 28), 'text': ' ',
'center': [0, 5] }),
Mask(id=1, image=np.array([[1, 0, 0, 0, 0, 0]]), group=1,
attributes={ 'color': (108, 225, 132), 'text': 'L',
'center': [0, 0] }),
Mask(id=2, image=np.array([[0, 0, 0, 1, 1, 0]]), group=1,
attributes={ 'color': (82, 174, 214), 'text': 'o',
'center': [0, 3] }),
Mask(id=3, image=np.array([[0, 1, 1, 0, 0, 0]]), group=0,
attributes={ 'color': (241, 73, 144), 'text': 'P',
'center': [0, 1] }),
Mask(image=np.array([[0, 0, 0, 0, 0, 1]]), group=0,
attributes={ 'index': 3, 'color': '183 6 28', 'text': ' ',
'center': '0 5' }),
Mask(image=np.array([[1, 0, 0, 0, 0, 0]]), group=1,
attributes={ 'index': 0, 'color': '108 225 132', 'text': 'L',
'center': '0 0' }),
Mask(image=np.array([[0, 0, 0, 1, 1, 0]]), group=1,
attributes={ 'index': 1, 'color': '82 174 214', 'text': 'o',
'center': '0 3' }),
Mask(image=np.array([[0, 1, 1, 0, 0, 0]]), group=0,
attributes={ 'index': 2, 'color': '241 73 144', 'text': 'P',
'center': '0 1' }),
]),
])

Expand Down
Loading