diff --git a/common/polystar/common/models/image.py b/common/polystar/common/models/image.py
index 1bf4fc61f065337992ef42e9f181f79729679b7c..33d63c22ab531203a2fa9b6439769640bb9729d7 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 5a10d97c38620fc4a6f8de2a1c1ce2cfd08c80b6..a13af165695e32f57c6c31c7a7e2c9f33ddf3881 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 0000000000000000000000000000000000000000..b3d47c380c48ee78c5a264568ed62ee8271f5ec7
--- /dev/null
+++ b/robots-at-runes/research/dataset/blend/examples/logo.xml
@@ -0,0 +1,12 @@
+<annotation>
+ <point>
+ <x>432</x>
+ <y>76</y>
+ <label>green</label>
+ </point>
+ <point>
+ <x>432</x>
+ <y>13</y>
+ <label>blue</label>
+ </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 969394a399862779da01609735917ad84d6f5fa0..c19202e3d159e3298ede9df9695cf388a5f20c7f 100644
--- a/robots-at-runes/research/dataset/blend/image_blender.py
+++ b/robots-at-runes/research/dataset/blend/image_blender.py
@@ -56,7 +56,7 @@ class ImageBlender:
@staticmethod
def _translate_poi(poi: PointOfInterest, x: int, y: int) -> PointOfInterest:
- return PointOfInterest(poi.x + x, poi.y + y)
+ return PointOfInterest(poi.x + x, poi.y + y, poi.label)
def _crop_background(self, background: Image) -> Image:
h, w, _ = background.shape
@@ -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,11 @@ 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.")
+ for poi in res.point_of_interests:
+ plt.plot([poi.x], [poi.y], f"{poi.label[0]}.")
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/blend/labeled_image_modifiers/labeled_image_rotator.py b/robots-at-runes/research/dataset/blend/labeled_image_modifiers/labeled_image_rotator.py
index b73c1eb06c6bbd84666ba74df4a1b1281af8610a..860972cbf0cc10519f982ef1a799c342bfc43361 100644
--- a/robots-at-runes/research/dataset/blend/labeled_image_modifiers/labeled_image_rotator.py
+++ b/robots-at-runes/research/dataset/blend/labeled_image_modifiers/labeled_image_rotator.py
@@ -25,7 +25,9 @@ class LabeledImageRotator(LabeledImageModifierABC):
prev_vector_to_center = np.array((poi.x - original_image.shape[1] / 2, poi.y - original_image.shape[0] / 2))
new_vector_to_center = np.dot(rotation_matrix, prev_vector_to_center)
return PointOfInterest(
- int(new_vector_to_center[0] + new_image.shape[1] / 2), int(new_vector_to_center[1] + new_image.shape[0] / 2)
+ int(new_vector_to_center[0] + new_image.shape[1] / 2),
+ int(new_vector_to_center[1] + new_image.shape[0] / 2),
+ poi.label,
)
def _get_value_from_factor(self, factor: float) -> float:
diff --git a/robots-at-runes/research/dataset/blend/labeled_image_modifiers/labeled_image_scaler.py b/robots-at-runes/research/dataset/blend/labeled_image_modifiers/labeled_image_scaler.py
index 73f9255d5f4dc537bc93e38291c3ea862328bd67..b98e25f4fa1a32e4197984864fe887dd7690b024 100644
--- a/robots-at-runes/research/dataset/blend/labeled_image_modifiers/labeled_image_scaler.py
+++ b/robots-at-runes/research/dataset/blend/labeled_image_modifiers/labeled_image_scaler.py
@@ -19,7 +19,7 @@ class LabeledImageScaler(LabeledImageModifierABC):
def _generate_modified_poi(
self, poi: PointOfInterest, original_image: Image, new_image: Image, scale: float
) -> PointOfInterest:
- return PointOfInterest(int(poi.x * scale), int(poi.y * scale))
+ return PointOfInterest(int(poi.x * scale), int(poi.y * scale), poi.label)
def _get_value_from_factor(self, factor: float) -> float:
intensity = (self.max_scale - 1) * abs(factor) + 1
diff --git a/robots-at-runes/research/dataset/labeled_image.py b/robots-at-runes/research/dataset/labeled_image.py
index fec391238b4fc7089a37575ca144777be49dc378..30d0dac6880f554786789cff8364ccac0d93a663 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
@@ -8,9 +13,49 @@ from polystar.common.models.image import Image
class PointOfInterest:
x: int
y: int
+ label: str
+
+ 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"]), label=d["label"])
+
+ @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"
+ ),
+ )