From 05b7f0b0642274a60bc32221e64726e9019a144b Mon Sep 17 00:00:00 2001
From: Mathieu Beligon <mathieu@feedly.com>
Date: Mon, 30 Mar 2020 00:53:24 -0400
Subject: [PATCH] [runes] (blend) Add way to save and load annotations
---
common/polystar/common/models/image.py | 4 ++
.../dataset/blend/examples/.gitignore | 4 +-
.../research/dataset/blend/examples/logo.xml | 10 ++++
.../research/dataset/blend/image_blender.py | 5 +-
.../research/dataset/labeled_image.py | 46 ++++++++++++++++++-
5 files changed, 64 insertions(+), 5 deletions(-)
create mode 100644 robots-at-runes/research/dataset/blend/examples/logo.xml
diff --git a/common/polystar/common/models/image.py b/common/polystar/common/models/image.py
index 1bf4fc6..33d63c2 100644
--- a/common/polystar/common/models/image.py
+++ b/common/polystar/common/models/image.py
@@ -8,3 +8,7 @@ class Image(Array[int, ..., ..., 3]):
@staticmethod
def from_path(image_path: Path) -> "Image":
return cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB)
+
+ def save(self, image_path: Path):
+ image_path.parent.mkdir(exist_ok=True, parents=True)
+ cv2.imwrite(str(image_path), cv2.cvtColor(self, cv2.COLOR_RGB2BGR))
diff --git a/robots-at-runes/research/dataset/blend/examples/.gitignore b/robots-at-runes/research/dataset/blend/examples/.gitignore
index 5a10d97..a13af16 100644
--- a/robots-at-runes/research/dataset/blend/examples/.gitignore
+++ b/robots-at-runes/research/dataset/blend/examples/.gitignore
@@ -1,2 +1,2 @@
-/image_*.jpg
-/labels.xml
\ No newline at end of file
+/image
+/image_annotation
\ No newline at end of file
diff --git a/robots-at-runes/research/dataset/blend/examples/logo.xml b/robots-at-runes/research/dataset/blend/examples/logo.xml
new file mode 100644
index 0000000..3257c42
--- /dev/null
+++ b/robots-at-runes/research/dataset/blend/examples/logo.xml
@@ -0,0 +1,10 @@
+<annotation>
+ <point>
+ <x>432</x>
+ <y>76</y>
+ </point>
+ <point>
+ <x>432</x>
+ <y>13</y>
+ </point>
+</annotation>
diff --git a/robots-at-runes/research/dataset/blend/image_blender.py b/robots-at-runes/research/dataset/blend/image_blender.py
index 969394a..8e30abd 100644
--- a/robots-at-runes/research/dataset/blend/image_blender.py
+++ b/robots-at-runes/research/dataset/blend/image_blender.py
@@ -76,7 +76,7 @@ if __name__ == "__main__":
_obj = LabeledImage(
cv2.cvtColor(cv2.imread(str(EXAMPLES_DIR / "logo.png"), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA),
- [PointOfInterest(432, 76), PointOfInterest(432, 13)],
+ PointOfInterest.from_annotation_file(EXAMPLES_DIR / "logo.xml"),
)
_bg = cv2.cvtColor(cv2.imread(str(EXAMPLES_DIR / "back1.jpg")), cv2.COLOR_BGR2RGB)
@@ -86,9 +86,10 @@ if __name__ == "__main__":
for i in range(10):
res = _blender.blend(_bg, _obj)
+ res.save(EXAMPLES_DIR, f"test_{i}")
+
plt.imshow(res.image)
plt.plot([poi.x for poi in res.point_of_interests], [poi.y for poi in res.point_of_interests], "r.")
plt.axis("off")
plt.tight_layout()
- plt.savefig(str(EXAMPLES_DIR / f"image_{i}.jpg"))
plt.show()
diff --git a/robots-at-runes/research/dataset/labeled_image.py b/robots-at-runes/research/dataset/labeled_image.py
index fec3912..2d66cc3 100644
--- a/robots-at-runes/research/dataset/labeled_image.py
+++ b/robots-at-runes/research/dataset/labeled_image.py
@@ -1,5 +1,10 @@
from dataclasses import dataclass, field
-from typing import List
+from pathlib import Path
+from typing import List, Dict, Any
+from xml.dom.minidom import parseString
+
+import xmltodict
+from dicttoxml import dicttoxml
from polystar.common.models.image import Image
@@ -9,8 +14,47 @@ class PointOfInterest:
x: int
y: int
+ def to_dict(self) -> Dict[str, int]:
+ return self.__dict__
+
+ @classmethod
+ def from_dict(cls, d: Dict[str, Any]):
+ return cls(x=int(d["x"]), y=int(d["y"]))
+
+ @classmethod
+ def from_annotation_file(cls, annotation_path: Path) -> List["PointOfInterest"]:
+ points = xmltodict.parse(annotation_path.read_text())["annotation"]["point"]
+ return [PointOfInterest.from_dict(p) for p in points]
+
@dataclass
class LabeledImage:
image: Image
point_of_interests: List[PointOfInterest] = field(default_factory=list)
+
+ def save(self, directory_path: Path, name: str):
+ Image.save(self.image, directory_path / "image" / f"{name}.jpg")
+ self._save_annotation(directory_path / "image_annotation" / f"{name}.xml")
+
+ def _save_annotation(self, annotation_path: Path):
+ annotation_path.parent.mkdir(exist_ok=True, parents=True)
+ xml = parseString(
+ dicttoxml(
+ {"annotation": {"point": [p.to_dict() for p in self.point_of_interests]}},
+ attr_type=False,
+ root="annotation",
+ item_func=lambda x: x,
+ )
+ .replace(b"<point><point>", b"<point>")
+ .replace(b"</point></point>", b"</point>")
+ ).toprettyxml()
+ annotation_path.write_text(xml)
+
+ @staticmethod
+ def from_directory(directory_path: Path, name: str) -> "LabeledImage":
+ return LabeledImage(
+ image=Image.from_path(directory_path / "image" / f"{name}.jpg"),
+ point_of_interests=PointOfInterest.from_annotation_file(
+ directory_path / "image_annotation" / f"{name}.xml"
+ ),
+ )
--
GitLab