Skip to content

Commit

Permalink
Adding dump and load support for MOT CSV format. (#830)
Browse files Browse the repository at this point in the history
* Adding dump and load support for MOT CSV format.
* Updated test cases to use correct track annotations for MOT format.
* Removed behaviour of MOT loader which would duplicate the last track
shape prior to setting outside=True.
  • Loading branch information
telenachos authored and nmanovic committed Nov 22, 2019
1 parent 112bab1 commit 3aa4abf
Show file tree
Hide file tree
Showing 5 changed files with 115 additions and 0 deletions.
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

0 comments on commit 3aa4abf

Please sign in to comment.