From fa8949fb8131f44a32735aa311a4645a0cec05f0 Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Fri, 18 Sep 2020 17:35:21 +0300 Subject: [PATCH 1/9] Add mots format --- cvat/apps/dataset_manager/formats/mots.py | 82 +++++++++++++++++++ cvat/apps/dataset_manager/formats/registry.py | 1 + .../dataset_manager/tests/_test_formats.py | 3 + cvat/apps/engine/tests/_test_rest_api.py | 67 ++++++++++++--- 4 files changed, 140 insertions(+), 13 deletions(-) create mode 100644 cvat/apps/dataset_manager/formats/mots.py diff --git a/cvat/apps/dataset_manager/formats/mots.py b/cvat/apps/dataset_manager/formats/mots.py new file mode 100644 index 000000000000..8e5ba11b1b42 --- /dev/null +++ b/cvat/apps/dataset_manager/formats/mots.py @@ -0,0 +1,82 @@ +# 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 +from datumaro.components.project import Dataset + +from .registry import dm_env, exporter, importer + + +@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(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='rectangle', + 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) + # 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) diff --git a/cvat/apps/dataset_manager/formats/registry.py b/cvat/apps/dataset_manager/formats/registry.py index ed4defc559df..c84d60fc603e 100644 --- a/cvat/apps/dataset_manager/formats/registry.py +++ b/cvat/apps/dataset_manager/formats/registry.py @@ -87,6 +87,7 @@ def make_exporter(name): import cvat.apps.dataset_manager.formats.labelme import cvat.apps.dataset_manager.formats.mask import cvat.apps.dataset_manager.formats.mot +import cvat.apps.dataset_manager.formats.mots import cvat.apps.dataset_manager.formats.pascal_voc import cvat.apps.dataset_manager.formats.tfrecord import cvat.apps.dataset_manager.formats.yolo \ No newline at end of file diff --git a/cvat/apps/dataset_manager/tests/_test_formats.py b/cvat/apps/dataset_manager/tests/_test_formats.py index 3b4ea5487ae0..00b1277bbda0 100644 --- a/cvat/apps/dataset_manager/tests/_test_formats.py +++ b/cvat/apps/dataset_manager/tests/_test_formats.py @@ -327,6 +327,7 @@ def test_export_formats_query(self): 'Datumaro 1.0', 'LabelMe 3.0', 'MOT 1.1', + 'MOTS PNG 1.0', 'PASCAL VOC 1.1', 'Segmentation mask 1.1', 'TFRecord 1.0', @@ -342,6 +343,7 @@ def test_import_formats_query(self): 'CVAT 1.1', 'LabelMe 3.0', 'MOT 1.1', + 'MOTS PNG 1.0', 'PASCAL VOC 1.1', 'Segmentation mask 1.1', 'TFRecord 1.0', @@ -376,6 +378,7 @@ def test_empty_images_are_exported(self): ('Datumaro 1.0', 'datumaro_project'), ('LabelMe 3.0', 'label_me'), # ('MOT 1.1', 'mot_seq'), # does not support + # ('MOTS PNG 1.0', 'mots_png'), # does not support ('PASCAL VOC 1.1', 'voc'), ('Segmentation mask 1.1', 'voc'), ('TFRecord 1.0', 'tf_detection_api'), diff --git a/cvat/apps/engine/tests/_test_rest_api.py b/cvat/apps/engine/tests/_test_rest_api.py index 374a4cd2e213..333f27280fe6 100644 --- a/cvat/apps/engine/tests/_test_rest_api.py +++ b/cvat/apps/engine/tests/_test_rest_api.py @@ -3137,6 +3137,39 @@ def _get_initial_annotation(annotation_format): } ] }] + polygon_tracks_wo_attrs = [{ + "frame": 0, + "label_id": task["labels"][1]["id"], + "group": 0, + "source": "manual", + "attributes": [], + "shapes": [ + { + "frame": 0, + "attributes": [], + "points": [1.0, 2.1, 50.2, 36.6, 7.0, 10.0], + "type": "polygon", + "occluded": False, + "outside": False, + }, + { + "frame": 1, + "attributes": [], + "points": [1.0, 2.1, 51, 36.6, 8.0, 11.0], + "type": "polygon", + "occluded": False, + "outside": False + }, + { + "frame": 2, + "attributes": [], + "points": [1.0, 2.1, 51, 36.6, 14.0, 15.0], + "type": "polygon", + "occluded": False, + "outside": True, + } + ] + }] rectangle_shapes_with_attrs = [{ "frame": 0, @@ -3241,11 +3274,15 @@ def _get_initial_annotation(annotation_format): "tracks": [], } if annotation_format == "CVAT for video 1.1": - annotations["tracks"] = rectangle_tracks_with_attrs + rectangle_tracks_wo_attrs + annotations["tracks"] = rectangle_tracks_with_attrs \ + + rectangle_tracks_wo_attrs \ + + polygon_tracks_wo_attrs elif annotation_format == "CVAT for images 1.1": - annotations["shapes"] = rectangle_shapes_with_attrs + rectangle_shapes_wo_attrs \ - + polygon_shapes_wo_attrs + polygon_shapes_with_attrs + annotations["shapes"] = rectangle_shapes_with_attrs \ + + rectangle_shapes_wo_attrs \ + + polygon_shapes_wo_attrs \ + + polygon_shapes_with_attrs annotations["tags"] = tags_with_attrs + tags_wo_attrs elif annotation_format == "PASCAL VOC 1.1": @@ -3260,24 +3297,28 @@ def _get_initial_annotation(annotation_format): annotations["shapes"] = polygon_shapes_wo_attrs elif annotation_format == "Segmentation mask 1.1": - annotations["shapes"] = rectangle_shapes_wo_attrs + polygon_shapes_wo_attrs + annotations["shapes"] = rectangle_shapes_wo_attrs \ + + polygon_shapes_wo_attrs annotations["tracks"] = rectangle_tracks_wo_attrs elif annotation_format == "MOT 1.1": annotations["shapes"] = rectangle_shapes_wo_attrs annotations["tracks"] = rectangle_tracks_wo_attrs + elif annotation_format == "MOTS PNG 1.0": + annotations["tracks"] = polygon_tracks_wo_attrs + elif annotation_format == "LabelMe 3.0": - annotations["shapes"] = rectangle_shapes_with_attrs + \ - rectangle_shapes_wo_attrs + \ - polygon_shapes_wo_attrs + \ - polygon_shapes_with_attrs + annotations["shapes"] = rectangle_shapes_with_attrs \ + + rectangle_shapes_wo_attrs \ + + polygon_shapes_wo_attrs \ + + polygon_shapes_with_attrs elif annotation_format == "Datumaro 1.0": - annotations["shapes"] = rectangle_shapes_with_attrs + \ - rectangle_shapes_wo_attrs + \ - polygon_shapes_wo_attrs + \ - polygon_shapes_with_attrs + annotations["shapes"] = rectangle_shapes_with_attrs \ + + rectangle_shapes_wo_attrs \ + + polygon_shapes_wo_attrs \ + + polygon_shapes_with_attrs annotations["tags"] = tags_with_attrs + tags_wo_attrs else: @@ -3376,7 +3417,7 @@ def _get_initial_annotation(annotation_format): self.assertEqual(response.status_code, HTTP_201_CREATED) # 7. check annotation - if import_format == "Segmentation mask 1.1": + if import_format in {"Segmentation mask 1.1", "MOTS PNG 1.0"}: continue # can't really predict the result to check response = self._get_api_v1_tasks_id_annotations(task["id"], annotator) self.assertEqual(response.status_code, HTTP_200_OK) From 72d122125d2279c542065a43d4d57d342f3c8aaf Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Fri, 18 Sep 2020 17:58:42 +0300 Subject: [PATCH 2/9] fix upload --- cvat/apps/dataset_manager/formats/mots.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/cvat/apps/dataset_manager/formats/mots.py b/cvat/apps/dataset_manager/formats/mots.py index 8e5ba11b1b42..76a843d272e7 100644 --- a/cvat/apps/dataset_manager/formats/mots.py +++ b/cvat/apps/dataset_manager/formats/mots.py @@ -9,16 +9,22 @@ 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 +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')) @@ -53,7 +59,7 @@ def _import(src_file, task_data): track_id = ann.attributes['track_id'] shape = task_data.TrackedShape( - type='rectangle', + type='polygon', points=ann.points, occluded=ann.attributes.get('occluded') == True, outside=False, From 17fe656ac115fe7482b902393d0aa1e2282067a5 Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Fri, 18 Sep 2020 18:06:55 +0300 Subject: [PATCH 3/9] update docs --- cvat/apps/dataset_manager/formats/README.md | 37 +++++++++++++++++++-- 1 file changed, 35 insertions(+), 2 deletions(-) diff --git a/cvat/apps/dataset_manager/formats/README.md b/cvat/apps/dataset_manager/formats/README.md index 2879eaa4b3da..ec412c3eec72 100644 --- a/cvat/apps/dataset_manager/formats/README.md +++ b/cvat/apps/dataset_manager/formats/README.md @@ -13,6 +13,7 @@ - [CVAT](#cvat) - [LabelMe](#labelme) - [MOT](#mot) + - [MOTS](#mots) - [COCO](#coco) - [PASCAL VOC and mask](#voc) - [YOLO](#yolo) @@ -708,8 +709,8 @@ Downloaded file: a zip archive of the following structure: ``` bash taskname.zip/ ├── img1/ -| ├── imgage1.jpg -| └── imgage2.jpg +| ├── image1.jpg +| └── image2.jpg └── gt/ ├── labels.txt └── gt.txt @@ -742,6 +743,38 @@ taskname.zip/ - supported annotations: Rectangle tracks +### [MOTS PNG](https://www.vision.rwth-aachen.de/page/mots) + +#### MOTS PNG Dumper + +Downloaded file: a zip archive of the following structure: + +``` bash +taskname.zip/ +└── / + | images/ + | ├── image1.jpg + | └── image2.jpg + └── instances/ + ├── labels.txt + ├── image1.png + └── image2.png + +# labels.txt +cat +dog +person +... +``` + +- supported annotations: Rectangle and Polygon tracks + +#### MOT Loader + +Uploaded file: a zip archive of the structure above + +- supported annotations: Polygon tracks + ### [LabelMe](http://labelme.csail.mit.edu/Release3.0) #### LabelMe Dumper From a556b0f784a654758b633b728357fc2be42aa2e8 Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Fri, 18 Sep 2020 18:13:45 +0300 Subject: [PATCH 4/9] update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index ec03b776985e..4098884b4f05 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - [Datumaro] CLI command for dataset equality comparison () - [Datumaro] Merging of datasets with different labels () - Add FBRS interactive segmentation serverless function () +- MOTS png mask format support () ### Changed - UI models (like DEXTR) were redesigned to be more interactive () From 221d33cb3496870bb66071b3882a8c7e441309dc Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Thu, 24 Sep 2020 17:05:09 +0300 Subject: [PATCH 5/9] Update datumaro dependency --- cvat/requirements/base.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/requirements/base.txt b/cvat/requirements/base.txt index ff8111ca351b..b1781436e995 100644 --- a/cvat/requirements/base.txt +++ b/cvat/requirements/base.txt @@ -44,4 +44,4 @@ tensorflow==2.2.0 # Optional requirement of Datumaro # archives. Don't use as a python module because it has GPL license. patool==1.12 diskcache==5.0.2 -git+https://github.com/openvinotoolkit/datumaro@v0.1.0 \ No newline at end of file +git+https://github.com/openvinotoolkit/datumaro@v0.1.1 \ No newline at end of file From d1062125cc81f125df07719b6ada9913a0ea0913 Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Thu, 24 Sep 2020 17:22:25 +0300 Subject: [PATCH 6/9] fix header --- cvat/apps/dataset_manager/formats/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/dataset_manager/formats/README.md b/cvat/apps/dataset_manager/formats/README.md index ec412c3eec72..ecd703fca1b0 100644 --- a/cvat/apps/dataset_manager/formats/README.md +++ b/cvat/apps/dataset_manager/formats/README.md @@ -769,7 +769,7 @@ person - supported annotations: Rectangle and Polygon tracks -#### MOT Loader +#### MOTS PNG Loader Uploaded file: a zip archive of the structure above From b71280e983a5136634f5eb0d750a3efa3bc544a2 Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Mon, 5 Oct 2020 16:09:21 +0300 Subject: [PATCH 7/9] update dm dependency --- cvat/requirements/base.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/requirements/base.txt b/cvat/requirements/base.txt index 0b8463c2e63d..60613b4356ca 100644 --- a/cvat/requirements/base.txt +++ b/cvat/requirements/base.txt @@ -44,4 +44,4 @@ tensorflow==2.2.1 # Optional requirement of Datumaro # archives. Don't use as a python module because it has GPL license. patool==1.12 diskcache==5.0.2 -git+https://github.com/openvinotoolkit/datumaro@v0.1.1 \ No newline at end of file +git+https://github.com/openvinotoolkit/datumaro@v0.1.2 \ No newline at end of file From 191dd1f1a1c335775c3bbb1508fe2a1579a8725e Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Tue, 6 Oct 2020 16:04:25 +0300 Subject: [PATCH 8/9] Support importing with outside property in mot and mots --- cvat/apps/dataset_manager/formats/mot.py | 14 ++++++++++++++ cvat/apps/dataset_manager/formats/mots.py | 14 ++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/cvat/apps/dataset_manager/formats/mot.py b/cvat/apps/dataset_manager/formats/mot.py index b854ddb2b8bb..81131dc15040 100644 --- a/cvat/apps/dataset_manager/formats/mot.py +++ b/cvat/apps/dataset_manager/formats/mot.py @@ -79,6 +79,20 @@ def _import(src_file, task_data): for track in tracks.values(): # MOT annotations do not require frames to be ordered 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 <= \ diff --git a/cvat/apps/dataset_manager/formats/mots.py b/cvat/apps/dataset_manager/formats/mots.py index 76a843d272e7..52bf0fa623ac 100644 --- a/cvat/apps/dataset_manager/formats/mots.py +++ b/cvat/apps/dataset_manager/formats/mots.py @@ -78,6 +78,20 @@ def _import(src_file, task_data): 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 <= \ From 2799f4b4cdb9d160ca71e964555397dc302ef5a2 Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Tue, 6 Oct 2020 16:06:40 +0300 Subject: [PATCH 9/9] fix track exporting --- CHANGELOG.md | 2 + cvat/apps/dataset_manager/annotation.py | 5 +- .../dataset_manager/tests/test_annotation.py | 67 +++++++++++++------ 3 files changed, 50 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 31ed442a0272..4e289a8cea85 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,7 @@ filters and searching the nearest frame without any annotations ( Alt] () - Updated `docker-compose` file version from `2.3` to `3.3`() - Added auto inference of url schema from host in CLI, if provided () +- Track frames in skips between annotation is presented in MOT and MOTS formats are marked `outside` () ### Deprecated - @@ -53,6 +54,7 @@ filters and searching the nearest frame without any annotations () - Fixed issues from #2112 () - Git application name (renamed to dataset_repo) () +- A problem in exporting of tracks, where tracks could be truncated () ### Security - diff --git a/cvat/apps/dataset_manager/annotation.py b/cvat/apps/dataset_manager/annotation.py index 30393dd56f14..b6dbe1c17482 100644 --- a/cvat/apps/dataset_manager/annotation.py +++ b/cvat/apps/dataset_manager/annotation.py @@ -729,7 +729,6 @@ def interpolate(shape0, shape1): if track.get("interpolated_shapes"): return track["interpolated_shapes"] - # TODO: should be return an iterator? shapes = [] curr_frame = track["shapes"][0]["frame"] prev_shape = {} @@ -747,9 +746,7 @@ def interpolate(shape0, shape1): curr_frame = shape["frame"] prev_shape = shape - # TODO: Need to modify a client and a database (append "outside" shapes for polytracks) - if not prev_shape["outside"] and (prev_shape["type"] == ShapeType.RECTANGLE - or prev_shape["type"] == ShapeType.POINTS or prev_shape["type"] == ShapeType.CUBOID): + if not prev_shape["outside"]: shape = copy(prev_shape) shape["frame"] = end_frame shapes.extend(interpolate(prev_shape, shape)) diff --git a/cvat/apps/dataset_manager/tests/test_annotation.py b/cvat/apps/dataset_manager/tests/test_annotation.py index 7604a4dc4cda..4f4dc84c123c 100644 --- a/cvat/apps/dataset_manager/tests/test_annotation.py +++ b/cvat/apps/dataset_manager/tests/test_annotation.py @@ -8,6 +8,17 @@ class TrackManagerTest(TestCase): + def _check_interpolation(self, track): + interpolated = TrackManager.get_interpolated_shapes(track, 0, 7) + + self.assertEqual(len(interpolated), 6) + self.assertTrue(interpolated[0]["keyframe"]) + self.assertFalse(interpolated[1]["keyframe"]) + self.assertTrue(interpolated[2]["keyframe"]) + self.assertTrue(interpolated[3]["keyframe"]) + self.assertFalse(interpolated[4]["keyframe"]) + self.assertFalse(interpolated[5]["keyframe"]) + def test_point_interpolation(self): track = { "frame": 0, @@ -32,14 +43,18 @@ def test_point_interpolation(self): "occluded": False, "outside": True }, + { + "frame": 4, + "attributes": [], + "points": [3.0, 4.0, 5.0, 6.0], + "type": "points", + "occluded": False, + "outside": False + }, ] } - interpolated = TrackManager.get_interpolated_shapes(track, 0, 2) - - self.assertEqual(len(interpolated), 3) - self.assertTrue(interpolated[0]["keyframe"]) - self.assertFalse(interpolated[1]["keyframe"]) + self._check_interpolation(track) def test_polygon_interpolation(self): track = { @@ -65,14 +80,18 @@ def test_polygon_interpolation(self): "occluded": False, "outside": True }, + { + "frame": 4, + "attributes": [], + "points": [3.0, 4.0, 5.0, 6.0, 7.0, 6.0, 4.0, 5.0], + "type": "polygon", + "occluded": False, + "outside": False + }, ] } - interpolated = TrackManager.get_interpolated_shapes(track, 0, 2) - - self.assertEqual(len(interpolated), 3) - self.assertTrue(interpolated[0]["keyframe"]) - self.assertFalse(interpolated[1]["keyframe"]) + self._check_interpolation(track) def test_bbox_interpolation(self): track = { @@ -98,14 +117,18 @@ def test_bbox_interpolation(self): "occluded": False, "outside": True }, + { + "frame": 4, + "attributes": [], + "points": [3.0, 4.0, 5.0, 6.0], + "type": "rectangle", + "occluded": False, + "outside": False + }, ] } - interpolated = TrackManager.get_interpolated_shapes(track, 0, 2) - - self.assertEqual(len(interpolated), 3) - self.assertTrue(interpolated[0]["keyframe"]) - self.assertFalse(interpolated[1]["keyframe"]) + self._check_interpolation(track) def test_line_interpolation(self): track = { @@ -131,11 +154,15 @@ def test_line_interpolation(self): "occluded": False, "outside": True }, + { + "frame": 4, + "attributes": [], + "points": [3.0, 4.0, 5.0, 6.0], + "type": "polyline", + "occluded": False, + "outside": False + }, ] } - interpolated = TrackManager.get_interpolated_shapes(track, 0, 2) - - self.assertEqual(len(interpolated), 3) - self.assertTrue(interpolated[0]["keyframe"]) - self.assertFalse(interpolated[1]["keyframe"]) \ No newline at end of file + self._check_interpolation(track) \ No newline at end of file