Add MOTS png format (#2198)
* Add mots format * fix upload * update docs * update changelog * Update datumaro dependency * fix header * update dm dependency * Support importing with outside property in mot and mots * fix track exporting Co-authored-by: Boris Sekachev <boris.sekachev@yandex.ru>main
Maxim Zhiltsov
5 years ago
committed by
GitHub
parent
d957d6a9ff
commit
d4129f2868
@ -0,0 +1,102 @@
|
|||||||
|
# Copyright (C) 2019 Intel Corporation
|
||||||
|
#
|
||||||
|
# SPDX-License-Identifier: MIT
|
||||||
|
|
||||||
|
from tempfile import TemporaryDirectory
|
||||||
|
|
||||||
|
from pyunpack import Archive
|
||||||
|
|
||||||
|
from cvat.apps.dataset_manager.bindings import (CvatTaskDataExtractor,
|
||||||
|
find_dataset_root, match_dm_item)
|
||||||
|
from cvat.apps.dataset_manager.util import make_zip_archive
|
||||||
|
from datumaro.components.extractor import AnnotationType, Transform
|
||||||
|
from datumaro.components.project import Dataset
|
||||||
|
|
||||||
|
from .registry import dm_env, exporter, importer
|
||||||
|
|
||||||
|
|
||||||
|
class KeepTracks(Transform):
|
||||||
|
def transform_item(self, item):
|
||||||
|
return item.wrap(annotations=[a for a in item.annotations
|
||||||
|
if 'track_id' in a.attributes])
|
||||||
|
|
||||||
|
@exporter(name='MOTS PNG', ext='ZIP', version='1.0')
|
||||||
|
def _export(dst_file, task_data, save_images=False):
|
||||||
|
extractor = CvatTaskDataExtractor(task_data, include_images=save_images)
|
||||||
|
envt = dm_env.transforms
|
||||||
|
extractor = extractor.transform(KeepTracks) # can only export tracks
|
||||||
|
extractor = extractor.transform(envt.get('polygons_to_masks'))
|
||||||
|
extractor = extractor.transform(envt.get('boxes_to_masks'))
|
||||||
|
extractor = extractor.transform(envt.get('merge_instance_segments'))
|
||||||
|
extractor = Dataset.from_extractors(extractor) # apply lazy transforms
|
||||||
|
with TemporaryDirectory() as temp_dir:
|
||||||
|
dm_env.converters.get('mots_png').convert(extractor,
|
||||||
|
save_dir=temp_dir, save_images=save_images)
|
||||||
|
|
||||||
|
make_zip_archive(temp_dir, dst_file)
|
||||||
|
|
||||||
|
@importer(name='MOTS PNG', ext='ZIP', version='1.0')
|
||||||
|
def _import(src_file, task_data):
|
||||||
|
with TemporaryDirectory() as tmp_dir:
|
||||||
|
Archive(src_file.name).extractall(tmp_dir)
|
||||||
|
|
||||||
|
dataset = dm_env.make_importer('mots')(tmp_dir).make_dataset()
|
||||||
|
masks_to_polygons = dm_env.transforms.get('masks_to_polygons')
|
||||||
|
dataset = dataset.transform(masks_to_polygons)
|
||||||
|
|
||||||
|
tracks = {}
|
||||||
|
label_cat = dataset.categories()[AnnotationType.label]
|
||||||
|
|
||||||
|
root_hint = find_dataset_root(dataset, task_data)
|
||||||
|
|
||||||
|
for item in dataset:
|
||||||
|
frame_number = task_data.abs_frame_id(
|
||||||
|
match_dm_item(item, task_data, root_hint=root_hint))
|
||||||
|
|
||||||
|
for ann in item.annotations:
|
||||||
|
if ann.type != AnnotationType.polygon:
|
||||||
|
continue
|
||||||
|
|
||||||
|
track_id = ann.attributes['track_id']
|
||||||
|
shape = task_data.TrackedShape(
|
||||||
|
type='polygon',
|
||||||
|
points=ann.points,
|
||||||
|
occluded=ann.attributes.get('occluded') == True,
|
||||||
|
outside=False,
|
||||||
|
keyframe=True,
|
||||||
|
z_order=ann.z_order,
|
||||||
|
frame=frame_number,
|
||||||
|
attributes=[],
|
||||||
|
source='manual',
|
||||||
|
)
|
||||||
|
|
||||||
|
# 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, 'manual', [])
|
||||||
|
tracks[track_id].shapes.append(shape)
|
||||||
|
|
||||||
|
for track in tracks.values():
|
||||||
|
track.shapes.sort(key=lambda t: t.frame)
|
||||||
|
|
||||||
|
# insert outside=True in skips between the frames track is visible
|
||||||
|
prev_shape_idx = 0
|
||||||
|
prev_shape = track.shapes[0]
|
||||||
|
for shape in track.shapes[1:]:
|
||||||
|
has_skip = task_data.frame_step < shape.frame - prev_shape.frame
|
||||||
|
if has_skip and not prev_shape.outside:
|
||||||
|
prev_shape = prev_shape._replace(outside=True,
|
||||||
|
frame=prev_shape.frame + task_data.frame_step)
|
||||||
|
prev_shape_idx += 1
|
||||||
|
track.shapes.insert(prev_shape_idx, prev_shape)
|
||||||
|
prev_shape = shape
|
||||||
|
prev_shape_idx += 1
|
||||||
|
|
||||||
|
# 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)
|
||||||
Loading…
Reference in New Issue