Skip to content

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

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

Merged
merged 6 commits into from
Nov 22, 2019
Merged
Show file tree
Hide file tree
Changes from 4 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).
- Added Mask-RCNN Auto Annotation Script in OpenVINO format
- Added Yolo Auto Annotation Script
- Auto segmentation using Mask_RCNN component (Keras+Tensorflow Mask R-CNN Segmentation)
- Added MOT CSV format support

### 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 |

## Links
- [Intel AI blog: New Computer Vision Tool Accelerates Annotation of Digital Images and Video](https://www.intel.ai/introducing-cvat)
Expand Down
110 changes: 110 additions & 0 deletions cvat/apps/annotation/mot.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
# 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' or shape.outside:
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():
# MOT doesn't support the outside flag so we duplicate each
# track's last shape, increment frame_id and set outside=True
# to end the track
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 + 1,
attributes=track.shapes[-1].attributes,
)
track.shapes.append(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,4 +12,5 @@
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'),
)