diff --git a/src/polystar/models/roco_image_annotation.py b/src/polystar/models/roco_image_annotation.py
index 32e4dec6326c895cd7f0da0bcb78284893ecb686..6c3b272b01eab459f8cffe7de312aab7ebe79e13 100644
--- a/src/polystar/models/roco_image_annotation.py
+++ b/src/polystar/models/roco_image_annotation.py
@@ -8,7 +8,7 @@ import xmltodict
from dicttoxml import dicttoxml
from polystar.models.image import Image, load_image, save_image
-from polystar.models.roco_object import ObjectFactory, ROCOObject
+from polystar.models.roco_object import ROCOObject, ROCOObjectFactory
@dataclass
@@ -40,7 +40,7 @@ class ROCOImageAnnotation:
json_objects = annotation.get("object", []) or []
json_objects = json_objects if isinstance(json_objects, list) else [json_objects]
roco_json_objects = [obj_json for obj_json in json_objects if not obj_json["name"].startswith("rune")]
- objects = [ObjectFactory.from_json(obj_json) for obj_json in roco_json_objects]
+ objects = [ROCOObjectFactory.from_json(obj_json) for obj_json in roco_json_objects]
return ROCOImageAnnotation(
width=int(annotation["size"]["width"]),
@@ -61,7 +61,7 @@ class ROCOImageAnnotation:
{
"annotation": {
"size": {"width": self.width, "height": self.height},
- "object": [ObjectFactory.to_json(obj) for obj in self.objects],
+ "object": [ROCOObjectFactory.to_json(obj) for obj in self.objects],
}
},
attr_type=False,
diff --git a/src/polystar/models/roco_object.py b/src/polystar/models/roco_object.py
index bd9a096343ecd0d13e42bf5d5695490249559967..844ebd5f4bd9ec5079d6f69bb5aa5e1e46e6a2bc 100644
--- a/src/polystar/models/roco_object.py
+++ b/src/polystar/models/roco_object.py
@@ -1,14 +1,12 @@
from dataclasses import dataclass
from enum import auto
-from typing import Any, Dict, NewType
+from typing import Any, Dict, NewType, Set
from polystar.models.box import Box
from polystar.utils.no_case_enum import NoCaseEnum
Json = NewType("Json", Dict[str, Any])
-ArmorNumber = NewType("ArmorNumber", int)
-
class ArmorColor(NoCaseEnum):
GREY = auto()
@@ -25,6 +23,9 @@ class ArmorColor(NoCaseEnum):
return self.name[0] if self != ArmorColor.UNKNOWN else "?"
+VALID_NUMBERS_2021: Set[int] = {1, 3, 4} # University League # CHANGING
+
+
class ArmorDigit(NoCaseEnum): # CHANGING
# Those have real numbers
HERO = auto()
@@ -55,6 +56,16 @@ class ArmorDigit(NoCaseEnum): # CHANGING
def digit(self) -> int:
return self.value + (self.value >= 2) # hacky, but a number is missing (2)
+ @staticmethod
+ def from_number(n: int) -> "ArmorDigit":
+ if n == 0:
+ return ArmorDigit.UNKNOWN
+
+ if n not in VALID_NUMBERS_2021:
+ return ArmorDigit.OUTDATED
+
+ return ArmorDigit(n - (n >= 3)) # hacky, but digit 2 is absent
+
class ObjectType(NoCaseEnum):
CAR = auto()
@@ -75,15 +86,15 @@ class ROCOObject:
@dataclass
class Armor(ROCOObject):
- number: ArmorNumber
- type: ArmorDigit
+ number: int
+ digit: ArmorDigit
color: ArmorColor
def __repr__(self):
return f"<{self} {self.color} {self.number}>"
-class ObjectFactory:
+class ROCOObjectFactory:
@staticmethod
def from_json(json: Json) -> ROCOObject:
t: ObjectType = ObjectType(json["name"])
@@ -100,10 +111,14 @@ class ObjectFactory:
if t is not ObjectType.ARMOR:
return ROCOObject(type=t, box=Box.from_size(x, y, w, h=h))
- armor_number = ArmorNumber(int(json["armor_class"])) if json["armor_class"] != "none" else 0
+ armor_number = int(json["armor_class"]) if json["armor_class"] != "none" else 0
return Armor(
- type=t, box=Box.from_size(x, y, w, h=h), number=armor_number, color=ArmorColor(json["armor_color"])
+ type=t,
+ box=Box.from_size(x, y, w, h=h),
+ number=armor_number,
+ digit=ArmorDigit.from_number(armor_number),
+ color=ArmorColor(json["armor_color"]),
)
@staticmethod
diff --git a/src/research/common/datasets/roco/roco_annotation.py b/src/research/common/datasets/roco/roco_annotation.py
index 89107bd564fc27a04992ab0893049c32fce1d63d..c1e47c6d749a2dc7d75177921fd0de8b4ffc5167 100644
--- a/src/research/common/datasets/roco/roco_annotation.py
+++ b/src/research/common/datasets/roco/roco_annotation.py
@@ -7,7 +7,7 @@ from xml.dom.minidom import parseString
import xmltodict
from dicttoxml import dicttoxml
-from polystar.models.roco_object import Armor, ObjectFactory, ROCOObject
+from polystar.models.roco_object import Armor, ROCOObject, ROCOObjectFactory
@dataclass
@@ -36,7 +36,7 @@ class ROCOAnnotation:
json_objects = xml_dict.get("object", []) or []
json_objects = json_objects if isinstance(json_objects, list) else [json_objects]
roco_json_objects = [obj_json for obj_json in json_objects if not obj_json["name"].startswith("rune")]
- objects = [ObjectFactory.from_json(obj_json) for obj_json in roco_json_objects]
+ objects = [ROCOObjectFactory.from_json(obj_json) for obj_json in roco_json_objects]
return ROCOAnnotation(
objects=objects,
@@ -51,7 +51,7 @@ class ROCOAnnotation:
{
"annotation": {
"size": {"width": self.w, "height": self.h},
- "object": [ObjectFactory.to_json(obj) for obj in self.objects],
+ "object": [ROCOObjectFactory.to_json(obj) for obj in self.objects],
}
},
attr_type=False,
diff --git a/src/research/common/datasets/roco/roco_dataset_descriptor.py b/src/research/common/datasets/roco/roco_dataset_descriptor.py
index 666522f6ffd8ebec944cf9ca8f147a7bdc07efa7..2ea388e699bccb2fee2d2f63b9fc9ffba7ae3fcd 100644
--- a/src/research/common/datasets/roco/roco_dataset_descriptor.py
+++ b/src/research/common/datasets/roco/roco_dataset_descriptor.py
@@ -64,6 +64,7 @@ def make_markdown_dataset_report(dataset: LazyROCOFileDataset, report_dir: Path)
]
)
mf.table(DataFrame(stats.armors_color2num2count))
+ print(f"file://{report_path}")
if __name__ == "__main__":
diff --git a/src/research/robots/armor_digit/armor_digit_dataset.py b/src/research/robots/armor_digit/armor_digit_dataset.py
index 45c1abfba586acf38fd6589bb8ac2635f76a9349..ad179678c69e06174ecc68cb1c50c4f0a20ecdec 100644
--- a/src/research/robots/armor_digit/armor_digit_dataset.py
+++ b/src/research/robots/armor_digit/armor_digit_dataset.py
@@ -1,5 +1,5 @@
from itertools import islice
-from typing import List, Set, Tuple
+from typing import List, Tuple
from polystar.filters.exclude_filter import ExcludeFilter
from polystar.models.image import FileImage
@@ -9,8 +9,6 @@ from research.common.datasets.roco.zoo.roco_dataset_zoo import ROCODatasetsZoo
from research.robots.dataset.armor_value_dataset_generator import ArmorValueDatasetGenerator
from research.robots.dataset.armor_value_target_factory import ArmorValueTargetFactory
-VALID_NUMBERS_2021: Set[int] = {1, 3, 4} # University League
-
def default_armor_digit_datasets() -> Tuple[
List[Dataset[FileImage, ArmorDigit]], List[Dataset[FileImage, ArmorDigit]], List[Dataset[FileImage, ArmorDigit]]
@@ -24,15 +22,10 @@ def make_armor_digit_dataset_generator() -> ArmorValueDatasetGenerator[ArmorDigi
class ArmorDigitTargetFactory(ArmorValueTargetFactory[ArmorDigit]):
def from_str(self, label: str) -> ArmorDigit:
- n = int(label)
-
- if n in VALID_NUMBERS_2021: # CHANGING
- return ArmorDigit(n - (n >= 3)) # hacky, but digit 2 is absent
-
- return ArmorDigit.OUTDATED
+ return ArmorDigit.from_number(int(label))
def from_armor(self, armor: Armor) -> ArmorDigit:
- return ArmorDigit(armor.number) if armor.number else ArmorDigit.UNKNOWN
+ return ArmorDigit(armor.digit)
if __name__ == "__main__":