-
Notifications
You must be signed in to change notification settings - Fork 2.9k
/
mot.py
87 lines (71 loc) · 3.37 KB
/
mot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
# Copyright (C) 2019 Intel Corporation
#
# SPDX-License-Identifier: MIT
from tempfile import TemporaryDirectory
from pyunpack import Archive
import datumaro.components.extractor as datumaro
from cvat.apps.dataset_manager.bindings import CvatTaskDataExtractor
from cvat.apps.dataset_manager.util import make_zip_archive
from datumaro.components.project import Dataset
from .registry import dm_env, exporter, importer
@exporter(name='MOT', ext='ZIP', version='1.1')
def _export(dst_file, task_data, save_images=False):
extractor = CvatTaskDataExtractor(task_data, include_images=save_images)
extractor = Dataset.from_extractors(extractor) # apply lazy transforms
with TemporaryDirectory() as temp_dir:
dm_env.converters.get('mot_seq_gt').convert(extractor,
save_dir=temp_dir, save_images=save_images)
make_zip_archive(temp_dir, dst_file)
@importer(name='MOT', ext='ZIP', version='1.1')
def _import(src_file, task_data):
with TemporaryDirectory() as tmp_dir:
Archive(src_file.name).extractall(tmp_dir)
dataset = dm_env.make_importer('mot_seq')(tmp_dir).make_dataset()
tracks = {}
label_cat = dataset.categories()[datumaro.AnnotationType.label]
for item in dataset:
frame_number = int(item.id) - 1 # NOTE: MOT frames start from 1
frame_number = task_data.abs_frame_id(frame_number)
for ann in item.annotations:
if ann.type != datumaro.AnnotationType.bbox:
continue
track_id = ann.attributes.get('track_id')
if track_id is None:
# Extension. Import regular boxes:
task_data.add_shape(task_data.LabeledShape(
type='rectangle',
label=label_cat.items[ann.label].name,
points=ann.points,
occluded=ann.attributes.get('occluded') == True,
z_order=ann.z_order,
group=0,
frame=frame_number,
attributes=[],
))
continue
shape = task_data.TrackedShape(
type='rectangle',
points=ann.points,
occluded=ann.attributes.get('occluded') == True,
outside=False,
keyframe=True,
z_order=ann.z_order,
frame=frame_number,
attributes=[],
)
# build trajectories as lists of shapes in track dict
if track_id not in tracks:
tracks[track_id] = task_data.Track(
label_cat.items[ann.label].name, 0, [])
tracks[track_id].shapes.append(shape)
for track in tracks.values():
# MOT annotations do not require frames to be ordered
track.shapes.sort(key=lambda t: t.frame)
# Append a shape with outside=True to finish the track
last_shape = track.shapes[-1]
if last_shape.frame + task_data.frame_step <= \
int(task_data.meta['task']['stop_frame']):
track.shapes.append(last_shape._replace(outside=True,
frame=last_shape.frame + task_data.frame_step)
)
task_data.add_track(track)