From 7a3f6e6286fc9abba49be25544aec3e3200e293a Mon Sep 17 00:00:00 2001
From: Mathieu Beligon <mathieu@feedly.com>
Date: Tue, 31 Mar 2020 20:46:21 -0400
Subject: [PATCH] [robots] (armor dataset factory) Add lock to be sure that all
images were extracted last time
---
.../dataset/armor_image_dataset_factory.py | 16 ++++++++++++----
1 file changed, 12 insertions(+), 4 deletions(-)
diff --git a/robots-at-robots/research/dataset/armor_image_dataset_factory.py b/robots-at-robots/research/dataset/armor_image_dataset_factory.py
index 85dd6cc..7bff809 100644
--- a/robots-at-robots/research/dataset/armor_image_dataset_factory.py
+++ b/robots-at-robots/research/dataset/armor_image_dataset_factory.py
@@ -1,5 +1,5 @@
+import json
from abc import abstractmethod
-from dataclasses import dataclass
from pathlib import Path
from typing import TypeVar, Tuple, Iterable
@@ -7,6 +7,7 @@ import cv2
from polystar.common.models.image import Image
from polystar.common.models.object import ArmorColor
+from polystar.common.utils.time import create_time_id
from research.dataset.armor_dataset_factory import ArmorDatasetFactory
from research_common.dataset.directory_roco_dataset import DirectoryROCODataset
from research_common.image_pipeline_evaluation.image_dataset_generator import ImageDatasetGenerator
@@ -18,23 +19,30 @@ class ArmorImageDatasetGenerator(ImageDatasetGenerator[T]):
task_name: str
def from_roco_dataset(self, dataset: DirectoryROCODataset) -> Iterable[Tuple[Image, T]]:
- if not (dataset.dataset_path / self.task_name).exists():
+ if not (dataset.dataset_path / self.task_name / ".lock").exists():
self._create_labelized_armor_images_from_roco(dataset)
return self._get_saved_images_and_labels(dataset)
def _create_labelized_armor_images_from_roco(self, dataset):
dset_path = dataset.dataset_path / self.task_name
- dset_path.mkdir()
+ dset_path.mkdir(exist_ok=True)
for (armor_img, color, digit, k, path) in ArmorDatasetFactory.from_dataset(dataset):
label = self._label_from_armor_info(color, digit, k, path)
cv2.imwrite(str(dset_path / f"{path.stem}-{k}-{label}.jpg"), cv2.cvtColor(armor_img, cv2.COLOR_RGB2BGR))
+ (dataset.dataset_path / self.task_name / ".lock").write_text(
+ json.dumps({"version": "0.0", "date": create_time_id()})
+ )
def _get_saved_images_and_labels(self, dataset: DirectoryROCODataset) -> Iterable[Tuple[Image, T]]:
return (
- (Image.from_path(image_path), self._label_from_str(image_path.stem.split("-")[-1]))
+ (Image.from_path(image_path), self._label_from_filepath(image_path))
for image_path in (dataset.dataset_path / self.task_name).glob("*.jpg")
+ if self._valid_label(self._label_from_filepath(image_path))
)
+ def _label_from_filepath(self, image_path: Path) -> T:
+ return self._label_from_str(image_path.stem.split("-")[-1])
+
@abstractmethod
def _label_from_armor_info(self, color: ArmorColor, digit: int, k: int, path: Path) -> T:
pass
--
GitLab