Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding dump and load support for MOT CSV format. #830

Merged
merged 6 commits into from
Nov 22, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
109 changes: 109 additions & 0 deletions cvat/apps/annotation/mot.py
Original file line number Diff line number Diff line change
@@ -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)
1 change: 1 addition & 0 deletions cvat/apps/annotation/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'),
)
3 changes: 3 additions & 0 deletions cvat/apps/engine/tests/test_rest_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down