zhiltsov-max
6 years ago
committed by
GitHub
parent
8efaf58f11
commit
0744c6ae16
@ -0,0 +1,341 @@
|
||||
|
||||
# Copyright (C) 2020 Intel Corporation
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
# The Multiple Object Tracking Benchmark challenge format support
|
||||
# Format description: https://arxiv.org/pdf/1906.04567.pdf
|
||||
# Another description: https://motchallenge.net/instructions
|
||||
|
||||
from collections import OrderedDict
|
||||
import csv
|
||||
from enum import Enum
|
||||
import logging as log
|
||||
import os
|
||||
import os.path as osp
|
||||
|
||||
from datumaro.components.extractor import (SourceExtractor,
|
||||
DatasetItem, AnnotationType, Bbox, LabelCategories
|
||||
)
|
||||
from datumaro.components.extractor import Importer
|
||||
from datumaro.components.converter import Converter
|
||||
from datumaro.components.cli_plugin import CliPlugin
|
||||
from datumaro.util import cast
|
||||
from datumaro.util.image import Image, save_image
|
||||
|
||||
|
||||
MotLabel = Enum('MotLabel', [
|
||||
('pedestrian', 1),
|
||||
('person on vehicle', 2),
|
||||
('car', 3),
|
||||
('bicycle', 4),
|
||||
('motorbike', 5),
|
||||
('non motorized vehicle', 6),
|
||||
('static person', 7),
|
||||
('distractor', 8),
|
||||
('occluder', 9),
|
||||
('occluder on the ground', 10),
|
||||
('occluder full', 11),
|
||||
('reflection', 12),
|
||||
])
|
||||
|
||||
class MotPath:
|
||||
IMAGE_DIR = 'img1'
|
||||
SEQINFO_FILE = 'seqinfo.ini'
|
||||
LABELS_FILE = 'labels.txt'
|
||||
GT_FILENAME = 'gt.txt'
|
||||
DET_FILENAME = 'det.txt'
|
||||
|
||||
IMAGE_EXT = '.jpg'
|
||||
|
||||
FIELDS = [
|
||||
'frame_id',
|
||||
'track_id',
|
||||
'x',
|
||||
'y',
|
||||
'w',
|
||||
'h',
|
||||
'confidence', # or 'not ignored' flag for GT anns
|
||||
'class_id',
|
||||
'visibility'
|
||||
]
|
||||
|
||||
|
||||
class MotSeqExtractor(SourceExtractor):
|
||||
def __init__(self, path, labels=None, occlusion_threshold=0, is_gt=None):
|
||||
super().__init__()
|
||||
|
||||
assert osp.isfile(path)
|
||||
self._path = path
|
||||
seq_root = osp.dirname(osp.dirname(path))
|
||||
|
||||
self._image_dir = ''
|
||||
if osp.isdir(osp.join(seq_root, MotPath.IMAGE_DIR)):
|
||||
self._image_dir = osp.join(seq_root, MotPath.IMAGE_DIR)
|
||||
|
||||
seq_info = osp.join(seq_root, MotPath.SEQINFO_FILE)
|
||||
if osp.isfile(seq_info):
|
||||
seq_info = self._parse_seq_info(seq_info)
|
||||
self._image_dir = osp.join(seq_root, seq_info['imdir'])
|
||||
else:
|
||||
seq_info = None
|
||||
self._seq_info = seq_info
|
||||
|
||||
self._occlusion_threshold = float(occlusion_threshold)
|
||||
|
||||
assert is_gt in {None, True, False}
|
||||
if is_gt is None:
|
||||
if osp.basename(path) == MotPath.DET_FILENAME:
|
||||
is_gt = False
|
||||
else:
|
||||
is_gt = True
|
||||
self._is_gt = is_gt
|
||||
|
||||
self._subset = None
|
||||
|
||||
if labels is None:
|
||||
if osp.isfile(osp.join(seq_root, MotPath.LABELS_FILE)):
|
||||
labels = osp.join(seq_root, MotPath.LABELS_FILE)
|
||||
else:
|
||||
labels = [lbl.name for lbl in MotLabel]
|
||||
if isinstance(labels, str):
|
||||
labels = self._parse_labels(labels)
|
||||
elif isinstance(labels, list):
|
||||
assert all(isinstance(lbl, str) for lbl in labels), labels
|
||||
else:
|
||||
raise TypeError("Unexpected type of 'labels' argument: %s" % labels)
|
||||
self._categories = self._load_categories(labels)
|
||||
self._items = self._load_items(path)
|
||||
|
||||
def categories(self):
|
||||
return self._categories
|
||||
|
||||
def __iter__(self):
|
||||
for item in self._items.values():
|
||||
yield item
|
||||
|
||||
def __len__(self):
|
||||
return len(self._items)
|
||||
|
||||
def subsets(self):
|
||||
if self._subset:
|
||||
return [self._subset]
|
||||
return None
|
||||
|
||||
def get_subset(self, name):
|
||||
if name != self._subset:
|
||||
return None
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def _parse_labels(path):
|
||||
with open(path, encoding='utf-8') as labels_file:
|
||||
return [s.strip() for s in labels_file]
|
||||
|
||||
def _load_categories(self, labels):
|
||||
attributes = ['track_id']
|
||||
if self._is_gt:
|
||||
attributes += ['occluded', 'visibility', 'ignored']
|
||||
else:
|
||||
attributes += ['score']
|
||||
label_cat = LabelCategories(attributes=attributes)
|
||||
for label in labels:
|
||||
label_cat.add(label)
|
||||
|
||||
return { AnnotationType.label: label_cat }
|
||||
|
||||
def _load_items(self, path):
|
||||
labels_count = len(self._categories[AnnotationType.label].items)
|
||||
items = OrderedDict()
|
||||
|
||||
if self._seq_info:
|
||||
for frame_id in range(self._seq_info['seqlength']):
|
||||
items[frame_id] = DatasetItem(
|
||||
id=frame_id,
|
||||
subset=self._subset,
|
||||
image=Image(
|
||||
path=osp.join(self._image_dir,
|
||||
'%06d%s' % (frame_id, self._seq_info['imext'])),
|
||||
size=(self._seq_info['imheight'], self._seq_info['imwidth'])
|
||||
)
|
||||
)
|
||||
elif osp.isdir(self._image_dir):
|
||||
for p in os.listdir(self._image_dir):
|
||||
if p.endswith(MotPath.IMAGE_EXT):
|
||||
frame_id = int(osp.splitext(p)[0])
|
||||
items[frame_id] = DatasetItem(
|
||||
id=frame_id,
|
||||
subset=self._subset,
|
||||
image=osp.join(self._image_dir, p),
|
||||
)
|
||||
|
||||
with open(path, newline='', encoding='utf-8') as csv_file:
|
||||
# NOTE: Different MOT files have different count of fields
|
||||
# (7, 9 or 10). This is handled by reader:
|
||||
# - all extra fields go to a separate field
|
||||
# - all unmet fields have None values
|
||||
for row in csv.DictReader(csv_file, fieldnames=MotPath.FIELDS):
|
||||
frame_id = int(row['frame_id'])
|
||||
item = items.get(frame_id)
|
||||
if item is None:
|
||||
item = DatasetItem(id=frame_id, subset=self._subset)
|
||||
annotations = item.annotations
|
||||
|
||||
x, y = float(row['x']), float(row['y'])
|
||||
w, h = float(row['w']), float(row['h'])
|
||||
label_id = row.get('class_id')
|
||||
if label_id and label_id != '-1':
|
||||
label_id = int(label_id) - 1
|
||||
assert label_id < labels_count, label_id
|
||||
else:
|
||||
label_id = None
|
||||
|
||||
attributes = {}
|
||||
|
||||
# Annotations for detection task are not related to any track
|
||||
track_id = int(row['track_id'])
|
||||
if 0 < track_id:
|
||||
attributes['track_id'] = track_id
|
||||
|
||||
confidence = cast(row.get('confidence'), float, 1)
|
||||
visibility = cast(row.get('visibility'), float, 1)
|
||||
if self._is_gt:
|
||||
attributes['visibility'] = visibility
|
||||
attributes['occluded'] = \
|
||||
visibility <= self._occlusion_threshold
|
||||
attributes['ignored'] = confidence == 0
|
||||
else:
|
||||
attributes['score'] = float(confidence)
|
||||
|
||||
annotations.append(Bbox(x, y, w, h, label=label_id,
|
||||
attributes=attributes))
|
||||
|
||||
items[frame_id] = item
|
||||
return items
|
||||
|
||||
@classmethod
|
||||
def _parse_seq_info(cls, path):
|
||||
fields = {}
|
||||
with open(path, encoding='utf-8') as f:
|
||||
for line in f:
|
||||
entry = line.lower().strip().split('=', maxsplit=1)
|
||||
if len(entry) == 2:
|
||||
fields[entry[0]] = entry[1]
|
||||
cls._check_seq_info(fields)
|
||||
for k in { 'framerate', 'seqlength', 'imwidth', 'imheight' }:
|
||||
fields[k] = int(fields[k])
|
||||
return fields
|
||||
|
||||
@staticmethod
|
||||
def _check_seq_info(seq_info):
|
||||
assert set(seq_info) == {'name', 'imdir', 'framerate', 'seqlength', 'imwidth', 'imheight', 'imext'}, seq_info
|
||||
|
||||
class MotSeqImporter(Importer):
|
||||
_EXTRACTOR_NAME = 'mot_seq'
|
||||
|
||||
@classmethod
|
||||
def detect(cls, path):
|
||||
return len(cls.find_subsets(path)) != 0
|
||||
|
||||
def __call__(self, path, **extra_params):
|
||||
from datumaro.components.project import Project # cyclic import
|
||||
project = Project()
|
||||
|
||||
subsets = self.find_subsets(path)
|
||||
if len(subsets) == 0:
|
||||
raise Exception("Failed to find 'mot' dataset at '%s'" % path)
|
||||
|
||||
for ann_file in subsets:
|
||||
log.info("Found a dataset at '%s'" % ann_file)
|
||||
|
||||
source_name = osp.splitext(osp.basename(ann_file))[0]
|
||||
project.add_source(source_name, {
|
||||
'url': ann_file,
|
||||
'format': self._EXTRACTOR_NAME,
|
||||
'options': extra_params,
|
||||
})
|
||||
|
||||
return project
|
||||
|
||||
@staticmethod
|
||||
def find_subsets(path):
|
||||
subsets = []
|
||||
if path.endswith('.txt') and osp.isfile(path):
|
||||
subsets = [path]
|
||||
elif osp.isdir(path):
|
||||
p = osp.join(path, 'gt', MotPath.GT_FILENAME)
|
||||
if osp.isfile(p):
|
||||
subsets.append(p)
|
||||
return subsets
|
||||
|
||||
class MotSeqGtConverter(Converter, CliPlugin):
|
||||
@classmethod
|
||||
def build_cmdline_parser(cls, **kwargs):
|
||||
parser = super().__init__(**kwargs)
|
||||
parser.add_argument('--save-images', action='store_true',
|
||||
help="Save images (default: %(default)s)")
|
||||
return parser
|
||||
|
||||
def __init__(self, save_images=False):
|
||||
super().__init__()
|
||||
|
||||
self._save_images = save_images
|
||||
|
||||
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)
|
||||
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
|
||||
|
||||
writer.writerow({
|
||||
'frame_id': frame_id,
|
||||
'track_id': int(anno.attributes.get('track_id', -1)),
|
||||
'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:
|
||||
self._save_image(item, index=frame_id)
|
||||
else:
|
||||
log.debug("Item '%s' has no image" % item.id)
|
||||
|
||||
labels_file = osp.join(save_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 _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)
|
||||
@ -0,0 +1,146 @@
|
||||
import numpy as np
|
||||
|
||||
from unittest import TestCase
|
||||
|
||||
from datumaro.components.extractor import (Extractor, DatasetItem,
|
||||
AnnotationType, Bbox, LabelCategories
|
||||
)
|
||||
from datumaro.plugins.mot_format import MotSeqGtConverter, MotSeqImporter
|
||||
from datumaro.util.test_utils import TestDir, compare_datasets
|
||||
|
||||
|
||||
class MotConverterTest(TestCase):
|
||||
def _test_save_and_load(self, source_dataset, converter, test_dir,
|
||||
target_dataset=None, importer_args=None):
|
||||
converter(source_dataset, test_dir)
|
||||
|
||||
if importer_args is None:
|
||||
importer_args = {}
|
||||
parsed_dataset = MotSeqImporter()(test_dir, **importer_args) \
|
||||
.make_dataset()
|
||||
|
||||
if target_dataset is None:
|
||||
target_dataset = source_dataset
|
||||
|
||||
compare_datasets(self, expected=target_dataset, actual=parsed_dataset)
|
||||
|
||||
def test_can_save_bboxes(self):
|
||||
class SrcExtractor(Extractor):
|
||||
def __iter__(self):
|
||||
return iter([
|
||||
DatasetItem(id=1, subset='train',
|
||||
image=np.ones((16, 16, 3)),
|
||||
annotations=[
|
||||
Bbox(0, 4, 4, 8, label=2, attributes={
|
||||
'occluded': True,
|
||||
}),
|
||||
Bbox(0, 4, 4, 4, label=3, attributes={
|
||||
'visibility': 0.4,
|
||||
}),
|
||||
Bbox(2, 4, 4, 4, attributes={
|
||||
'ignored': True
|
||||
}),
|
||||
]
|
||||
),
|
||||
|
||||
DatasetItem(id=2, subset='val',
|
||||
image=np.ones((8, 8, 3)),
|
||||
annotations=[
|
||||
Bbox(1, 2, 4, 2, label=3),
|
||||
]
|
||||
),
|
||||
|
||||
DatasetItem(id=3, subset='test',
|
||||
image=np.ones((5, 4, 3)) * 3,
|
||||
),
|
||||
])
|
||||
|
||||
def categories(self):
|
||||
label_cat = LabelCategories()
|
||||
for label in range(10):
|
||||
label_cat.add('label_' + str(label))
|
||||
return {
|
||||
AnnotationType.label: label_cat,
|
||||
}
|
||||
|
||||
class DstExtractor(Extractor):
|
||||
def __iter__(self):
|
||||
return iter([
|
||||
DatasetItem(id=1,
|
||||
image=np.ones((16, 16, 3)),
|
||||
annotations=[
|
||||
Bbox(0, 4, 4, 8, label=2, attributes={
|
||||
'occluded': True,
|
||||
'visibility': 0.0,
|
||||
'ignored': False,
|
||||
}),
|
||||
Bbox(0, 4, 4, 4, label=3, attributes={
|
||||
'occluded': False,
|
||||
'visibility': 0.4,
|
||||
'ignored': False,
|
||||
}),
|
||||
Bbox(2, 4, 4, 4, attributes={
|
||||
'occluded': False,
|
||||
'visibility': 1.0,
|
||||
'ignored': True,
|
||||
}),
|
||||
]
|
||||
),
|
||||
|
||||
DatasetItem(id=2,
|
||||
image=np.ones((8, 8, 3)),
|
||||
annotations=[
|
||||
Bbox(1, 2, 4, 2, label=3, attributes={
|
||||
'occluded': False,
|
||||
'visibility': 1.0,
|
||||
'ignored': False,
|
||||
}),
|
||||
]
|
||||
),
|
||||
|
||||
DatasetItem(id=3,
|
||||
image=np.ones((5, 4, 3)) * 3,
|
||||
),
|
||||
])
|
||||
|
||||
def categories(self):
|
||||
label_cat = LabelCategories()
|
||||
for label in range(10):
|
||||
label_cat.add('label_' + str(label))
|
||||
return {
|
||||
AnnotationType.label: label_cat,
|
||||
}
|
||||
|
||||
with TestDir() as test_dir:
|
||||
self._test_save_and_load(
|
||||
SrcExtractor(), MotSeqGtConverter(save_images=True),
|
||||
test_dir, target_dataset=DstExtractor())
|
||||
|
||||
class MotImporterTest(TestCase):
|
||||
def test_can_detect(self):
|
||||
class TestExtractor(Extractor):
|
||||
def __iter__(self):
|
||||
return iter([
|
||||
DatasetItem(id=1, subset='train',
|
||||
image=np.ones((16, 16, 3)),
|
||||
annotations=[
|
||||
Bbox(0, 4, 4, 8, label=2),
|
||||
]
|
||||
),
|
||||
])
|
||||
|
||||
def categories(self):
|
||||
label_cat = LabelCategories()
|
||||
for label in range(10):
|
||||
label_cat.add('label_' + str(label))
|
||||
return {
|
||||
AnnotationType.label: label_cat,
|
||||
}
|
||||
|
||||
def generate_dummy_dataset(path):
|
||||
MotSeqGtConverter()(TestExtractor(), save_dir=path)
|
||||
|
||||
with TestDir() as test_dir:
|
||||
generate_dummy_dataset(test_dir)
|
||||
|
||||
self.assertTrue(MotSeqImporter.detect(test_dir))
|
||||
Loading…
Reference in New Issue