[common] (datasets) robots@robots: adapt to the new format

common/polystar/common/filters/pass_through_filter.py

+from typing import Any
+
+from polystar.common.filters.filter_abc import FilterABC
+
+
+class PassThroughFilter(FilterABC[Any]):
+    def validate_single(self, example: Any) -> bool:
+        return True

common/research/common/dataset/improvement/zoom.py

 from copy import copy
 from dataclasses import dataclass
+from itertools import islice
 from typing import Iterable, List, Tuple
 
 from polystar.common.models.box import Box
 from polystar.common.models.image import Image
 from polystar.common.target_pipeline.objects_validators.in_box_validator import InBoxValidator
 from polystar.common.view.plt_results_viewer import PltResultViewer
-from research.common.datasets.roco.roco_annotation import ROCOAnnotation
-from research.common.datasets.roco.zoo.roco_datasets_zoo import ROCODatasetsZoo
+from research.common.datasets_v3.roco.roco_annotation import ROCOAnnotation
+from research.common.datasets_v3.roco.zoo.roco_dataset_zoo import ROCODatasetsZoo
 
 
 def crop_image_annotation(
@@ -119,13 +120,10 @@ class Zoomer:
 
 
 if __name__ == "__main__":
-    zoomer = Zoomer(854, 480, 0.15, 0.5)
+    _zoomer = Zoomer(854, 480, 0.15, 0.5)
 
-    for k, (img, annot) in enumerate(ROCODatasetsZoo.DJI.NorthChina):
-        viewer = PltResultViewer(f"img {i}")
+    for _img, _annot, _name in islice(ROCODatasetsZoo.DJI.NORTH_CHINA.lazy(), 0, 3):
+        _viewer = PltResultViewer(f"img {_name}")
 
-        for (cropped_image, cropped_annotation) in zoomer.zoom(img, annot):
-            viewer.display_image_with_objects(cropped_image, cropped_annotation.objects)
-
-        if k == 2:
-            break
+        for (_cropped_image, _cropped_annotation, _cropped_name) in _zoomer.zoom(_img, _annot, _name):
+            _viewer.display_image_with_objects(_cropped_image, _cropped_annotation.objects)

common/research/common/dataset/tensorflow_record.py

@@ -4,24 +4,24 @@ from shutil import move
 from typing import List
 
 import tensorflow as tf
+from tensorflow_core.python.lib.io import python_io
+
 from polystar.common.models.label_map import label_map
+from polystar.common.utils.tqdm import smart_tqdm
 from research.common.constants import TENSORFLOW_RECORDS_DIR
-from research.common.datasets.roco.directory_roco_dataset import \
-    DirectoryROCODataset
-from research.common.datasets.roco.roco_annotation import ROCOAnnotation
-from tensorflow_core.python.lib.io import python_io
-from tqdm import tqdm
+from research.common.datasets_v3.roco.roco_annotation import ROCOAnnotation
+from research.common.datasets_v3.roco.roco_datasets import ROCODatasets
 
 
 class TensorflowRecordFactory:
     @staticmethod
-    def from_datasets(datasets: List[DirectoryROCODataset], prefix: str = ""):
+    def from_datasets(datasets: List[ROCODatasets], prefix: str = ""):
         record_name = prefix + "_".join(d.name for d in datasets)
         writer = python_io.TFRecordWriter(str(TENSORFLOW_RECORDS_DIR / f"{record_name}.record"))
         c = 0
-        for dataset in tqdm(datasets, desc=record_name, total=len(datasets), unit="dataset"):
-            for image_path, annotation, _ in tqdm(
-                dataset, desc=dataset.name, total=len(dataset), unit="img", leave=False
+        for dataset in smart_tqdm(datasets, desc=record_name, unit="dataset"):
+            for image_path, annotation, _ in smart_tqdm(
+                dataset.lazy_files(), desc=dataset.name, unit="img", leave=False
             ):
                 writer.write(_example_from_image_annotation(image_path, annotation).SerializeToString())
                 c += 1
@@ -32,7 +32,7 @@ class TensorflowRecordFactory:
         )
 
     @staticmethod
-    def from_dataset(dataset: DirectoryROCODataset, prefix: str = ""):
+    def from_dataset(dataset: ROCODatasets, prefix: str = ""):
         TensorflowRecordFactory.from_datasets([dataset], prefix)
 
 

common/research/common/datasets/dataset.py

-from abc import ABC, abstractmethod
-from collections import deque
-from typing import Callable, Generic, Iterable, Iterator, Tuple, TypeVar
-
-from more_itertools import ilen
-from polystar.common.utils.iterable_utils import smart_len
-from polystar.common.utils.misc import identity
-
-ExampleT = TypeVar("ExampleT")
-TargetT = TypeVar("TargetT")
-ExampleU = TypeVar("ExampleU")
-TargetU = TypeVar("TargetU")
-
-
-class Dataset(Generic[ExampleT, TargetT], Iterable[Tuple[ExampleT, TargetT, str]], ABC):
-    def __init__(self, name: str):
-        self.name = name
-
-    @property
-    @abstractmethod
-    def examples(self) -> Iterable[ExampleT]:
-        pass
-
-    @property
-    @abstractmethod
-    def targets(self) -> Iterable[TargetT]:
-        pass
-
-    @property
-    @abstractmethod
-    def names(self) -> Iterable[TargetT]:
-        pass
-
-    @abstractmethod
-    def __iter__(self) -> Iterator[Tuple[ExampleT, TargetT, str]]:
-        pass
-
-    @abstractmethod
-    def __len__(self):
-        pass
-
-    def transform_examples(self, example_transformer: Callable[[ExampleT], ExampleU]) -> "Dataset[ExampleU, TargetT]":
-        return self.transform(example_transformer, identity)
-
-    def transform_targets(
-        self, target_transformer: Callable[[TargetT], TargetU] = identity
-    ) -> "Dataset[ExampleT, TargetU]":
-        return self.transform(identity, target_transformer)
-
-    def transform(
-        self, example_transformer: Callable[[ExampleT], ExampleU], target_transformer: Callable[[TargetT], TargetU]
-    ) -> "Dataset[ExampleU, TargetU]":
-        return GeneratorDataset(
-            self.name,
-            lambda: (
-                (example_transformer(example), target_transformer(target), name) for example, target, name in self
-            ),
-        )
-
-    def __str__(self):
-        return f"<{self.__class__.__name__} {self.name}>"
-
-    __repr__ = __str__
-
-    def check_consistency(self):
-        assert smart_len(self.targets) == smart_len(self.examples) == smart_len(self.names)
-
-
-class LazyUnzipper:
-    def __init__(self, iterator: Iterator[Tuple], n: int):
-        self._iterator = iterator
-        self._memory = [deque() for _ in range(n)]
-
-    def empty(self, i: int):
-        return self._iterator is None and not self._memory[i]
-
-    def elements(self, i: int):
-        while True:
-            if self._memory[i]:
-                yield self._memory[i].popleft()
-            elif self._iterator is None:
-                return
-            else:
-                try:
-                    elements = next(self._iterator)
-                    for k in range(len(elements)):
-                        if k != i:
-                            self._memory[k].append(elements[k])
-                    yield elements[i]
-                except StopIteration:
-                    self._iterator = None
-                    return
-
-
-class LazyDataset(Dataset[ExampleT, TargetT], ABC):
-    def __init__(self, name: str):
-        super().__init__(name)
-        self._unzipper = None
-
-    @property
-    def examples(self) -> Iterable[ExampleT]:
-        return self._elements(0)
-
-    @property
-    def targets(self) -> Iterable[TargetT]:
-        return self._elements(1)
-
-    @property
-    def names(self) -> Iterable[str]:
-        return self._elements(2)
-
-    def __len__(self):
-        return ilen(self)
-
-    def _elements(self, i: int) -> Iterable:
-        if self._unzipper is None or self._unzipper.empty(i):
-            self._unzipper = LazyUnzipper(iter(self), 3)
-        return self._unzipper.elements(i)
-
-
-class GeneratorDataset(LazyDataset[ExampleT, TargetT]):
-    def __init__(self, name: str, generator: Callable[[], Iterator[Tuple[ExampleT, TargetT, str]]]):
-        self.generator = generator
-        super().__init__(name)
-
-    def __iter__(self) -> Iterator[Tuple[ExampleT, TargetT, str]]:
-        return self.generator()

common/research/common/datasets/filtered_dataset.py

-from polystar.common.filters.filter_abc import FilterABC
-from research.common.datasets.dataset import Dataset, ExampleT, TargetT
-from research.common.datasets.simple_dataset import SimpleDataset
-
-
-class FilteredTargetsDataset(SimpleDataset[ExampleT, TargetT]):
-    def __init__(self, dataset: Dataset[ExampleT, TargetT], targets_filter: FilterABC[TargetT]):
-        targets, examples, names = targets_filter.filter_with_siblings(
-            list(dataset.targets), list(dataset.examples), list(dataset.names)
-        )
-        super().__init__(examples, targets, names, dataset.name)
-
-
-class FilteredExamplesDataset(SimpleDataset[ExampleT, TargetT]):
-    def __init__(self, dataset: Dataset[ExampleT, TargetT], examples_filter: FilterABC[ExampleT]):
-        super().__init__(
-            *examples_filter.filter_with_siblings(list(dataset.examples), list(dataset.targets), list(dataset.names)),
-            dataset.name,
-        )

common/research/common/datasets/image_dataset.py

-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Iterator, List, Tuple
-
-from memoized_property import memoized_property
-from more_itertools import ilen
-from polystar.common.models.image import Image, load_image
-from research.common.datasets.dataset import Dataset, LazyDataset, TargetT
-
-ImageDataset = Dataset[Image, TargetT]
-
-
-class ImageFileDataset(LazyDataset[Path, TargetT], ABC):
-    def __iter__(self) -> Iterator[Tuple[Path, TargetT, str]]:
-        for image_file in self.image_files:
-            yield image_file, self.target_from_image_file(image_file), image_file.stem
-
-    @abstractmethod
-    def target_from_image_file(self, image_file: Path) -> TargetT:
-        pass
-
-    @property
-    @abstractmethod
-    def image_files(self) -> Iterator[Path]:
-        pass
-
-    def open(self) -> ImageDataset:
-        return open_file_dataset(self)
-
-    def __len__(self):
-        return ilen(self.image_files)
-
-
-def open_file_dataset(dataset: Dataset[Path, TargetT]) -> ImageDataset:
-    return dataset.transform_examples(load_image)
-
-
-class ImageDirectoryDataset(ImageFileDataset[TargetT], ABC):
-    def __init__(self, images_dir: Path, name: str, extension: str = "jpg"):
-        super().__init__(name)
-        self.extension = extension
-        self.images_dir = images_dir
-
-    @memoized_property
-    def image_files(self) -> List[Path]:
-        return list(sorted(self.images_dir.glob(f"*.{self.extension}")))
-
-    def __len__(self):
-        return len(self.image_files)

common/research/common/datasets/roco/directory_roco_dataset.py

-from pathlib import Path
-
-from polystar.common.models.image import Image, save_image
-from research.common.datasets.image_dataset import ImageDirectoryDataset
-from research.common.datasets.roco.roco_annotation import ROCOAnnotation
-
-
-class DirectoryROCODataset(ImageDirectoryDataset[ROCOAnnotation]):
-    def __init__(self, dataset_path: Path, name: str):
-        super().__init__(dataset_path / "image", name)
-        self.main_dir = dataset_path
-        self.annotations_dir: Path = self.main_dir / "image_annotation"
-
-    def target_from_image_file(self, image_file: Path) -> ROCOAnnotation:
-        return ROCOAnnotation.from_xml_file(self.annotations_dir / f"{image_file.stem}.xml")
-
-    def create(self):
-        self.main_dir.mkdir(parents=True)
-        self.images_dir.mkdir()
-        self.annotations_dir.mkdir()
-
-    def add(self, image: Image, annotation: ROCOAnnotation, name: str):
-        save_image(image, self.images_dir / f"{name}.jpg")
-        (self.annotations_dir / f"{name}.xml").write_text(annotation.to_xml())

common/research/common/datasets/roco/roco_dataset.py

-from polystar.common.models.image import Image
-from research.common.datasets.dataset import Dataset
-from research.common.datasets.image_dataset import ImageFileDataset
-from research.common.datasets.roco.roco_annotation import ROCOAnnotation
-
-ROCODataset = Dataset[Image, ROCOAnnotation]
-ROCOFileDataset = ImageFileDataset[ROCOAnnotation]

common/research/common/datasets/roco/roco_datasets.py

-from abc import abstractmethod
-from pathlib import Path
-from typing import Any, ClassVar, Iterable, Iterator, List, Tuple
-
-from research.common.datasets.roco.directory_roco_dataset import \
-    DirectoryROCODataset
-
-
-class ROCODatasets(Iterable[DirectoryROCODataset]):
-    name: ClassVar[str]
-    datasets: ClassVar[List[DirectoryROCODataset]]
-    directory: ClassVar[Path]
-
-    @classmethod
-    @abstractmethod
-    def make_dataset(cls, dataset_name: str, *args: Any) -> DirectoryROCODataset:
-        pass
-
-    def __init_subclass__(cls, **kwargs):
-        cls.datasets: List[DirectoryROCODataset] = []
-        for dataset_name, args in cls.__dict__.items():
-            if not dataset_name.islower():
-                if not isinstance(args, Tuple):
-                    args = (args,)
-                dataset = cls.make_dataset(dataset_name, *args)
-                setattr(cls, dataset_name, dataset)
-                cls.datasets.append(dataset)
-
-        cls.name = cls.__name__[: -len("ROCODatasets")]
-
-    def __iter__(self) -> Iterator[DirectoryROCODataset]:
-        return self.datasets.__iter__()

common/research/common/datasets/roco/zoo/dji.py

-from research.common.constants import DJI_ROCO_DSET_DIR
-from research.common.datasets.roco.directory_roco_dataset import \
-    DirectoryROCODataset
-from research.common.datasets.roco.roco_datasets import ROCODatasets
-
-
-class DJIROCODatasets(ROCODatasets):
-    directory = DJI_ROCO_DSET_DIR
-
-    CentralChina = "robomaster_Central China Regional Competition"
-    NorthChina = "robomaster_North China Regional Competition"
-    SouthChina = "robomaster_South China Regional Competition"
-    Final = "robomaster_Final Tournament"
-
-    @classmethod
-    def make_dataset(cls, dataset_name: str, competition_name: str) -> DirectoryROCODataset:
-        return DirectoryROCODataset(cls.directory / competition_name, dataset_name)

common/research/common/datasets/roco/zoo/dji_zoomed.py

-from typing import Any
-
-from polystar.common.utils.str_utils import camel2snake
-from research.common.constants import DJI_ROCO_ZOOMED_DSET_DIR
-from research.common.datasets.roco.directory_roco_dataset import \
-    DirectoryROCODataset
-from research.common.datasets.roco.roco_datasets import ROCODatasets
-
-
-class DJIROCOZoomedDatasets(ROCODatasets):
-    directory = DJI_ROCO_ZOOMED_DSET_DIR
-
-    CentralChina = ()
-    NorthChina = ()
-    SouthChina = ()
-    Final = ()
-
-    @classmethod
-    def make_dataset(cls, dataset_name: str, *args: Any) -> DirectoryROCODataset:
-        return DirectoryROCODataset(cls.directory / camel2snake(dataset_name), f"{dataset_name}ZoomedV2")

common/research/common/datasets/roco/zoo/roco_datasets_zoo.py

-from typing import Iterable
-
-from research.common.datasets.roco.roco_datasets import ROCODatasets
-from research.common.datasets.roco.zoo.dji import DJIROCODatasets
-from research.common.datasets.roco.zoo.dji_zoomed import DJIROCOZoomedDatasets
-from research.common.datasets.roco.zoo.twitch import TwitchROCODatasets
-
-
-class ROCODatasetsZoo(Iterable[ROCODatasets]):
-    DJI_ZOOMED = DJIROCOZoomedDatasets()
-    DJI = DJIROCODatasets()
-    TWITCH = TwitchROCODatasets()
-
-    def __iter__(self):
-        return (self.DJI, self.DJI_ZOOMED, self.TWITCH).__iter__()

common/research/common/datasets/roco/zoo/twitch.py

-from typing import Any
-
-from research.common.constants import TWITCH_DSET_DIR
-from research.common.datasets.roco.directory_roco_dataset import \
-    DirectoryROCODataset
-from research.common.datasets.roco.roco_datasets import ROCODatasets
-
-
-class TwitchROCODatasets(ROCODatasets):
-    directory = TWITCH_DSET_DIR / "v1"
-
-    T470149568 = ()
-    T470150052 = ()
-    T470151286 = ()
-    T470152289 = ()
-    T470152730 = ()
-    T470152838 = ()
-    T470153081 = ()
-    T470158483 = ()
-
-    @classmethod
-    def make_dataset(cls, dataset_name: str, *args: Any) -> DirectoryROCODataset:
-        twitch_id = dataset_name[len("T") :]
-        return DirectoryROCODataset(cls.directory / twitch_id, dataset_name)

common/research/common/datasets/simple_dataset.py

-from typing import Iterable, Iterator, List, Tuple
-
-from research.common.datasets.dataset import Dataset, ExampleT, TargetT
-
-
-class SimpleDataset(Dataset[ExampleT, TargetT]):
-    def __init__(self, examples: Iterable[ExampleT], targets: Iterable[TargetT], names: Iterable[str], name: str):
-        super().__init__(name)
-        self._examples = list(examples)
-        self._targets = list(targets)
-        self._names = list(names)
-        self.check_consistency()
-
-    @property
-    def examples(self) -> List[ExampleT]:
-        return self._examples
-
-    @property
-    def targets(self) -> List[TargetT]:
-        return self._targets
-
-    @property
-    def names(self) -> List[TargetT]:
-        return self._names
-
-    def __iter__(self) -> Iterator[Tuple[ExampleT, TargetT, str]]:
-        return zip(self.examples, self.targets, self.names)
-
-    def __len__(self):
-        return len(self.examples)

common/research/common/datasets/union_dataset.py

-from typing import Iterable, Iterator, Tuple
-
-from research.common.datasets.dataset import (Dataset, ExampleT, LazyDataset,
-                                              TargetT)
-
-
-class UnionDataset(LazyDataset[ExampleT, TargetT]):
-    def __init__(self, datasets: Iterable[Dataset[ExampleT, TargetT]], name: str = None):
-        self.datasets = list(datasets)
-        super().__init__(name or "_".join(d.name for d in self.datasets))
-
-    def __iter__(self) -> Iterator[Tuple[ExampleT, TargetT, str]]:
-        for dataset in self.datasets:
-            yield from dataset
-
-    def __len__(self):
-        return sum(map(len, self.datasets))

common/research/common/datasets_v3/capped_dataset.py

+from itertools import islice
+from typing import Iterator, Tuple
+
+from research.common.datasets_v3.lazy_dataset import ExampleT, LazyDataset, TargetT
+
+
+class CappedDataset(LazyDataset):
+    def __init__(self, source: LazyDataset[ExampleT, TargetT], n: int):
+        super().__init__(source.name)
+        self.n = n
+        self.source = source
+
+    def __iter__(self) -> Iterator[Tuple[ExampleT, TargetT]]:
+        return islice(self.source, self.n)

common/research/common/datasets_v3/dataset_builder.py

@@ -3,6 +3,7 @@ from typing import Callable, Generic, Iterable, Iterator, Tuple
 from polystar.common.filters.filter_abc import FilterABC
 from polystar.common.filters.pass_through_filter import PassThroughFilter
 from polystar.common.utils.misc import identity
+from research.common.datasets_v3.capped_dataset import CappedDataset
 from research.common.datasets_v3.dataset import Dataset
 from research.common.datasets_v3.filter_dataset import ExampleU, FilterDataset, TargetU
 from research.common.datasets_v3.lazy_dataset import ExampleT, LazyDataset, TargetT
@@ -54,6 +55,10 @@ class DatasetBuilder(Generic[ExampleT, TargetT], Iterable[Tuple[ExampleT, Target
         self.dataset = TransformDataset(self.dataset, identity, target_transformer)
         return self
 
+    def cap(self, n: int) -> "DatasetBuilder[ExampleT, TargetT]":
+        self.dataset = CappedDataset(self.dataset, n)
+        return self
+
     @property
     def name(self) -> str:
         return self.dataset.name