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)