def test_open(self):
with TemporaryDirectory() as dataset_dir:
dataset = DirectoryROCODataset(Path(dataset_dir), "fake")
annotation = ROCOAnnotation("frame_1",
objects=[],
has_rune=False,
w=160,
h=90)
image = asarray([[[250, 0, 0], [250, 0, 0]],
[[250, 0, 0], [250, 0, 0]]]).astype(float32)
dataset.annotations_dir.mkdir()
dataset.images_dir.mkdir()
(dataset.annotations_dir / "frame_1.xml").write_text(
annotation.to_xml())
save_image(image, dataset.images_dir / "frame_1.jpg")
image_dataset = dataset.open()
self.assertEqual([annotation], list(image_dataset.targets))
images = list(image_dataset.examples)
self.assertEqual(1, len(images))
assert_array_almost_equal(image / 256, images[0] / 256,
decimal=2) # jpeg precision
def test_file(self):
with TemporaryDirectory() as dataset_dir:
dataset_dir = Path(dataset_dir)
annotation = ROCOAnnotation(objects=[],
has_rune=False,
w=160,
h=90)
images_dir, annotations_dir = self._setup_dir(dataset_dir)
(annotations_dir / "frame_1.xml").write_text(annotation.to_xml())
(images_dir / "frame_1.jpg").write_text("")
dataset = ROCODatasetBuilder(dataset_dir, "fake").build_lazy()
self.assertEqual(
[(images_dir / "frame_1.jpg", annotation, "frame_1")],
list(dataset))
def clear_small_bases(annotation: ROCOAnnotation):
small_bases, annotation.objects = SMALL_BASE_FILTER.split(annotation.objects)
if not small_bases:
return
armors, robots = ARMORS_FILTER.split(annotation.objects)
for base in small_bases:
armors = (-InBoxObjectFilter(base.box, 0.5)).filter(armors)
annotation.objects = robots + armors
def test_targets(self):
with TemporaryDirectory() as dataset_dir:
dataset = DirectoryROCODataset(Path(dataset_dir), "fake")
annotation = ROCOAnnotation("frame_1",
objects=[],
has_rune=False,
w=160,
h=90)
dataset.annotations_dir.mkdir()
dataset.images_dir.mkdir()
(dataset.annotations_dir / "frame_1.xml").write_text(
annotation.to_xml())
(dataset.images_dir / "frame_1.jpg").write_text("")
self.assertEqual([annotation], list(dataset.targets))
self.assertEqual([dataset.images_dir / "frame_1.jpg"],
list(dataset.examples))
def _generate_from_single(
self, image: Image, annotation: ROCOAnnotation,
name: str) -> Iterator[Tuple[Image, ROCOAnnotation, str]]:
annotation.objects, robots = ARMORS_FILTER.split(annotation.objects)
for i, robot in enumerate(self.robots_filter.filter(robots)):
yield crop_image_annotation(
image,
annotation,
robot.box,
min_coverage=0.75,
name=f"{name}-{i}-{robot.type.name.lower()}")
def add_annotated_images_to_dataset(images_dir: Path, annotations_dir: Path, dataset_dir: Path):
corrector = AnnotationFileCorrector(save_before=False)
for annotation_file in tqdm(list(annotations_dir.glob("**/*.xml"))):
sub_dataset_dir = dataset_dir / annotation_file.stem.split("-")[0]
image_file = (images_dir / annotation_file.relative_to(annotations_dir)).with_suffix(".jpg")
copy_file(image_file, sub_dataset_dir / "image")
annotation_file = copy_file(annotation_file, sub_dataset_dir / "image_annotation")
corrector.correct_annotation_file(annotation_file)
if ROCOAnnotation.from_xml_file(annotation_file).has_rune:
move_image_and_annotation_from_directory(sub_dataset_dir, dataset_dir / "runes", annotation_file.stem)
def _scale_annotation(annotation: ROCOAnnotation, height: int, width: int):
vertical_ratio, horizontal_ratio = height / annotation.h, width / annotation.w
for obj in annotation.objects:
obj.box = Box.from_positions(
x1=int(obj.box.x1 * horizontal_ratio),
y1=int(obj.box.y1 * vertical_ratio),
x2=int(obj.box.x2 * horizontal_ratio),
y2=int(obj.box.y2 * vertical_ratio),
)
annotation.w, annotation.h = width, height
def test_image(self):
with TemporaryDirectory() as dataset_dir:
dataset_dir = Path(dataset_dir)
annotation = ROCOAnnotation(objects=[],
has_rune=False,
w=160,
h=90)
image = asarray([[[250, 0, 0], [250, 0, 0]],
[[250, 0, 0], [250, 0, 0]]]).astype(float32)
images_dir, annotations_dir = self._setup_dir(dataset_dir)
(annotations_dir / "frame_1.xml").write_text(annotation.to_xml())
save_image(image, images_dir / "frame_1.jpg")
dataset = ROCODatasetBuilder(dataset_dir,
"fake").to_images().build()
self.assertEqual([annotation], list(dataset.targets))
images = list(dataset.examples)
self.assertEqual(1, len(images))
assert_array_almost_equal(image / 256, images[0] / 256,
decimal=2) # jpeg precision
def crop_image_annotation(
image: Image, annotation: ROCOAnnotation, box: Box, min_coverage: float, name: str
) -> Tuple[Image, ROCOAnnotation, str]:
objects = InBoxObjectFilter(box, min_coverage).filter(annotation.objects)
objects = [copy(o) for o in objects]
for obj in objects:
obj.box = Box.from_positions(
x1=max(0, obj.box.x1 - box.x1),
y1=max(0, obj.box.y1 - box.y1),
x2=min(box.x2, obj.box.x2 - box.x1),
y2=min(box.y2, obj.box.y2 - box.y1),
)
return (
image[box.y1 : box.y2, box.x1 : box.x2],
ROCOAnnotation(w=box.w, h=box.h, objects=objects, has_rune=False),
name,
)
def _roco_annotation_from_image_file(self,
image_file: Path) -> ROCOAnnotation:
return ROCOAnnotation.from_xml_file(self.annotations_dir /
f"{image_file.stem}.xml")
def _save_one(self, img: Image, annotation: ROCOAnnotation, name: str):
annotation.save_with_image(self.cache_dir, img, name)