|
| 1 | +# Copyright (C) 2019 Intel Corporation |
| 2 | +# |
| 3 | +# SPDX-License-Identifier: MIT |
| 4 | + |
| 5 | +from tempfile import TemporaryDirectory |
| 6 | + |
| 7 | +from pyunpack import Archive |
| 8 | + |
| 9 | +from cvat.apps.dataset_manager.bindings import (CvatTaskDataExtractor, |
| 10 | + find_dataset_root, match_dm_item) |
| 11 | +from cvat.apps.dataset_manager.util import make_zip_archive |
| 12 | +from datumaro.components.extractor import AnnotationType, Transform |
| 13 | +from datumaro.components.project import Dataset |
| 14 | + |
| 15 | +from .registry import dm_env, exporter, importer |
| 16 | + |
| 17 | + |
| 18 | +class KeepTracks(Transform): |
| 19 | + def transform_item(self, item): |
| 20 | + return item.wrap(annotations=[a for a in item.annotations |
| 21 | + if 'track_id' in a.attributes]) |
| 22 | + |
| 23 | +@exporter(name='MOTS PNG', ext='ZIP', version='1.0') |
| 24 | +def _export(dst_file, task_data, save_images=False): |
| 25 | + extractor = CvatTaskDataExtractor(task_data, include_images=save_images) |
| 26 | + envt = dm_env.transforms |
| 27 | + extractor = extractor.transform(KeepTracks) # can only export tracks |
| 28 | + extractor = extractor.transform(envt.get('polygons_to_masks')) |
| 29 | + extractor = extractor.transform(envt.get('boxes_to_masks')) |
| 30 | + extractor = extractor.transform(envt.get('merge_instance_segments')) |
| 31 | + extractor = Dataset.from_extractors(extractor) # apply lazy transforms |
| 32 | + with TemporaryDirectory() as temp_dir: |
| 33 | + dm_env.converters.get('mots_png').convert(extractor, |
| 34 | + save_dir=temp_dir, save_images=save_images) |
| 35 | + |
| 36 | + make_zip_archive(temp_dir, dst_file) |
| 37 | + |
| 38 | +@importer(name='MOTS PNG', ext='ZIP', version='1.0') |
| 39 | +def _import(src_file, task_data): |
| 40 | + with TemporaryDirectory() as tmp_dir: |
| 41 | + Archive(src_file.name).extractall(tmp_dir) |
| 42 | + |
| 43 | + dataset = dm_env.make_importer('mots')(tmp_dir).make_dataset() |
| 44 | + masks_to_polygons = dm_env.transforms.get('masks_to_polygons') |
| 45 | + dataset = dataset.transform(masks_to_polygons) |
| 46 | + |
| 47 | + tracks = {} |
| 48 | + label_cat = dataset.categories()[AnnotationType.label] |
| 49 | + |
| 50 | + root_hint = find_dataset_root(dataset, task_data) |
| 51 | + |
| 52 | + for item in dataset: |
| 53 | + frame_number = task_data.abs_frame_id( |
| 54 | + match_dm_item(item, task_data, root_hint=root_hint)) |
| 55 | + |
| 56 | + for ann in item.annotations: |
| 57 | + if ann.type != AnnotationType.polygon: |
| 58 | + continue |
| 59 | + |
| 60 | + track_id = ann.attributes['track_id'] |
| 61 | + shape = task_data.TrackedShape( |
| 62 | + type='polygon', |
| 63 | + points=ann.points, |
| 64 | + occluded=ann.attributes.get('occluded') == True, |
| 65 | + outside=False, |
| 66 | + keyframe=True, |
| 67 | + z_order=ann.z_order, |
| 68 | + frame=frame_number, |
| 69 | + attributes=[], |
| 70 | + source='manual', |
| 71 | + ) |
| 72 | + |
| 73 | + # build trajectories as lists of shapes in track dict |
| 74 | + if track_id not in tracks: |
| 75 | + tracks[track_id] = task_data.Track( |
| 76 | + label_cat.items[ann.label].name, 0, 'manual', []) |
| 77 | + tracks[track_id].shapes.append(shape) |
| 78 | + |
| 79 | + for track in tracks.values(): |
| 80 | + track.shapes.sort(key=lambda t: t.frame) |
| 81 | + |
| 82 | + # insert outside=True in skips between the frames track is visible |
| 83 | + prev_shape_idx = 0 |
| 84 | + prev_shape = track.shapes[0] |
| 85 | + for shape in track.shapes[1:]: |
| 86 | + has_skip = task_data.frame_step < shape.frame - prev_shape.frame |
| 87 | + if has_skip and not prev_shape.outside: |
| 88 | + prev_shape = prev_shape._replace(outside=True, |
| 89 | + frame=prev_shape.frame + task_data.frame_step) |
| 90 | + prev_shape_idx += 1 |
| 91 | + track.shapes.insert(prev_shape_idx, prev_shape) |
| 92 | + prev_shape = shape |
| 93 | + prev_shape_idx += 1 |
| 94 | + |
| 95 | + # Append a shape with outside=True to finish the track |
| 96 | + last_shape = track.shapes[-1] |
| 97 | + if last_shape.frame + task_data.frame_step <= \ |
| 98 | + int(task_data.meta['task']['stop_frame']): |
| 99 | + track.shapes.append(last_shape._replace(outside=True, |
| 100 | + frame=last_shape.frame + task_data.frame_step) |
| 101 | + ) |
| 102 | + task_data.add_track(track) |
0 commit comments