Skip to content

Commit

Permalink
[HOTFIX] Include bounding box annotation data when splitting videos (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
andreaazzini authored Oct 21, 2021
1 parent 560eccc commit 30ee564
Show file tree
Hide file tree
Showing 6 changed files with 120 additions and 58 deletions.
34 changes: 30 additions & 4 deletions darwin/datatypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,12 +101,27 @@ def make_tag(class_name: str, subs: Optional[List[SubAnnotation]] = None):
return Annotation(AnnotationClass(class_name, "tag"), {}, subs or [])


def make_polygon(class_name: str, point_path: List[Point], subs: Optional[List[SubAnnotation]] = None):
return Annotation(AnnotationClass(class_name, "polygon"), {"path": point_path}, subs or [])
def make_polygon(
class_name: str, point_path: List[Point], bounding_box: Optional[Dict], subs: Optional[List[SubAnnotation]] = None
):
return Annotation(
AnnotationClass(class_name, "polygon"),
_maybe_add_bounding_box_data({"path": point_path}, bounding_box),
subs or [],
)


def make_complex_polygon(class_name: str, point_paths: List[List[Point]], subs: Optional[List[SubAnnotation]] = None):
return Annotation(AnnotationClass(class_name, "complex_polygon", "polygon"), {"paths": point_paths}, subs or [])
def make_complex_polygon(
class_name: str,
point_paths: List[List[Point]],
bounding_box: Optional[Dict],
subs: Optional[List[SubAnnotation]] = None,
):
return Annotation(
AnnotationClass(class_name, "complex_polygon", "polygon"),
_maybe_add_bounding_box_data({"paths": point_paths}, bounding_box),
subs or [],
)


def make_keypoint(class_name: str, x: float, y: float, subs: Optional[List[SubAnnotation]] = None):
Expand Down Expand Up @@ -169,3 +184,14 @@ def make_video_annotation(frames, keyframes, segments, interpolated):
raise ValueError("invalid argument to make_video_annotation")

return VideoAnnotation(first_annotation.annotation_class, frames, keyframes, segments, interpolated)


def _maybe_add_bounding_box_data(data: Dict[str, Any], bounding_box: Optional[Dict]) -> Dict[str, Any]:
if bounding_box:
data["bounding_box"] = {
"x": bounding_box["x"],
"y": bounding_box["y"],
"w": bounding_box["w"],
"h": bounding_box["h"],
}
return data
1 change: 1 addition & 0 deletions darwin/exporter/formats/coco.py
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,7 @@ def build_annotation(annotation_file, annotation_id, annotation: dt.Annotation,
dt.make_polygon(
annotation.annotation_class.name,
[{"x": x, "y": y}, {"x": x + w, "y": y}, {"x": x + w, "y": y + h}, {"x": x, "y": y + h}],
None,
annotation.subs,
),
categories,
Expand Down
10 changes: 10 additions & 0 deletions darwin/exporter/formats/darwin.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,13 @@ def build_image_annotation(annotation_file: dt.AnnotationFile):
annotation.annotation_class.annotation_type: build_annotation_data(annotation),
"name": annotation.annotation_class.name,
}

if (
annotation.annotation_class.annotation_type == "complex_polygon"
or annotation.annotation_class.annotation_type == "polygon"
) and "bounding_box" in annotation.data:
payload["bounding_box"] = annotation.data["bounding_box"]

annotations.append(payload)

return {
Expand All @@ -25,4 +32,7 @@ def build_annotation_data(annotation: dt.Annotation):
if annotation.annotation_class.annotation_type == "complex_polygon":
return {"path": annotation.data["paths"]}

if annotation.annotation_class.annotation_type == "polygon":
return dict(filter(lambda item: item[0] != "bounding_box", annotation.data.items()))

return dict(annotation.data)
12 changes: 7 additions & 5 deletions darwin/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -286,17 +286,19 @@ def parse_darwin_annotation(annotation: Dict[str, Any]):
name = annotation["name"]
main_annotation = None
if "polygon" in annotation:
bounding_box = annotation.get("bounding_box")
if "additional_paths" in annotation["polygon"]:
paths = [annotation["polygon"]["path"]] + annotation["polygon"]["additional_paths"]
main_annotation = dt.make_complex_polygon(name, paths)
main_annotation = dt.make_complex_polygon(name, paths, bounding_box)
else:
main_annotation = dt.make_polygon(name, annotation["polygon"]["path"])
main_annotation = dt.make_polygon(name, annotation["polygon"]["path"], bounding_box)
elif "complex_polygon" in annotation:
bounding_box = annotation.get("bounding_box")
if "additional_paths" in annotation["complex_polygon"]:
paths = annotation["complex_polygon"]["path"] + annotation["complex_polygon"]["additional_paths"]
main_annotation = dt.make_complex_polygon(name, paths)
main_annotation = dt.make_complex_polygon(name, paths, bounding_box)
else:
main_annotation = dt.make_complex_polygon(name, annotation["complex_polygon"]["path"])
main_annotation = dt.make_complex_polygon(name, annotation["complex_polygon"]["path"], bounding_box)
elif "bounding_box" in annotation:
bounding_box = annotation["bounding_box"]
main_annotation = dt.make_bounding_box(
Expand Down Expand Up @@ -350,7 +352,7 @@ def split_video_annotation(annotation):
for i, frame_url in enumerate(annotation.frame_urls):
annotations = [a.frames[i] for a in annotation.annotations if i in a.frames]
annotation_classes = set([annotation.annotation_class for annotation in annotations])
filename = f"{Path(annotation.filename).stem}/{i:07d}.jpg"
filename = f"{Path(annotation.filename).stem}/{i:07d}.png"

frame_annotations.append(
dt.AnnotationFile(
Expand Down
6 changes: 3 additions & 3 deletions tests/darwin/dataset/remote_dataset_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,21 +333,21 @@ def it_works_on_videos(
"annotations": [
{"name": "test_class", "polygon": {"path": [{"x": 0, "y": 0}, {"x": 1, "y": 1}, {"x": 1, "y": 0}]}}
],
"image": {"filename": "test_video/0000000.jpg", "height": 1080, "url": "frame_1.jpg", "width": 1920},
"image": {"filename": "test_video/0000000.png", "height": 1080, "url": "frame_1.jpg", "width": 1920},
}

with (video_path / "0000001.json").open() as f:
assert json.load(f) == {
"annotations": [],
"image": {"filename": "test_video/0000001.jpg", "height": 1080, "url": "frame_2.jpg", "width": 1920},
"image": {"filename": "test_video/0000001.png", "height": 1080, "url": "frame_2.jpg", "width": 1920},
}

with (video_path / "0000002.json").open() as f:
assert json.load(f) == {
"annotations": [
{"name": "test_class", "polygon": {"path": [{"x": 5, "y": 5}, {"x": 6, "y": 6}, {"x": 6, "y": 5}]}}
],
"image": {"filename": "test_video/0000002.jpg", "height": 1080, "url": "frame_3.jpg", "width": 1920},
"image": {"filename": "test_video/0000002.png", "height": 1080, "url": "frame_3.jpg", "width": 1920},
}


Expand Down
115 changes: 69 additions & 46 deletions tests/darwin/utils_test.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from unittest.mock import MagicMock, patch

from darwin.datatypes import AnnotationFile
import darwin.datatypes as dt
from darwin.utils import (
is_extension_allowed,
is_image_extension_allowed,
Expand Down Expand Up @@ -114,7 +114,7 @@ def it_parses_darwin_images_correctly(tmp_path):
import_file = directory / "darwin-file.json"
import_file.write_text(content)

annotation_file: AnnotationFile = parse_darwin_json(import_file, None)
annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None)

assert annotation_file.path == import_file
assert annotation_file.filename == "P49-RediPad-ProPlayLEFTY_442.jpg"
Expand All @@ -131,7 +131,7 @@ def it_parses_darwin_images_correctly(tmp_path):

def it_parses_darwin_videos_correctly(tmp_path):
content = """
{
{
"dataset": "my-dataset",
"image": {
"width": 3840,
Expand All @@ -149,56 +149,56 @@ def it_parses_darwin_videos_correctly(tmp_path):
},
"annotations": [
{
"frames": {
"3": {
"bounding_box": {
"h": 547.0,
"w": 400.0,
"x": 363.0,
"y": 701.0
},
"instance_id": {
"value": 119
},
"keyframe": true,
"polygon": {
"path": [
{
"x": 748.0,
"y": 732.0
},
{
"x": 751.0,
"y": 735.0
},
{
"x": 748.0,
"y": 733.0
"frames": {
"3": {
"bounding_box": {
"h": 547.0,
"w": 400.0,
"x": 363.0,
"y": 701.0
},
"instance_id": {
"value": 119
},
"keyframe": true,
"polygon": {
"path": [
{
"x": 748.0,
"y": 732.0
},
{
"x": 751.0,
"y": 735.0
},
{
"x": 748.0,
"y": 733.0
}
]
}
}
},
"interpolate_algorithm": "linear-1.1",
"interpolated": true,
"name": "Hand",
"segments": [
[
3,
46
]
}
}
},
"interpolate_algorithm": "linear-1.1",
"interpolated": true,
"name": "Hand",
"segments": [
[
3,
46
]
]
}
]
}
}
"""

directory = tmp_path / "imports"
directory.mkdir()
import_file = directory / "darwin-file.json"
import_file.write_text(content)

annotation_file: AnnotationFile = parse_darwin_json(import_file, None)
annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None)

assert annotation_file.path == import_file
assert annotation_file.filename == "above tractor.mp4"
Expand All @@ -213,6 +213,29 @@ def it_parses_darwin_videos_correctly(tmp_path):
assert annotation_file.frame_urls == ["https://my-website.com/api/videos/209/frames/0"]
assert annotation_file.remote_path == "/"

assert annotation_file.annotations == [
dt.VideoAnnotation(
annotation_class=dt.AnnotationClass(
name="Hand", annotation_type="polygon", annotation_internal_type=None
),
frames={
3: dt.Annotation(
annotation_class=dt.AnnotationClass(
name="Hand", annotation_type="polygon", annotation_internal_type=None
),
data={
"path": [{"x": 748.0, "y": 732.0}, {"x": 751.0, "y": 735.0}, {"x": 748.0, "y": 733.0}],
"bounding_box": {"x": 363.0, "y": 701.0, "w": 400.0, "h": 547.0},
},
subs=[dt.SubAnnotation(annotation_type="instance_id", data=119)],
)
},
keyframes={3: True},
segments=[[3, 46]],
interpolated=True,
)
]

def it_returns_None_if_no_annotations_exist(tmp_path):
content = """
{
Expand All @@ -232,7 +255,7 @@ def it_returns_None_if_no_annotations_exist(tmp_path):
import_file = directory / "darwin-file.json"
import_file.write_text(content)

annotation_file: AnnotationFile = parse_darwin_json(import_file, None)
annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None)

assert not annotation_file

Expand Down Expand Up @@ -260,11 +283,11 @@ def it_uses_a_default_path_if_one_is_missing(tmp_path):
import_file = directory / "darwin-file.json"
import_file.write_text(content)

annotation_file: AnnotationFile = parse_darwin_json(import_file, None)
annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None)

assert annotation_file.remote_path == "/"

def it_imports_a_skeleteton(tmp_path):
def it_imports_a_skeleton(tmp_path):
content = """
{
"dataset": "cars",
Expand Down Expand Up @@ -325,7 +348,7 @@ def it_imports_a_skeleteton(tmp_path):
import_file = directory / "darwin-file.json"
import_file.write_text(content)

annotation_file: AnnotationFile = parse_darwin_json(import_file, None)
annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None)

assert annotation_file.annotations[0].annotation_class.annotation_type == "polygon"
assert annotation_file.annotations[1].annotation_class.annotation_type == "skeleton"
Expand Down Expand Up @@ -422,7 +445,7 @@ def it_imports_multiple_skeletetons(tmp_path):
import_file = directory / "darwin-file.json"
import_file.write_text(content)

annotation_file: AnnotationFile = parse_darwin_json(import_file, None)
annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None)

assert annotation_file.annotations[0].annotation_class.annotation_type == "polygon"
assert annotation_file.annotations[1].annotation_class.annotation_type == "skeleton"
Expand Down

0 comments on commit 30ee564

Please sign in to comment.