diff --git a/CHANGELOG.md b/CHANGELOG.md index 6fd0f3a3db3..b4daaa00ee9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ https://github.com/opencv/cvat/issues/750). - Mask-RCNN Auto Annotation Script in OpenVINO format - Yolo Auto Annotation Script - Auto segmentation using Mask_RCNN component (Keras+Tensorflow Mask R-CNN Segmentation) +- Added MOT CSV format support - Ability to dump/load annotations in LabelMe format from UI ### Changed diff --git a/README.md b/README.md index 9bacba82e98..b3d0480acfd 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,7 @@ Format selection is possible after clicking on the Upload annotation / Dump anno | [MS COCO Object Detection](http://cocodataset.org/#format-data) | X | X | | PNG mask | X | | | [TFrecord](https://www.tensorflow.org/tutorials/load_data/tf_records) | X | X | +| [MOT](https://motchallenge.net/) | X | X | | [LabelMe](http://labelme.csail.mit.edu/Release3.0) | X | X | ## Links diff --git a/cvat/apps/annotation/mot.py b/cvat/apps/annotation/mot.py new file mode 100644 index 00000000000..b7f63d1e79b --- /dev/null +++ b/cvat/apps/annotation/mot.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: MIT +format_spec = { + "name": "MOT", + "dumpers": [ + { + "display_name": "{name} {format} {version}", + "format": "CSV", + "version": "1.0", + "handler": "dump" + }, + ], + "loaders": [ + { + "display_name": "{name} {format} {version}", + "format": "CSV", + "version": "1.0", + "handler": "load", + } + ], +} + + +MOT = [ + "frame_id", + "track_id", + "xtl", + "ytl", + "width", + "height", + "confidence", + "class_id", + "visibility" +] + + +def dump(file_object, annotations): + """ Export track shapes in MOT CSV format. Due to limitations of the MOT + format, this process only supports rectangular interpolation mode + annotations. + """ + import csv + import io + + # csv requires a text buffer + with io.TextIOWrapper(file_object, encoding="utf-8") as csv_file: + writer = csv.DictWriter(csv_file, fieldnames=MOT) + for i, track in enumerate(annotations.tracks): + for shape in track.shapes: + # MOT doesn't support polygons or 'outside' property + if shape.type != 'rectangle': + continue + writer.writerow({ + "frame_id": shape.frame, + "track_id": i, + "xtl": shape.points[0], + "ytl": shape.points[1], + "width": shape.points[2] - shape.points[0], + "height": shape.points[3] - shape.points[1], + "confidence": 1, + "class_id": track.label, + "visibility": 1 - int(shape.occluded) + }) + + +def load(file_object, annotations): + """ Read MOT CSV format and convert objects to annotated tracks. + """ + import csv + import io + tracks = {} + # csv requires a text buffer + with io.TextIOWrapper(file_object, encoding="utf-8") as csv_file: + reader = csv.DictReader(csv_file, fieldnames=MOT) + for row in reader: + # create one shape per row + xtl = float(row["xtl"]) + ytl = float(row["ytl"]) + xbr = xtl + float(row["width"]) + ybr = ytl + float(row["height"]) + shape = annotations.TrackedShape( + type="rectangle", + points=[xtl, ytl, xbr, ybr], + occluded=float(row["visibility"]) == 0, + outside=False, + keyframe=False, + z_order=0, + frame=int(row["frame_id"]), + attributes=[], + ) + # build trajectories as lists of shapes in track dict + track_id = int(row["track_id"]) + if track_id not in tracks: + tracks[track_id] = annotations.Track(row["class_id"], track_id, []) + tracks[track_id].shapes.append(shape) + for track in tracks.values(): + # Set outside=True for the last shape since MOT has no support + # for this flag + last = annotations.TrackedShape( + type=track.shapes[-1].type, + points=track.shapes[-1].points, + occluded=track.shapes[-1].occluded, + outside=True, + keyframe=track.shapes[-1].keyframe, + z_order=track.shapes[-1].z_order, + frame=track.shapes[-1].frame, + attributes=track.shapes[-1].attributes, + ) + track.shapes[-1] = last + annotations.add_track(track) diff --git a/cvat/apps/annotation/settings.py b/cvat/apps/annotation/settings.py index 1c42bf3b04a..e1e5f82b42c 100644 --- a/cvat/apps/annotation/settings.py +++ b/cvat/apps/annotation/settings.py @@ -12,5 +12,6 @@ os.path.join(path_prefix, 'coco.py'), os.path.join(path_prefix, 'mask.py'), os.path.join(path_prefix, 'tfrecord.py'), + os.path.join(path_prefix, 'mot.py'), os.path.join(path_prefix, 'labelme.py'), ) diff --git a/cvat/apps/engine/tests/test_rest_api.py b/cvat/apps/engine/tests/test_rest_api.py index ae109bc7272..15bfea7abad 100644 --- a/cvat/apps/engine/tests/test_rest_api.py +++ b/cvat/apps/engine/tests/test_rest_api.py @@ -2662,6 +2662,9 @@ def _get_initial_annotation(annotation_format): annotations["shapes"] = rectangle_shapes_with_attrs + rectangle_shapes_wo_attrs + polygon_shapes_wo_attrs annotations["tracks"] = rectangle_tracks_with_attrs + rectangle_tracks_wo_attrs + elif annotation_format == "MOT CSV 1.0": + annotations["tracks"] = rectangle_tracks_wo_attrs + return annotations response = self._get_annotation_formats(annotator)