diff --git a/dataset/dji_roco/robomaster_Central China Regional Competition/digits/.changes b/dataset/dji_roco/robomaster_Central China Regional Competition/digits/.changes
new file mode 100644
index 0000000000000000000000000000000000000000..9e26dfeeb6e641a33dae4961196235bdb965b21b
Binary files /dev/null and b/dataset/dji_roco/robomaster_Central China Regional Competition/digits/.changes differ
diff --git a/dataset/dji_roco/robomaster_South China Regional Competition/digits/.changes b/dataset/dji_roco/robomaster_South China Regional Competition/digits/.changes
new file mode 100644
index 0000000000000000000000000000000000000000..9e26dfeeb6e641a33dae4961196235bdb965b21b
Binary files /dev/null and b/dataset/dji_roco/robomaster_South China Regional Competition/digits/.changes differ
diff --git a/poetry.lock b/poetry.lock
index d853b41c27feb9d0eaed9b23e008abaeed50ccfb..65d8fc1aef4096e6e0f754701d1b9f0652a7731a 100644
Binary files a/poetry.lock and b/poetry.lock differ
diff --git a/polystar_cv/config/settings.toml b/polystar_cv/config/settings.toml
index c81c609731b41f3cdeaab76d868056e0641be6ef..8145dfe306963de2b34a1777eaeae541a47c122a 100644
--- a/polystar_cv/config/settings.toml
+++ b/polystar_cv/config/settings.toml
@@ -1,9 +1,8 @@
[default]
-CAMERA_WIDTH = 1920
-CAMERA_HEIGHT = 1080
-CAMERA_HORIZONTAL_FOV = 120
+CAMERA = 'RASPI_V2'
-MODEL_NAME = 'robots/TRT_ssd_mobilenet_v2_roco.bin'
+OBJECTS_DETECTION_MODEL = "ssd_mobilenet_v2_roco_2018_03_29_20200314_015411_TWITCH_TEMP_733_IMGS_29595steps"
+ARMOR_DIGIT_MODEL = "20210117_145856_kd_cnn/distiled - temp 4.1e+01 - cnn - (32) - lr 7.8e-04 - drop 63.pkl"
[development]
diff --git a/polystar_cv/polystar/common/dependency_injection.py b/polystar_cv/polystar/common/dependency_injection.py
index 88eca3aadbaef5d55dbce6cb2492254d271eb408..5482b1ce49695a0612cea7fa0a6b41d71af0aca7 100644
--- a/polystar_cv/polystar/common/dependency_injection.py
+++ b/polystar_cv/polystar/common/dependency_injection.py
@@ -1,13 +1,36 @@
from dataclasses import dataclass
-from math import pi
+from typing import List
-from dynaconf import LazySettings
from injector import Injector, Module, multiprovider, provider, singleton
+from numpy.core._multiarray_umath import deg2rad
-from polystar.common.constants import LABEL_MAP_PATH
+from polystar.common.communication.print_target_sender import PrintTargetSender
+from polystar.common.communication.target_sender_abc import TargetSenderABC
+from polystar.common.constants import LABEL_MAP_PATH, MODELS_DIR
+from polystar.common.frame_generators.camera_frame_generator import RaspiV2CameraFrameGenerator, WebcamFrameGenerator
+from polystar.common.frame_generators.frames_generator_abc import FrameGeneratorABC
from polystar.common.models.camera import Camera
from polystar.common.models.label_map import LabelMap
-from polystar.common.settings import settings
+from polystar.common.settings import Settings, settings
+from polystar.common.target_pipeline.armors_descriptors.armors_color_descriptor import ArmorsColorDescriptor
+from polystar.common.target_pipeline.armors_descriptors.armors_descriptor_abc import ArmorsDescriptorABC
+from polystar.common.target_pipeline.armors_descriptors.armors_digit_descriptor import ArmorsDigitDescriptor
+from polystar.common.target_pipeline.detected_objects.detected_objects_factory import DetectedObjectFactory
+from polystar.common.target_pipeline.detected_objects.detected_robot import DetectedRobot
+from polystar.common.target_pipeline.object_selectors.closest_object_selector import ClosestObjectSelector
+from polystar.common.target_pipeline.object_selectors.object_selector_abc import ObjectSelectorABC
+from polystar.common.target_pipeline.objects_detectors.objects_detector_abc import ObjectsDetectorABC
+from polystar.common.target_pipeline.objects_detectors.tf_model_objects_detector import TFModelObjectsDetector
+from polystar.common.target_pipeline.objects_linker.objects_linker_abs import ObjectsLinkerABC
+from polystar.common.target_pipeline.objects_linker.simple_objects_linker import SimpleObjectsLinker
+from polystar.common.target_pipeline.objects_validators.confidence_object_validator import ConfidenceObjectValidator
+from polystar.common.target_pipeline.objects_validators.objects_validator_abc import ObjectsValidatorABC
+from polystar.common.target_pipeline.target_factories.ratio_simple_target_factory import RatioSimpleTargetFactory
+from polystar.common.target_pipeline.target_factories.target_factory_abc import TargetFactoryABC
+from polystar.common.utils.serialization import pkl_load
+from research.common.constants import PIPELINES_DIR
+from research.robots.armor_color.pipeline import ArmorColorPipeline
+from research.robots.armor_color.scripts.benchmark import MeanChannels, RedBlueComparisonClassifier
def make_injector() -> Injector:
@@ -16,16 +39,72 @@ def make_injector() -> Injector:
@dataclass
class CommonModule(Module):
- settings: LazySettings
+ settings: Settings
@provider
@singleton
def provide_camera(self) -> Camera:
- return Camera(
- self.settings.CAMERA_HORIZONTAL_FOV / 180 * pi, self.settings.CAMERA_WIDTH, self.settings.CAMERA_HEIGHT
- )
+ if settings.CAMERA == "RASPI_V2":
+ return Camera(
+ horizontal_fov=deg2rad(62.2),
+ vertical_fov=deg2rad(48.8),
+ pixel_size_m=1.12e-6,
+ focal_m=3.04e-3,
+ vertical_resolution=720,
+ horizontal_resolution=1_280,
+ )
+ raise ValueError(f"Camera {settings.CAMERA} not recognized")
@multiprovider
@singleton
def provide_label_map(self) -> LabelMap:
return LabelMap.from_file(LABEL_MAP_PATH)
+
+ @provider
+ @singleton
+ def provide_objects_detector(self, object_factory: DetectedObjectFactory) -> ObjectsDetectorABC:
+ if self.settings.is_dev:
+ return TFModelObjectsDetector(object_factory, MODELS_DIR / "robots" / settings.OBJECTS_DETECTION_MODEL)
+ import pycuda.autoinit # This is needed for initializing CUDA driver
+
+ raise NotImplementedError()
+
+ @multiprovider
+ @singleton
+ def provide_armor_descriptors(self) -> List[ArmorsDescriptorABC]:
+ return [
+ ArmorsColorDescriptor(ArmorColorPipeline.from_pipes([MeanChannels(), RedBlueComparisonClassifier()])),
+ ArmorsDigitDescriptor(pkl_load(PIPELINES_DIR / "armor-digit" / settings.ARMOR_DIGIT_MODEL)),
+ ]
+
+ @multiprovider
+ @singleton
+ def provide_objects_validators(self) -> List[ObjectsValidatorABC[DetectedRobot]]:
+ return [ConfidenceObjectValidator(0.6)]
+
+ @provider
+ @singleton
+ def provide_object_selector(self) -> ObjectSelectorABC:
+ return ClosestObjectSelector()
+
+ @provider
+ @singleton
+ def provide_target_factory(self, camera: Camera) -> TargetFactoryABC:
+ return RatioSimpleTargetFactory(camera, 0.3, 0.1)
+
+ @provider
+ @singleton
+ def provide_target_sender(self) -> TargetSenderABC:
+ return PrintTargetSender()
+
+ @provider
+ @singleton
+ def provide_objects_linker(self) -> ObjectsLinkerABC:
+ return SimpleObjectsLinker(min_percentage_intersection=0.8)
+
+ @provider
+ @singleton
+ def provide_webcam(self) -> FrameGeneratorABC:
+ if self.settings.is_prod:
+ return RaspiV2CameraFrameGenerator(1_280, 720)
+ return WebcamFrameGenerator()
diff --git a/polystar_cv/polystar/common/frame_generators/camera_frame_generator.py b/polystar_cv/polystar/common/frame_generators/camera_frame_generator.py
index eb3d4153847100bc476beb04e134b2470c60fcb3..853b7d09cee28561806c69534336388315376067 100644
--- a/polystar_cv/polystar/common/frame_generators/camera_frame_generator.py
+++ b/polystar_cv/polystar/common/frame_generators/camera_frame_generator.py
@@ -1,13 +1,13 @@
+from dataclasses import dataclass
from typing import Any, Iterable
import cv2
-from dataclasses import dataclass
from polystar.common.frame_generators.cv2_frame_generator_abc import CV2FrameGeneratorABC
@dataclass
-class CameraFrameGenerator(CV2FrameGeneratorABC):
+class RaspiV2CameraFrameGenerator(CV2FrameGeneratorABC):
width: int
height: int
@@ -23,3 +23,9 @@ class CameraFrameGenerator(CV2FrameGeneratorABC):
"videoconvert ! appsink",
cv2.CAP_GSTREAMER,
)
+
+
+@dataclass
+class WebcamFrameGenerator(CV2FrameGeneratorABC):
+ def _capture_params(self) -> Iterable[Any]:
+ return (0,)
diff --git a/polystar_cv/polystar/common/frame_generators/cv2_frame_generator_abc.py b/polystar_cv/polystar/common/frame_generators/cv2_frame_generator_abc.py
index c185c1882f7c0fe842dc8c46896329fd82b19a18..10ec363259b01dee55ea39a4a8b6f24f3b98ffbd 100644
--- a/polystar_cv/polystar/common/frame_generators/cv2_frame_generator_abc.py
+++ b/polystar_cv/polystar/common/frame_generators/cv2_frame_generator_abc.py
@@ -1,8 +1,8 @@
from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
from typing import Any, Iterable
import cv2
-from dataclasses import dataclass, field
from polystar.common.frame_generators.frames_generator_abc import FrameGeneratorABC
from polystar.common.models.image import Image
@@ -15,6 +15,7 @@ class CV2FrameGeneratorABC(FrameGeneratorABC, ABC):
def __enter__(self):
self._cap = cv2.VideoCapture(*self._capture_params())
+ assert self._cap.isOpened()
def __exit__(self, exc_type, exc_val, exc_tb):
self._cap.release()
diff --git a/polystar_cv/polystar/common/models/camera.py b/polystar_cv/polystar/common/models/camera.py
index 9e0c28f7c00e079e1d4808057f6353079def7df1..7da33a32b0ee61fa59f89afe27a16200f8aea707 100644
--- a/polystar_cv/polystar/common/models/camera.py
+++ b/polystar_cv/polystar/common/models/camera.py
@@ -3,7 +3,11 @@ from dataclasses import dataclass
@dataclass
class Camera:
- horizontal_angle: float
+ horizontal_fov: float
+ vertical_fov: float
- w: int
- h: int
+ pixel_size_m: float
+ focal_m: float
+
+ vertical_resolution: int
+ horizontal_resolution: int
diff --git a/polystar_cv/polystar/common/models/object.py b/polystar_cv/polystar/common/models/object.py
index 68860b6d5fc1ccdef9bfea31a77b131faa60ec18..1689a45cb32e06f02b806736cbd45820b2812feb 100644
--- a/polystar_cv/polystar/common/models/object.py
+++ b/polystar_cv/polystar/common/models/object.py
@@ -45,19 +45,23 @@ class ArmorDigit(NoCaseEnum): # CHANGING
# if self.value <= 5:
# return f"{self.value} ({self.name.title()})"
# return self.name.title()
- return f"{self.value + (self.value >= 2)} ({self.name.title()})" # hacky, but a number is missing (2)
+ return f"{self.digit} ({self.name.title()})" # hacky, but a number is missing (2)
@property
def short(self) -> str:
- return self.name[0] if self != ArmorDigit.UNKNOWN else "?"
+ return self.digit if self != ArmorDigit.UNKNOWN else "?"
+
+ @property
+ def digit(self) -> int:
+ return self.value + (self.value >= 2) # hacky, but a number is missing (2)
class ObjectType(NoCaseEnum):
- Car = auto()
- Watcher = auto()
- Base = auto()
- Armor = auto()
- Ignore = auto()
+ CAR = auto()
+ WATCHER = auto()
+ BASE = auto()
+ ARMOR = auto()
+ IGNORE = auto()
@dataclass
@@ -93,7 +97,7 @@ class ObjectFactory:
x, y = max(0, x), max(0, y)
- if t is not ObjectType.Armor:
+ if t is not ObjectType.ARMOR:
return Object(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
diff --git a/polystar_cv/polystar/common/models/tf_model.py b/polystar_cv/polystar/common/models/tf_model.py
index 468bb49b444edc88a9f9256c4543fef64e8475cd..3ed29141270e65da8666af81a2ee080e8ce363e2 100644
--- a/polystar_cv/polystar/common/models/tf_model.py
+++ b/polystar_cv/polystar/common/models/tf_model.py
@@ -1,5 +1,5 @@
from typing import NewType
-from tensorflow_core.python.eager.wrap_function import WrappedFunction
+from tensorflow.python.eager.wrap_function import WrappedFunction
TFModel = NewType("TFModel", WrappedFunction)
diff --git a/polystar_cv/polystar/common/pipeline/classification/classification_pipeline.py b/polystar_cv/polystar/common/pipeline/classification/classification_pipeline.py
index c85c7d666a7149fbcb5e83c5178aa9ba583e4591..1c1a6a97be1598fd1d84ede43ee0250e49ab591c 100644
--- a/polystar_cv/polystar/common/pipeline/classification/classification_pipeline.py
+++ b/polystar_cv/polystar/common/pipeline/classification/classification_pipeline.py
@@ -40,6 +40,8 @@ class ClassificationPipeline(Pipeline, Generic[IT, EnumT], ABC):
return self.predict_proba_and_classes(x)[1]
def predict_proba(self, x: Sequence[IT]) -> ndarray:
+ if not len(x):
+ return asarray([])
proba = super().predict_proba(x)
missing_classes = self.classifier.n_classes - proba.shape[1]
if not missing_classes:
diff --git a/polystar_cv/polystar/common/pipeline/classification/keras_classification_pipeline.py b/polystar_cv/polystar/common/pipeline/classification/keras_classification_pipeline.py
index 388e33dcbad7083144dcd902bf98e50c3c644330..d2bcfc09aacbb3edd68a169b459dc6595ec3ad1f 100644
--- a/polystar_cv/polystar/common/pipeline/classification/keras_classification_pipeline.py
+++ b/polystar_cv/polystar/common/pipeline/classification/keras_classification_pipeline.py
@@ -121,6 +121,7 @@ class KerasClassificationPipeline(ClassificationPipeline):
input_shape, conv_blocks=conv_blocks, dense_size=dense_size, output_size=cls.n_classes, dropout=dropout,
),
trainer=KerasTrainer(
+ data_preparator=KerasDataPreparator(batch_size=32, steps=100),
model_preparator=Distiller(temperature=temperature, teacher_model=teacher_pipeline.classifier.model),
compilation_parameters=KerasCompilationParameters(
loss=DistillationLoss(temperature=temperature, n_classes=cls.n_classes),
diff --git a/polystar_cv/polystar/common/pipeline/keras/classifier.py b/polystar_cv/polystar/common/pipeline/keras/classifier.py
index 8568a878692d293ec2423d74573ddd66600ff7e7..05d55ba7d464d4babbba1151fe55e64225bfc28d 100644
--- a/polystar_cv/polystar/common/pipeline/keras/classifier.py
+++ b/polystar_cv/polystar/common/pipeline/keras/classifier.py
@@ -3,12 +3,14 @@ from tempfile import NamedTemporaryFile
from typing import Dict, List, Optional, Sequence
from numpy import asarray
+from tensorflow import Graph, Session
from tensorflow.python.keras.models import Model, load_model
from tensorflow.python.keras.utils.np_utils import to_categorical
from polystar.common.models.image import Image
from polystar.common.pipeline.classification.classifier_abc import ClassifierABC
from polystar.common.pipeline.keras.trainer import KerasTrainer
+from polystar.common.settings import settings
from polystar.common.utils.registry import registry
@@ -30,6 +32,9 @@ class KerasClassifier(ClassifierABC):
return self
def predict_proba(self, examples: List[Image]) -> Sequence[float]:
+ if settings.is_prod: # FIXME
+ with self.graph.as_default(), self.session.as_default():
+ return self.model.predict(asarray(examples))
return self.model.predict(asarray(examples))
def __getstate__(self) -> Dict:
@@ -43,10 +48,14 @@ class KerasClassifier(ClassifierABC):
def __setstate__(self, state: Dict):
self.__dict__.update(state)
+ self.graph = Graph()
with NamedTemporaryFile(suffix=".hdf5", delete=True) as fd:
fd.write(state.pop("model_str"))
fd.flush()
- self.model = load_model(fd.name)
+ with self.graph.as_default():
+ self.session = Session(graph=self.graph)
+ with self.session.as_default():
+ self.model = load_model(fd.name, compile=False)
self.trainer = None
@property
diff --git a/polystar_cv/polystar/common/pipeline/keras/cnn.py b/polystar_cv/polystar/common/pipeline/keras/cnn.py
index 8e3cecfaacf29b9229d9e62827015c3cf509f250..4f3e45724ae1484fc5822b4164501e862b2bcbe0 100644
--- a/polystar_cv/polystar/common/pipeline/keras/cnn.py
+++ b/polystar_cv/polystar/common/pipeline/keras/cnn.py
@@ -1,6 +1,6 @@
from typing import Sequence, Tuple
-from tensorflow.python.keras import Input, Sequential
+from tensorflow.python.keras import Sequential
from tensorflow.python.keras.layers import Conv2D, Dense, Dropout, Flatten, MaxPooling2D, Softmax
@@ -12,10 +12,14 @@ def make_cnn_model(
dropout: float,
) -> Sequential:
model = Sequential()
- model.add(Input((*input_shape, 3)))
+ model.add(Conv2D(conv_blocks[0][0], (3, 3), activation="relu", input_shape=(*input_shape, 3)))
+ is_first = True
for conv_sizes in conv_blocks:
for size in conv_sizes:
+ if is_first:
+ is_first = False
+ continue
model.add(Conv2D(size, (3, 3), activation="relu"))
model.add(MaxPooling2D())
diff --git a/polystar_cv/polystar/common/pipeline/keras/distillation.py b/polystar_cv/polystar/common/pipeline/keras/distillation.py
index 79b42c630feaedb9e76f591d4748e54efac8922a..6b82771a990de291bc72452ee1f77c2182b53ddb 100644
--- a/polystar_cv/polystar/common/pipeline/keras/distillation.py
+++ b/polystar_cv/polystar/common/pipeline/keras/distillation.py
@@ -2,25 +2,22 @@ from typing import Callable
from tensorflow.python.keras import Input, Model, Sequential
from tensorflow.python.keras.layers import Softmax, concatenate
-from tensorflow.python.keras.losses import KLDivergence
-from tensorflow.python.keras.models import Model
+from tensorflow.python.keras.losses import Loss, kullback_leibler_divergence
from tensorflow.python.ops.nn_ops import softmax
from polystar.common.pipeline.keras.model_preparator import KerasModelPreparator
-class DistillationLoss(KLDivergence):
+class DistillationLoss(Loss):
def __init__(self, temperature: float, n_classes: int):
- super().__init__()
+ super().__init__(name="kd_loss")
self.n_classes = n_classes
self.temperature = temperature
- def __call__(self, y_true, y_pred, sample_weight=None):
+ def call(self, y_true, y_pred):
teacher_logits, student_logits = y_pred[:, : self.n_classes], y_pred[:, self.n_classes :]
- return super().__call__(
- softmax(teacher_logits / self.temperature, axis=1),
- softmax(student_logits / self.temperature, axis=1),
- sample_weight=sample_weight,
+ return kullback_leibler_divergence(
+ softmax(teacher_logits / self.temperature, axis=1), softmax(student_logits / self.temperature, axis=1)
)
diff --git a/polystar_cv/polystar/common/pipeline/keras/trainer.py b/polystar_cv/polystar/common/pipeline/keras/trainer.py
index 9db4ba7716ccb2efb09e03440cb0400ff4de40e3..16826b645940eb512f69e2d00b5ce78e8ae2f8b2 100644
--- a/polystar_cv/polystar/common/pipeline/keras/trainer.py
+++ b/polystar_cv/polystar/common/pipeline/keras/trainer.py
@@ -14,7 +14,7 @@ from polystar.common.pipeline.keras.model_preparator import KerasModelPreparator
class KerasTrainer:
compilation_parameters: KerasCompilationParameters
callbacks: List[Callback]
- data_preparator: KerasDataPreparator = field(default_factory=KerasDataPreparator)
+ data_preparator: KerasDataPreparator
model_preparator: KerasModelPreparator = field(default_factory=KerasModelPreparator)
max_epochs: int = 300
verbose: int = 0
diff --git a/polystar_cv/polystar/common/pipeline/keras/transfer_learning.py b/polystar_cv/polystar/common/pipeline/keras/transfer_learning.py
index a6a7b4747e2186f9f61672538359bd51e2fdd291..e3ff9dcfb63682cff9a48309b41cdf8f6e859f32 100644
--- a/polystar_cv/polystar/common/pipeline/keras/transfer_learning.py
+++ b/polystar_cv/polystar/common/pipeline/keras/transfer_learning.py
@@ -1,6 +1,7 @@
from typing import Callable, Tuple
-from tensorflow.python.keras import Input, Model, Sequential
+from tensorflow.python.keras import Sequential
+from tensorflow.python.keras.engine import InputLayer
from tensorflow.python.keras.layers import Dense, Dropout, Flatten, Softmax
from tensorflow.python.keras.models import Model
@@ -13,7 +14,7 @@ def make_transfer_learning_model(
return Sequential(
[
- Input(input_shape),
+ InputLayer(input_shape),
base_model,
Flatten(),
Dense(dense_size, activation="relu"),
diff --git a/polystar_cv/polystar/common/settings.py b/polystar_cv/polystar/common/settings.py
index 23b13f2a08ec612b40073c131cdc3e172147f516..44ad5c0c653519b7a55638321bd9060bde2a3410 100644
--- a/polystar_cv/polystar/common/settings.py
+++ b/polystar_cv/polystar/common/settings.py
@@ -1,5 +1,4 @@
from enum import Enum
-from pathlib import Path
from dynaconf import LazySettings
@@ -12,18 +11,23 @@ class Environment(str, Enum):
class Settings(LazySettings):
- def get_env(self) -> Environment:
+ @property
+ def env(self) -> Environment:
return Environment(self.current_env.lower())
+ @property
+ def is_prod(self) -> bool:
+ return self.env == Environment.PRODUCTION
-def _config_file_for_project(project_name: str) -> Path:
- return PROJECT_DIR / "config" / "settings.toml"
+ @property
+ def is_dev(self) -> bool:
+ return self.env == Environment.DEVELOPMENT
def make_settings() -> LazySettings:
- return LazySettings(
+ return Settings(
SILENT_ERRORS_FOR_DYNACONF=False,
- SETTINGS_FILE_FOR_DYNACONF=f"{PROJECT_DIR / 'config' / 'settings.toml'}",
+ SETTINGS_FILE_FOR_DYNACONF=f"{PROJECT_DIR / 'polystar_cv' / 'config' / 'settings.toml'}",
ENV_SWITCHER_FOR_DYNACONF="POLYSTAR_ENV",
)
diff --git a/polystar_cv/polystar/common/target_pipeline/armors_descriptors/armors_color_descriptor.py b/polystar_cv/polystar/common/target_pipeline/armors_descriptors/armors_color_descriptor.py
index c996df4da35561564156b9e293968a62078ef5ad..9d99d30f94012c97d372b9f6d2237be6ae93cc53 100644
--- a/polystar_cv/polystar/common/target_pipeline/armors_descriptors/armors_color_descriptor.py
+++ b/polystar_cv/polystar/common/target_pipeline/armors_descriptors/armors_color_descriptor.py
@@ -2,15 +2,18 @@ from dataclasses import dataclass
from typing import List
from polystar.common.models.image import Image
-from polystar.common.pipeline.classification.classification_pipeline import ClassificationPipeline
from polystar.common.target_pipeline.armors_descriptors.armors_descriptor_abc import ArmorsDescriptorABC
from polystar.common.target_pipeline.detected_objects.detected_armor import DetectedArmor
+from research.robots.armor_color.pipeline import ArmorColorPipeline
@dataclass
class ArmorsColorDescriptor(ArmorsDescriptorABC):
- image_pipeline: ClassificationPipeline
+ image_pipeline: ArmorColorPipeline
+
+ def __post_init__(self):
+ assert ArmorColorPipeline.classes == self.image_pipeline.classes
def _describe_armors_from_images(self, armors_images: List[Image], armors: List[DetectedArmor]):
colors_predictions = self.image_pipeline.predict_proba(armors_images)
diff --git a/polystar_cv/polystar/common/target_pipeline/armors_descriptors/armors_digit_descriptor.py b/polystar_cv/polystar/common/target_pipeline/armors_descriptors/armors_digit_descriptor.py
new file mode 100644
index 0000000000000000000000000000000000000000..79c80691758a0141fbc6bedc56126aee53b9a1eb
--- /dev/null
+++ b/polystar_cv/polystar/common/target_pipeline/armors_descriptors/armors_digit_descriptor.py
@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+from typing import List
+
+from polystar.common.models.image import Image
+from polystar.common.target_pipeline.armors_descriptors.armors_descriptor_abc import ArmorsDescriptorABC
+from polystar.common.target_pipeline.detected_objects.detected_armor import DetectedArmor
+from research.robots.armor_digit.pipeline import ArmorDigitPipeline
+
+
+@dataclass
+class ArmorsDigitDescriptor(ArmorsDescriptorABC):
+
+ image_pipeline: ArmorDigitPipeline
+
+ def __post_init__(self):
+ assert ArmorDigitPipeline.classes == self.image_pipeline.classes
+
+ def _describe_armors_from_images(self, armors_images: List[Image], armors: List[DetectedArmor]):
+ digit_predictions = self.image_pipeline.predict_proba(armors_images)
+ for digits_proba, armor in zip(digit_predictions, armors):
+ armor.digits_proba = digits_proba
diff --git a/polystar_cv/polystar/common/target_pipeline/debug_pipeline.py b/polystar_cv/polystar/common/target_pipeline/debug_pipeline.py
index e0e9ec8c0c3c90ddf9b637df96ddc58e9a6328a4..7a44ec74c2821ad735c89210afc87f173da73fa8 100644
--- a/polystar_cv/polystar/common/target_pipeline/debug_pipeline.py
+++ b/polystar_cv/polystar/common/target_pipeline/debug_pipeline.py
@@ -1,6 +1,8 @@
from dataclasses import dataclass, field
from typing import List
+from injector import inject
+
from polystar.common.models.image import Image
from polystar.common.target_pipeline.detected_objects.detected_armor import DetectedArmor
from polystar.common.target_pipeline.detected_objects.detected_object import DetectedObject
@@ -18,6 +20,7 @@ class DebugInfo:
target: TargetABC = field(init=False, default=None)
+@inject
@dataclass
class DebugTargetPipeline(TargetPipeline):
"""Wrap a pipeline with debug, to store debug infos"""
diff --git a/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_armor.py b/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_armor.py
index 3bfaf2b654f30647057b271891f414db70970413..2d1a538b480b5e66d906c9993dbcce18c157b04a 100644
--- a/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_armor.py
+++ b/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_armor.py
@@ -1,16 +1,17 @@
from dataclasses import dataclass, field
import numpy as np
-from numpy import argmax
from polystar.common.models.object import ArmorColor, ArmorDigit, ObjectType
from polystar.common.target_pipeline.detected_objects.detected_object import DetectedObject
+from research.robots.armor_color.pipeline import ArmorColorPipeline
+from research.robots.armor_digit.pipeline import ArmorDigitPipeline
@dataclass
class DetectedArmor(DetectedObject):
def __post_init__(self):
- assert self.type == ObjectType.Armor
+ assert self.type == ObjectType.ARMOR
colors_proba: np.ndarray = field(init=False, default=None)
digits_proba: np.ndarray = field(init=False, default=None)
@@ -24,7 +25,7 @@ class DetectedArmor(DetectedObject):
return self._color
if self.colors_proba is not None:
- self._color = ArmorColor(self.colors_proba.argmax() + 1)
+ self._color = ArmorDigitPipeline.classes[self.colors_proba.argmax()]
return self._color
return ArmorColor.UNKNOWN
@@ -34,8 +35,8 @@ class DetectedArmor(DetectedObject):
if self._digit is not None:
return self._digit
- if self.digits_proba:
- self._digit = ArmorDigit(argmax(self.colors_proba) + 1)
+ if self.digits_proba is not None:
+ self._digit = ArmorColorPipeline.classes[self.digits_proba.argmax()]
return self._digit
return ArmorDigit.UNKNOWN
diff --git a/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_objects_factory.py b/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_objects_factory.py
index 543ba72b93f96ada3aa65d363a3efcb31218f9a0..60398099cd23db51e8ef52583e8500722979095d 100644
--- a/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_objects_factory.py
+++ b/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_objects_factory.py
@@ -1,6 +1,8 @@
from dataclasses import dataclass
from typing import List, Tuple, Type, Union
+from injector import inject
+
from polystar.common.models.box import Box
from polystar.common.models.image import Image
from polystar.common.models.label_map import LabelMap
@@ -21,6 +23,7 @@ class ObjectParams:
class DetectedObjectFactory:
+ @inject
def __init__(self, label_map: LabelMap, armors_descriptors: List[ArmorsDescriptorABC]):
self.armors_descriptors = armors_descriptors
self.label_map = label_map
@@ -57,4 +60,4 @@ class DetectedObjectFactory:
@staticmethod
def _get_object_class_from_type(object_type: ObjectType) -> Type[Union[DetectedRobot, DetectedArmor]]:
- return DetectedArmor if object_type is ObjectType.Armor else DetectedRobot
+ return DetectedArmor if object_type is ObjectType.ARMOR else DetectedRobot
diff --git a/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_robot.py b/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_robot.py
index 372ed7f156226f47389564355e14fd400dcf88c0..7d9f7ff25a65ff9afd03e6f80296cf6b6b1a4863 100644
--- a/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_robot.py
+++ b/polystar_cv/polystar/common/target_pipeline/detected_objects/detected_robot.py
@@ -13,4 +13,4 @@ class DetectedRobot(DetectedObject):
class FakeDetectedRobot(DetectedRobot):
def __init__(self, armor: DetectedArmor):
- super().__init__(type=ObjectType.Car, box=armor.box, confidence=0, armors=[armor])
+ super().__init__(type=ObjectType.CAR, box=armor.box, confidence=0, armors=[armor])
diff --git a/polystar_cv/polystar/common/target_pipeline/objects_detectors/tf_model_objects_detector.py b/polystar_cv/polystar/common/target_pipeline/objects_detectors/tf_model_objects_detector.py
index c4282a55ade9d687d9dcf5268526ba72d358a200..2ef9a49d1e7ae9d5e257f3a166cf5f0475221dc6 100644
--- a/polystar_cv/polystar/common/target_pipeline/objects_detectors/tf_model_objects_detector.py
+++ b/polystar_cv/polystar/common/target_pipeline/objects_detectors/tf_model_objects_detector.py
@@ -1,14 +1,16 @@
-from dataclasses import dataclass
+from dataclasses import InitVar, dataclass
+from pathlib import Path
from typing import Dict, List, Tuple
import numpy as np
import tensorflow as tf
+from tensorflow import GraphDef, Session
+from tensorflow.python.platform.gfile import GFile
from polystar.common.models.image import Image
from polystar.common.models.tf_model import TFModel
from polystar.common.target_pipeline.detected_objects.detected_armor import DetectedArmor
-from polystar.common.target_pipeline.detected_objects.detected_objects_factory import (DetectedObjectFactory,
- ObjectParams)
+from polystar.common.target_pipeline.detected_objects.detected_objects_factory import ObjectParams
from polystar.common.target_pipeline.detected_objects.detected_robot import DetectedRobot
from polystar.common.target_pipeline.objects_detectors.objects_detector_abc import ObjectsDetectorABC
@@ -16,6 +18,46 @@ from polystar.common.target_pipeline.objects_detectors.objects_detector_abc impo
@dataclass
class TFModelObjectsDetector(ObjectsDetectorABC):
+ model_path: InitVar[Path]
+
+ def __post_init__(self, model_path: Path):
+ self.graph = tf.Graph()
+ with self.graph.as_default():
+ od_graph_def = GraphDef()
+ with GFile(str(model_path / "frozen_inference_graph.pb"), "rb") as fid:
+ serialized_graph = fid.read()
+ od_graph_def.ParseFromString(serialized_graph)
+ tf.import_graph_def(od_graph_def, name="")
+
+ def detect(self, image: Image) -> Tuple[List[DetectedRobot], List[DetectedArmor]]:
+ with self.graph.as_default(), Session(graph=self.graph) as session:
+ image_np_expanded = np.expand_dims(image, axis=0)
+ image_tensor = self.graph.get_tensor_by_name("image_tensor:0")
+ boxes = self.graph.get_tensor_by_name("detection_boxes:0")
+ scores = self.graph.get_tensor_by_name("detection_scores:0")
+ classes = self.graph.get_tensor_by_name("detection_classes:0")
+ num_detections = self.graph.get_tensor_by_name("num_detections:0")
+ boxes, scores, classes, num_detections = session.run(
+ [boxes, scores, classes, num_detections], feed_dict={image_tensor: image_np_expanded}
+ )
+ return self._construct_objects_from_tf_results(image, boxes[0], scores[0], classes[0])
+
+ def _construct_objects_from_tf_results(
+ self, image: Image, boxes: np.ndarray, scores: np.ndarray, classes: np.ndarray
+ ) -> Tuple[List[DetectedRobot], List[DetectedArmor]]:
+ return self.objects_factory.make_lists(
+ [
+ ObjectParams(ymin=ymin, xmin=xmin, ymax=ymax, xmax=xmax, score=score, object_class_id=object_class_id)
+ for (ymin, xmin, ymax, xmax), object_class_id, score in zip(boxes, classes, scores)
+ if score >= 0.1
+ ],
+ image,
+ )
+
+
+@dataclass
+class TFV2ModelObjectsDetector(ObjectsDetectorABC):
+
model: TFModel
def detect(self, image: Image) -> Tuple[List[DetectedRobot], List[DetectedArmor]]:
@@ -30,7 +72,7 @@ class TFModelObjectsDetector(ObjectsDetectorABC):
return input_tensor
def _make_single_prediction(self, input_tensor: tf.Tensor) -> Dict[str, np.array]:
- output_dict: Dict[str, tf.Tensor] = self.model(input_tensor) # typing is correct despite PyCharm's saying
+ output_dict: Dict[str, tf.Tensor] = self.model(input_tensor)
return self._normalize_prediction(output_dict)
@staticmethod
diff --git a/polystar_cv/polystar/common/target_pipeline/objects_linker/simple_objects_linker.py b/polystar_cv/polystar/common/target_pipeline/objects_linker/simple_objects_linker.py
index 8500e3eb3d0b5fe7c2cd8c2276f8099bdbf54884..86dd41d700442947e2e12d3c070bd5ac664f1d54 100644
--- a/polystar_cv/polystar/common/target_pipeline/objects_linker/simple_objects_linker.py
+++ b/polystar_cv/polystar/common/target_pipeline/objects_linker/simple_objects_linker.py
@@ -14,8 +14,8 @@ class SimpleObjectsLinker(ObjectsLinkerABC):
def __init__(self, min_percentage_intersection: float):
super().__init__()
self.min_percentage_intersection = min_percentage_intersection
- self.robots_filter = NegationValidator(TypeObjectValidator(ObjectType.Armor))
- self.armors_filter = TypeObjectValidator(ObjectType.Armor)
+ self.robots_filter = NegationValidator(TypeObjectValidator(ObjectType.ARMOR))
+ self.armors_filter = TypeObjectValidator(ObjectType.ARMOR)
def link_armors_to_robots(
self, robots: List[DetectedRobot], armors: List[DetectedArmor], image: Image
diff --git a/polystar_cv/polystar/common/target_pipeline/objects_validators/armor_digit_validator.py b/polystar_cv/polystar/common/target_pipeline/objects_validators/armor_digit_validator.py
new file mode 100644
index 0000000000000000000000000000000000000000..060d78c7c3fa6b3949fd384c8ad38b19e6b338d0
--- /dev/null
+++ b/polystar_cv/polystar/common/target_pipeline/objects_validators/armor_digit_validator.py
@@ -0,0 +1,15 @@
+from typing import Iterable
+
+from numpy.core.multiarray import ndarray
+
+from polystar.common.models.object import Armor
+from polystar.common.target_pipeline.detected_objects.detected_robot import DetectedRobot
+from polystar.common.target_pipeline.objects_validators.objects_validator_abc import ObjectsValidatorABC
+
+
+class ArmorDigitValidator(ObjectsValidatorABC[DetectedRobot]):
+ def __init__(self, digits: Iterable[int]):
+ self.digits = digits
+
+ def validate_single(self, armor: Armor, image: ndarray) -> bool:
+ return isinstance(armor, Armor) and armor.number in self.digits
diff --git a/polystar_cv/polystar/common/target_pipeline/objects_validators/objects_validator_abc.py b/polystar_cv/polystar/common/target_pipeline/objects_validators/objects_validator_abc.py
index e75721dbe4e791698b85801226e2d0b8ea355aef..075ce1521f5764bc98ba8fc10ae473514b83d0a8 100644
--- a/polystar_cv/polystar/common/target_pipeline/objects_validators/objects_validator_abc.py
+++ b/polystar_cv/polystar/common/target_pipeline/objects_validators/objects_validator_abc.py
@@ -1,5 +1,5 @@
-from abc import abstractmethod, ABC
-from typing import List, TypeVar, Generic
+from abc import ABC, abstractmethod
+from typing import Generic, List, TypeVar
import numpy as np
@@ -8,7 +8,7 @@ from polystar.common.models.object import Object
ObjectT = TypeVar("ObjectT", bound=Object)
-class ObjectsValidatorABC(Generic[ObjectT], ABC):
+class ObjectsValidatorABC(Generic[ObjectT], ABC): # FIXME Filter would do here
def filter(self, objects: List[ObjectT], image: np.ndarray) -> List[ObjectT]:
return [obj for obj, is_valid in zip(objects, self.validate(objects, image)) if is_valid]
diff --git a/polystar_cv/polystar/common/target_pipeline/target_factories/ratio_target_factory_abc.py b/polystar_cv/polystar/common/target_pipeline/target_factories/ratio_target_factory_abc.py
index 1f8a3917b372eb22ab265813558dba2ddfbad018..c58a2207bf118147b0e06b180a0a3e50cbf515a1 100644
--- a/polystar_cv/polystar/common/target_pipeline/target_factories/ratio_target_factory_abc.py
+++ b/polystar_cv/polystar/common/target_pipeline/target_factories/ratio_target_factory_abc.py
@@ -1,5 +1,5 @@
from abc import ABC
-from math import sin, asin, sqrt, atan2
+from math import atan2, pi, tan
from typing import Tuple
import numpy as np
@@ -10,17 +10,16 @@ from polystar.common.target_pipeline.target_factories.target_factory_abc import
class RatioTargetFactoryABC(TargetFactoryABC, ABC):
- def __init__(self, camera: Camera, real_width: float, real_height: float):
- self._ratio_w = real_width * camera.w // 2 / sin(camera.horizontal_angle)
- # self._ratio_h = real_height * camera.h // 2 / sin(camera.vertical_angle)
-
- self._coef_angle = sin(camera.horizontal_angle) / (camera.w // 2)
+ def __init__(self, camera: Camera, real_width_m: float, real_height_m: float):
+ self._vertical_distance_coef = camera.focal_m * real_height_m / camera.pixel_size_m
+ self._vertical_angle_distance = camera.vertical_resolution / (2 * tan(camera.vertical_fov))
+ self._horizontal_angle_distance = camera.horizontal_resolution / (2 * tan(camera.horizontal_fov))
def _calculate_angles(self, obj: DetectedObject, image: np.ndarray) -> Tuple[float, float]:
x, y = obj.box.x + obj.box.w // 2 - image.shape[1] // 2, image.shape[0] // 2 - obj.box.y + obj.box.h // 2
- phi = asin(sqrt(x ** 2 + y ** 2) * self._coef_angle)
- theta = atan2(y, x)
+ phi = -atan2(x, self._horizontal_angle_distance)
+ theta = pi / 2 - atan2(y, self._vertical_angle_distance)
return phi, theta
def _calculate_distance(self, obj: DetectedObject) -> float:
- return self._ratio_w / obj.box.w
+ return self._vertical_distance_coef / obj.box.h
diff --git a/polystar_cv/polystar/common/target_pipeline/target_pipeline.py b/polystar_cv/polystar/common/target_pipeline/target_pipeline.py
index f226380a189ce6643cee2cc5a3d78528e46cca4d..9b51450094f407ff103e1d0a5cd76cc41ffaabd7 100644
--- a/polystar_cv/polystar/common/target_pipeline/target_pipeline.py
+++ b/polystar_cv/polystar/common/target_pipeline/target_pipeline.py
@@ -1,6 +1,8 @@
from dataclasses import dataclass
from typing import List
+from injector import inject
+
from polystar.common.communication.target_sender_abc import TargetSenderABC
from polystar.common.models.image import Image
from polystar.common.target_pipeline.detected_objects.detected_object import DetectedObject
@@ -17,6 +19,7 @@ class NoTargetFoundException(Exception):
pass
+@inject
@dataclass
class TargetPipeline:
diff --git a/polystar_cv/polystar/common/utils/fps.py b/polystar_cv/polystar/common/utils/fps.py
new file mode 100644
index 0000000000000000000000000000000000000000..02e9883eaafb13f2e3a7b203e92d55a36086f05e
--- /dev/null
+++ b/polystar_cv/polystar/common/utils/fps.py
@@ -0,0 +1,23 @@
+from time import time
+
+
+class FPS:
+ def __init__(self, inertia: float = 0.9):
+ self.inertia = inertia
+ self.previous_time = time()
+ self.fps = 0
+
+ def __str__(self) -> str:
+ return str(self.fps)
+
+ def __format__(self, format_spec: str):
+ return format(self.fps, format_spec)
+
+ def tick(self) -> float:
+ t = time()
+ self.fps = self.inertia * self.fps + (1 - self.inertia) / (t - self.previous_time)
+ self.previous_time = t
+ return self.fps
+
+ def skip(self):
+ self.previous_time = time()
diff --git a/polystar_cv/polystar/common/utils/iterable_utils.py b/polystar_cv/polystar/common/utils/iterable_utils.py
index 01bc2da41ef0b4ba3894a14cdf989c745c6829fe..1efac7228ee1aea8d161613ade1422b74525ad13 100644
--- a/polystar_cv/polystar/common/utils/iterable_utils.py
+++ b/polystar_cv/polystar/common/utils/iterable_utils.py
@@ -27,3 +27,14 @@ def group_by(it: Iterable[T], key: Callable[[T], U]) -> Dict[U, List[T]]:
for item in it:
rv[key(item)].append(item)
return rv
+
+
+def chunk(it: Iterable[T], batch_size: float) -> Iterable[List[T]]:
+ batch = []
+ for el in it:
+ batch.append(el)
+ if len(batch) == batch_size:
+ yield batch
+ batch = []
+ if batch:
+ yield batch
diff --git a/polystar_cv/polystar/common/view/results_viewer_abc.py b/polystar_cv/polystar/common/view/results_viewer_abc.py
index 619b47e492abd518f1b3e23eb4bda82058d6e776..6e89ad2bacdd131719d32fe29107548efa70e40e 100644
--- a/polystar_cv/polystar/common/view/results_viewer_abc.py
+++ b/polystar_cv/polystar/common/view/results_viewer_abc.py
@@ -1,6 +1,6 @@
from abc import ABC, abstractmethod
from itertools import cycle
-from typing import Iterable, NewType, Sequence, Tuple
+from typing import Iterable, Sequence, Tuple
from polystar.common.models.image import Image
from polystar.common.models.image_annotation import ImageAnnotation
@@ -8,7 +8,7 @@ from polystar.common.models.object import Object
from polystar.common.target_pipeline.debug_pipeline import DebugInfo
from polystar.common.target_pipeline.detected_objects.detected_robot import DetectedRobot, FakeDetectedRobot
-ColorView = NewType("ColorView", Tuple[float, float, float])
+ColorView = Tuple[float, float, float]
class ResultViewerABC(ABC):
@@ -55,12 +55,15 @@ class ResultViewerABC(ABC):
self.display_image_with_objects(annotation.image, annotation.objects)
def display_debug_info(self, debug_info: DebugInfo):
+ self.add_debug_info(debug_info)
+ self.display()
+
+ def add_debug_info(self, debug_info: DebugInfo):
self.new(debug_info.image)
self.add_robots(debug_info.detected_robots, forced_color=(0.3, 0.3, 0.3))
self.add_robots(debug_info.validated_robots)
if debug_info.selected_armor is not None:
self.add_object(debug_info.selected_armor)
- self.display()
def add_robot(self, robot: DetectedRobot, forced_color: ColorView = None):
objects = robot.armors
diff --git a/polystar_cv/research/common/constants.py b/polystar_cv/research/common/constants.py
index 8d90b8acd48eaf627b2bbc55086b0a23996335b4..f2b7e8c7bd365474e9b9a00726c92917f08091d2 100644
--- a/polystar_cv/research/common/constants.py
+++ b/polystar_cv/research/common/constants.py
@@ -20,3 +20,4 @@ TWITCH_DSET_ROBOTS_VIEWS_DIR.mkdir(parents=True, exist_ok=True)
EVALUATION_DIR: Path = PROJECT_DIR / "experiments"
+PIPELINES_DIR = PROJECT_DIR / "pipelines"
diff --git a/polystar_cv/research/common/dataset/cleaning/dataset_changes.py b/polystar_cv/research/common/dataset/cleaning/dataset_changes.py
index bdeb4741876923230fc2f1f823bf1c6ddbec0062..937b834051fde928b06c5a31020cafa36ee67e6c 100644
--- a/polystar_cv/research/common/dataset/cleaning/dataset_changes.py
+++ b/polystar_cv/research/common/dataset/cleaning/dataset_changes.py
@@ -1,5 +1,5 @@
import json
-from contextlib import suppress
+import logging
from pathlib import Path
from typing import Dict, List, Set
@@ -10,13 +10,18 @@ from polystar.common.utils.time import create_time_id
from research.common.gcloud.gcloud_storage import GCStorages
INVALIDATED_KEY: str = "invalidated"
+logger = logging.getLogger(__name__)
class DatasetChanges:
def __init__(self, dataset_directory: Path):
self.changes_file: Path = dataset_directory / ".changes"
- with suppress(FileNotFoundError):
+ try:
GCStorages.DEV.download_file_if_missing(self.changes_file)
+ except FileNotFoundError:
+ self.changes_file.write_text("{}")
+ except ConnectionError:
+ logger.warning(f"Can't load {self.changes_file} because of no internet connection")
@property
def invalidated(self) -> Set[str]:
diff --git a/polystar_cv/research/common/datasets/dataset_builder.py b/polystar_cv/research/common/datasets/dataset_builder.py
index fd8650f91a2d7c8fb88bf916e35b9fdd626acfc9..7b831c86b1bc693a85063c8cd1491109d0f7efaa 100644
--- a/polystar_cv/research/common/datasets/dataset_builder.py
+++ b/polystar_cv/research/common/datasets/dataset_builder.py
@@ -6,6 +6,7 @@ from polystar.common.utils.misc import identity
from research.common.datasets.dataset import Dataset
from research.common.datasets.filter_dataset import ExampleU, FilterDataset, TargetU
from research.common.datasets.lazy_dataset import ExampleT, LazyDataset, TargetT
+from research.common.datasets.shuffle_dataset import ShuffleDataset
from research.common.datasets.slice_dataset import SliceDataset
from research.common.datasets.transform_dataset import TransformDataset
@@ -58,6 +59,10 @@ class DatasetBuilder(Generic[ExampleT, TargetT], Iterable[Tuple[ExampleT, Target
self.dataset = TransformDataset(self.dataset, identity, target_transformer)
return self
+ def shuffle(self) -> "DatasetBuilder[ExampleT, TargetU]":
+ self.dataset = ShuffleDataset(self.dataset)
+ return self
+
def cap(self, n: int) -> "DatasetBuilder[ExampleT, TargetT]":
self.dataset = SliceDataset(self.dataset, stop=n)
return self
diff --git a/polystar_cv/research/common/datasets/roco/roco_annotation_filters/roco_annotation_object_filter.py b/polystar_cv/research/common/datasets/roco/roco_annotation_filters/roco_annotation_object_filter.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b9d61db7f1f819f26a7f778ec427189b12a00be
--- /dev/null
+++ b/polystar_cv/research/common/datasets/roco/roco_annotation_filters/roco_annotation_object_filter.py
@@ -0,0 +1,11 @@
+from polystar.common.filters.filter_abc import FilterABC
+from polystar.common.target_pipeline.objects_validators.objects_validator_abc import ObjectsValidatorABC
+from research.common.datasets.roco.roco_annotation import ROCOAnnotation
+
+
+class ROCOAnnotationObjectFilter(FilterABC):
+ def __init__(self, object_validator: ObjectsValidatorABC):
+ self.object_validator = object_validator
+
+ def validate_single(self, annotation: ROCOAnnotation) -> bool:
+ return any(self.object_validator.validate(annotation.objects, None))
diff --git a/polystar_cv/research/common/datasets/roco/roco_dataset_descriptor.py b/polystar_cv/research/common/datasets/roco/roco_dataset_descriptor.py
index 070bc9fb9404db103899c04143e7ededdc15d83a..4f2a76535314daa368cea1fe78fefe30dbf63fd1 100644
--- a/polystar_cv/research/common/datasets/roco/roco_dataset_descriptor.py
+++ b/polystar_cv/research/common/datasets/roco/roco_dataset_descriptor.py
@@ -33,11 +33,11 @@ class ROCODatasetStats:
rv.n_images += 1
rv.n_runes += annotation.has_rune
for obj in annotation.objects:
- if obj.type == ObjectType.Car:
+ if obj.type == ObjectType.CAR:
rv.n_robots += 1
- elif obj.type == ObjectType.Base:
+ elif obj.type == ObjectType.BASE:
rv.n_bases += 1
- elif obj.type == ObjectType.Watcher:
+ elif obj.type == ObjectType.WATCHER:
rv.n_watchers += 1
elif isinstance(obj, Armor):
rv.armors_color2num2count[obj.color.name.lower()][obj.number] += 1
diff --git a/polystar_cv/research/common/datasets/roco/zoo/roco_dataset_zoo.py b/polystar_cv/research/common/datasets/roco/zoo/roco_dataset_zoo.py
index 12a284c1b977b14641abc4dea10a7f5f96218222..d9e546ec07c698f64e6980ea105d62e714bcd89e 100644
--- a/polystar_cv/research/common/datasets/roco/zoo/roco_dataset_zoo.py
+++ b/polystar_cv/research/common/datasets/roco/zoo/roco_dataset_zoo.py
@@ -11,6 +11,11 @@ class ROCODatasetsZoo(Iterable[Type[ROCODatasets]]):
DJI = DJIROCODatasets
TWITCH = TwitchROCODatasets
+ DEFAULT_TEST_DATASETS = [TWITCH.T470152838, TWITCH.T470151286]
+ DEFAULT_VALIDATION_DATASETS = [TWITCH.T470152289, TWITCH.T470149568]
+ TWITCH_TRAIN_DATASETS = [TWITCH.T470150052, TWITCH.T470152730, TWITCH.T470153081, TWITCH.T470158483]
+ DEFAULT_TRAIN_DATASETS = TWITCH_TRAIN_DATASETS + [DJI.FINAL, DJI.CENTRAL_CHINA, DJI.NORTH_CHINA, DJI.SOUTH_CHINA]
+
def __iter__(self) -> Iterator[Type[ROCODatasets]]:
return iter((self.DJI, self.DJI_ZOOMED, self.TWITCH))
diff --git a/polystar_cv/research/common/datasets/shuffle_dataset.py b/polystar_cv/research/common/datasets/shuffle_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..ad28020701a07907d64a2ebc48a74dfa73c1ccf9
--- /dev/null
+++ b/polystar_cv/research/common/datasets/shuffle_dataset.py
@@ -0,0 +1,15 @@
+from random import shuffle
+from typing import Iterator, Tuple
+
+from research.common.datasets.lazy_dataset import ExampleT, LazyDataset, TargetT
+
+
+class ShuffleDataset(LazyDataset):
+ def __init__(self, source: LazyDataset[ExampleT, TargetT]):
+ super().__init__(source.name)
+ self.source = source
+
+ def __iter__(self) -> Iterator[Tuple[ExampleT, TargetT]]:
+ data = list(self.source)
+ shuffle(data)
+ return iter(data)
diff --git a/polystar_cv/research/common/datasets/slice_dataset.py b/polystar_cv/research/common/datasets/slice_dataset.py
index ec06e46f0d5abbdfd7e7a94c18c16903f0f48989..348d8be3d3098120367a20634828971d35106e36 100644
--- a/polystar_cv/research/common/datasets/slice_dataset.py
+++ b/polystar_cv/research/common/datasets/slice_dataset.py
@@ -4,7 +4,7 @@ from typing import Iterator, Optional, Tuple
from research.common.datasets.lazy_dataset import ExampleT, LazyDataset, TargetT
-class SliceDataset(LazyDataset):
+class SliceDataset(LazyDataset[ExampleT, TargetT]):
def __init__(
self,
source: LazyDataset[ExampleT, TargetT],
diff --git a/polystar_cv/research/common/utils/logs.py b/polystar_cv/research/common/utils/logs.py
new file mode 100644
index 0000000000000000000000000000000000000000..0b580c6220344413e7e482bbbdb2ae0e477d1eba
--- /dev/null
+++ b/polystar_cv/research/common/utils/logs.py
@@ -0,0 +1,6 @@
+import logging
+
+
+def setup_dev_logs():
+ logging.getLogger().setLevel("INFO")
+ logging.info("logs started")
diff --git a/polystar_cv/research/robots/armor_color/datasets.py b/polystar_cv/research/robots/armor_color/datasets.py
index 5c53b9ef1e4cfc6ddc93ba8c53ebf07a97194e97..026aa425688392b2115ac5c57def0ae395d795bd 100644
--- a/polystar_cv/research/robots/armor_color/datasets.py
+++ b/polystar_cv/research/robots/armor_color/datasets.py
@@ -1,12 +1,20 @@
from typing import List, Tuple
+from polystar.common.models.image import FileImage
from polystar.common.models.object import Armor, ArmorColor
-from research.common.datasets.image_dataset import FileImageDataset
-from research.common.datasets.roco.zoo.roco_dataset_zoo import ROCODatasetsZoo
+from research.common.datasets.dataset import Dataset
from research.robots.dataset.armor_value_dataset_generator import ArmorValueDatasetGenerator
from research.robots.dataset.armor_value_target_factory import ArmorValueTargetFactory
+def make_armor_color_datasets(
+ include_dji: bool = True,
+) -> Tuple[
+ List[Dataset[FileImage, ArmorColor]], List[Dataset[FileImage, ArmorColor]], List[Dataset[FileImage, ArmorColor]]
+]:
+ return make_armor_color_dataset_generator().default_datasets(include_dji)
+
+
class ArmorColorTargetFactory(ArmorValueTargetFactory[ArmorColor]):
def from_str(self, label: str) -> ArmorColor:
return ArmorColor(label)
@@ -17,33 +25,3 @@ class ArmorColorTargetFactory(ArmorValueTargetFactory[ArmorColor]):
def make_armor_color_dataset_generator() -> ArmorValueDatasetGenerator[ArmorColor]:
return ArmorValueDatasetGenerator("colors", ArmorColorTargetFactory())
-
-
-def make_armor_color_datasets(
- include_dji: bool = True,
-) -> Tuple[List[FileImageDataset], List[FileImageDataset], List[FileImageDataset]]:
- color_dataset_generator = make_armor_color_dataset_generator()
-
- train_roco_datasets = [
- ROCODatasetsZoo.TWITCH.T470150052,
- ROCODatasetsZoo.TWITCH.T470152730,
- ROCODatasetsZoo.TWITCH.T470153081,
- ROCODatasetsZoo.TWITCH.T470158483,
- ]
- if include_dji:
- train_roco_datasets.extend(
- [
- ROCODatasetsZoo.DJI.FINAL,
- ROCODatasetsZoo.DJI.CENTRAL_CHINA,
- ROCODatasetsZoo.DJI.NORTH_CHINA,
- ROCODatasetsZoo.DJI.SOUTH_CHINA,
- ]
- )
-
- train_datasets, validation_datasets, test_datasets = color_dataset_generator.from_roco_datasets(
- train_roco_datasets,
- [ROCODatasetsZoo.TWITCH.T470149568, ROCODatasetsZoo.TWITCH.T470152289],
- [ROCODatasetsZoo.TWITCH.T470152838, ROCODatasetsZoo.TWITCH.T470151286],
- )
-
- return train_datasets, validation_datasets, test_datasets
diff --git a/polystar_cv/research/robots/armor_digit/armor_digit_dataset.py b/polystar_cv/research/robots/armor_digit/armor_digit_dataset.py
index f7f743707185a711e00379a4b7e529adb4aa18cd..7ff86d49e46e87937d115f8e87aaa7c91a4e0522 100644
--- a/polystar_cv/research/robots/armor_digit/armor_digit_dataset.py
+++ b/polystar_cv/research/robots/armor_digit/armor_digit_dataset.py
@@ -2,8 +2,9 @@ from itertools import islice
from typing import List, Set, Tuple
from polystar.common.filters.exclude_filter import ExcludeFilter
+from polystar.common.models.image import FileImage
from polystar.common.models.object import Armor, ArmorDigit
-from research.common.datasets.image_dataset import FileImageDataset
+from research.common.datasets.dataset import Dataset
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
@@ -11,34 +12,14 @@ from research.robots.dataset.armor_value_target_factory import ArmorValueTargetF
VALID_NUMBERS_2021: Set[int] = {1, 3, 4} # University League
-def make_armor_digit_dataset_generator() -> ArmorValueDatasetGenerator[ArmorDigit]:
- return ArmorValueDatasetGenerator("digits", ArmorDigitTargetFactory(), ExcludeFilter({ArmorDigit.OUTDATED}))
+def default_armor_digit_datasets() -> Tuple[
+ List[Dataset[FileImage, ArmorDigit]], List[Dataset[FileImage, ArmorDigit]], List[Dataset[FileImage, ArmorDigit]]
+]:
+ return make_armor_digit_dataset_generator().default_datasets()
-def default_armor_digit_datasets() -> Tuple[List[FileImageDataset], List[FileImageDataset], List[FileImageDataset]]:
- digit_dataset_generator = make_armor_digit_dataset_generator()
- train_datasets, validation_datasets, test_datasets = digit_dataset_generator.from_roco_datasets(
- [
- ROCODatasetsZoo.TWITCH.T470150052,
- ROCODatasetsZoo.TWITCH.T470152730,
- ROCODatasetsZoo.TWITCH.T470153081,
- ROCODatasetsZoo.TWITCH.T470158483,
- ROCODatasetsZoo.DJI.FINAL,
- ROCODatasetsZoo.DJI.CENTRAL_CHINA,
- ROCODatasetsZoo.DJI.NORTH_CHINA,
- ROCODatasetsZoo.DJI.SOUTH_CHINA,
- ],
- [ROCODatasetsZoo.TWITCH.T470149568, ROCODatasetsZoo.TWITCH.T470152289],
- [ROCODatasetsZoo.TWITCH.T470152838, ROCODatasetsZoo.TWITCH.T470151286],
- )
- # train_datasets.append(
- # digit_dataset_generator.from_roco_dataset(ROCODatasetsZoo.DJI.FINAL).to_file_images()
- # # .cap(2133 + 1764 + 1436)
- # # .skip(2133 + 176 + 1436)
- # # .cap(5_000)
- # .build()
- # )
- return train_datasets, validation_datasets, test_datasets
+def make_armor_digit_dataset_generator() -> ArmorValueDatasetGenerator[ArmorDigit]:
+ return ArmorValueDatasetGenerator("digits", ArmorDigitTargetFactory(), ExcludeFilter({ArmorDigit.OUTDATED}))
class ArmorDigitTargetFactory(ArmorValueTargetFactory[ArmorDigit]):
diff --git a/polystar_cv/research/robots/armor_digit/scripts/benchmark.py b/polystar_cv/research/robots/armor_digit/scripts/benchmark.py
index 2dde34dbd1f6dd7fd69493ad98cd86361cd41bae..9b349f0164811e1b28e9dfb5a914ba341cc2c0d8 100644
--- a/polystar_cv/research/robots/armor_digit/scripts/benchmark.py
+++ b/polystar_cv/research/robots/armor_digit/scripts/benchmark.py
@@ -2,7 +2,6 @@ import logging
import warnings
from pathlib import Path
-from polystar.common.constants import PROJECT_DIR
from polystar.common.pipeline.classification.random_model import RandomClassifier
from polystar.common.utils.serialization import pkl_load
from research.common.utils.experiment_dir import prompt_experiment_dir
@@ -34,9 +33,7 @@ if __name__ == "__main__":
# input_size=32, logs_dir=_report_dir, dropout=0, lr=0.00021, dense_size=64, model_factory=VGG16
# )
- _vgg16_pipeline = pkl_load(
- PROJECT_DIR / "pipelines/armor-digit/20201225_131957_vgg16/VGG16 (32) - lr 2.1e-04 - drop 0.pkl"
- )
+ _vgg16_pipeline = pkl_load(PIPELINES_DIR / "armor-digit/20201225_131957_vgg16/VGG16 (32) - lr 2.1e-04 - drop 0.pkl")
_vgg16_pipeline.name = "vgg16_tl"
_distiled_vgg16_into_cnn_pipeline = ArmorDigitKerasPipeline.from_distillation(
diff --git a/polystar_cv/research/robots/armor_digit/scripts/evaluate.py b/polystar_cv/research/robots/armor_digit/scripts/evaluate.py
new file mode 100644
index 0000000000000000000000000000000000000000..39c813414ea7f55250fa009a5dca301614dafa8a
--- /dev/null
+++ b/polystar_cv/research/robots/armor_digit/scripts/evaluate.py
@@ -0,0 +1,61 @@
+from pathlib import Path
+from time import time
+from typing import Iterable
+
+import seaborn
+from matplotlib.pyplot import show, title
+from pandas import DataFrame
+
+from polystar.common.models.image import FileImage, file_images_to_images
+from polystar.common.models.object import ArmorDigit
+from polystar.common.pipeline.classification.classification_pipeline import ClassificationPipeline
+from polystar.common.utils.iterable_utils import chunk
+from polystar.common.utils.serialization import pkl_load
+from research.common.constants import PIPELINES_DIR
+from research.common.datasets.dataset import Dataset
+from research.common.gcloud.gcloud_storage import GCStorages
+from research.common.utils.logs import setup_dev_logs
+from research.robots.armor_digit.armor_digit_dataset import make_armor_digit_dataset_generator
+
+
+def time_digit_pipeline(pipeline_path: Path):
+ seaborn.set()
+ GCStorages.DEV.download_file_if_missing(pipeline_path)
+ pipeline = pkl_load(pipeline_path)
+ test_datasets = make_armor_digit_dataset_generator().default_test_datasets()
+ pipeline.predict_proba_and_classes(file_images_to_images(test_datasets[0].examples[:5]))
+ df = DataFrame(
+ [
+ {"dataset": dataset.name, "time (s)": t, "batch_size": batch_size}
+ for dataset in test_datasets
+ for batch_size in range(1, 16)
+ for t in time_batch_inference(pipeline, dataset, batch_size)
+ ]
+ )
+ seaborn.violinplot(x="batch_size", y="time (s)", data=df, cut=0)
+ title(f"Time inference\n{pipeline.name}")
+ show()
+
+
+def time_batch_inference(
+ pipeline: ClassificationPipeline, dataset: Dataset[FileImage, ArmorDigit], batch_size: int
+) -> Iterable[float]:
+ for file_images in chunk(dataset.examples, batch_size):
+ images = file_images_to_images(file_images)
+ t = time()
+ pipeline.predict_proba_and_classes(images)
+ rv = time() - t
+ if rv > 0.1:
+ print(rv, *[f"file://{f.path}" for f in file_images])
+ yield rv
+
+
+if __name__ == "__main__":
+ setup_dev_logs()
+ time_digit_pipeline(
+ PIPELINES_DIR
+ / "armor-digit/20210117_145856_kd_cnn/distiled - temp 4.1e+01 - cnn - (32) - lr 7.8e-04 - drop 63.pkl"
+ )
+ time_digit_pipeline(
+ PIPELINES_DIR / "armor-digit/20210110_220816_vgg16_full_dset/wrapper (32) - lr 2.1e-04 - drop 0.pkl"
+ )
diff --git a/polystar_cv/research/robots/armor_digit/scripts/hyper_tune_distiled_vgg16_into_cnn.py b/polystar_cv/research/robots/armor_digit/scripts/hyper_tune_distiled_vgg16_into_cnn.py
index 3b035e72dad631fa6437c8e269c0c1cfe6195136..2994f9e3f1153ae3b67ec3b587d2444f0154c19f 100644
--- a/polystar_cv/research/robots/armor_digit/scripts/hyper_tune_distiled_vgg16_into_cnn.py
+++ b/polystar_cv/research/robots/armor_digit/scripts/hyper_tune_distiled_vgg16_into_cnn.py
@@ -4,7 +4,6 @@ from pathlib import Path
from optuna import Trial
-from polystar.common.constants import PROJECT_DIR
from polystar.common.utils.serialization import pkl_load
from research.common.utils.experiment_dir import make_experiment_dir
from research.robots.armor_digit.digit_benchmarker import make_default_digit_benchmarker
@@ -14,7 +13,7 @@ from research.robots.evaluation.hyper_tuner import HyperTuner
class DistilledPipelineFactory:
def __init__(self, teacher_name: str):
- self.teacher: ArmorDigitKerasPipeline = pkl_load(PROJECT_DIR / "pipelines/armor-digit" / teacher_name)
+ self.teacher: ArmorDigitKerasPipeline = pkl_load(PIPELINES_DIR / "armor-digit" / teacher_name)
def __call__(self, report_dir: Path, trial: Trial) -> ArmorDigitPipeline:
return ArmorDigitKerasPipeline.from_distillation(
diff --git a/polystar_cv/research/robots/armor_digit/scripts/train_kd_cnn.py b/polystar_cv/research/robots/armor_digit/scripts/train_kd_cnn.py
new file mode 100644
index 0000000000000000000000000000000000000000..13e6b17ccd16642f86a815181bc3c37b08832335
--- /dev/null
+++ b/polystar_cv/research/robots/armor_digit/scripts/train_kd_cnn.py
@@ -0,0 +1,26 @@
+from polystar.common.utils.serialization import pkl_load
+from polystar.common.utils.time import create_time_id
+from research.common.constants import PIPELINES_DIR
+from research.common.utils.logs import setup_dev_logs
+from research.robots.armor_digit.pipeline import ArmorDigitKerasPipeline
+from research.robots.armor_digit.training import train_report_and_upload_digit_pipeline
+
+if __name__ == "__main__":
+ setup_dev_logs()
+
+ _training_dir = PIPELINES_DIR / "armor-digit" / f"{create_time_id()}_kd_cnn"
+
+ _kd_cnn_pipeline = ArmorDigitKerasPipeline.from_distillation(
+ teacher_pipeline=pkl_load(
+ PIPELINES_DIR / "armor-digit/20210110_220816_vgg16_full_dset/wrapper (32) - lr 2.1e-04 - drop 0.pkl"
+ ),
+ conv_blocks=((32, 32), (64, 64)),
+ logs_dir=str(_training_dir),
+ dropout=0.63,
+ lr=0.000776,
+ dense_size=1024,
+ temperature=41.2,
+ verbose=1,
+ )
+
+ train_report_and_upload_digit_pipeline(_kd_cnn_pipeline, training_dir=_training_dir)
diff --git a/polystar_cv/research/robots/armor_digit/scripts/train_vgg16.py b/polystar_cv/research/robots/armor_digit/scripts/train_vgg16.py
index 650aabcf1f444a83ec2f79483c14bd1d7c8d0697..9c54d412a2a3c308676fb6775a2f844cd50274a8 100644
--- a/polystar_cv/research/robots/armor_digit/scripts/train_vgg16.py
+++ b/polystar_cv/research/robots/armor_digit/scripts/train_vgg16.py
@@ -1,21 +1,13 @@
-import logging
-import warnings
-
from tensorflow.python.keras.applications.vgg16 import VGG16
-from polystar.common.constants import PROJECT_DIR
-from polystar.common.utils.serialization import pkl_dump
from polystar.common.utils.time import create_time_id
-from research.robots.armor_digit.digit_benchmarker import make_default_digit_benchmarker
+from research.common.constants import PIPELINES_DIR
+from research.common.utils.logs import setup_dev_logs
from research.robots.armor_digit.pipeline import ArmorDigitKerasPipeline
-
-PIPELINES_DIR = PROJECT_DIR / "pipelines"
+from research.robots.armor_digit.training import train_report_and_upload_digit_pipeline
if __name__ == "__main__":
- logging.getLogger().setLevel("INFO")
- logging.getLogger("tensorflow").setLevel("ERROR")
- warnings.filterwarnings("ignore")
- logging.info("Training vgg16")
+ setup_dev_logs()
_training_dir = PIPELINES_DIR / "armor-digit" / f"{create_time_id()}_vgg16_full_dset"
@@ -29,9 +21,4 @@ if __name__ == "__main__":
verbose=1,
)
- logging.info(f"Run `tensorboard --logdir={_training_dir}` for realtime logs")
-
- _benchmarker = make_default_digit_benchmarker(_training_dir)
- _benchmarker.benchmark([_vgg16_pipeline])
-
- pkl_dump(_vgg16_pipeline, _training_dir / _vgg16_pipeline.name)
+ train_report_and_upload_digit_pipeline(_vgg16_pipeline, _training_dir)
diff --git a/polystar_cv/research/robots/armor_digit/training.py b/polystar_cv/research/robots/armor_digit/training.py
new file mode 100644
index 0000000000000000000000000000000000000000..c14f5f4a851ba2a899921011e070d40065b26de2
--- /dev/null
+++ b/polystar_cv/research/robots/armor_digit/training.py
@@ -0,0 +1,13 @@
+import pickle
+from pathlib import Path
+
+from research.common.gcloud.gcloud_storage import GCStorages
+from research.robots.armor_digit.digit_benchmarker import make_default_digit_benchmarker
+from research.robots.armor_digit.pipeline import ArmorDigitPipeline
+
+
+def train_report_and_upload_digit_pipeline(pipeline: ArmorDigitPipeline, training_dir: Path):
+ make_default_digit_benchmarker(training_dir).benchmark([pipeline])
+
+ with GCStorages.DEV.open((training_dir / pipeline.name).with_suffix(".pkl"), "wb") as f:
+ pickle.dump(pipeline, f)
diff --git a/polystar_cv/research/robots/dataset/armor_dataset_factory.py b/polystar_cv/research/robots/dataset/armor_dataset_factory.py
index 6a800aefa3a2f2342875b83418756c64473a4170..c14eca1d40b5837ffbb897fd24139ee3794088e7 100644
--- a/polystar_cv/research/robots/dataset/armor_dataset_factory.py
+++ b/polystar_cv/research/robots/dataset/armor_dataset_factory.py
@@ -23,7 +23,7 @@ class ArmorDataset(LazyDataset[Image, Armor]):
@staticmethod
def _generate_from_single(image: Image, annotation: ROCOAnnotation, name) -> Iterator[Tuple[Image, Armor, str]]:
- armors: List[Armor] = TypeObjectValidator(ObjectType.Armor).filter(annotation.objects, image)
+ armors: List[Armor] = TypeObjectValidator(ObjectType.ARMOR).filter(annotation.objects, image)
for i, obj in enumerate(armors):
croped_img = image[obj.box.y1 : obj.box.y2, obj.box.x1 : obj.box.x2]
diff --git a/polystar_cv/research/robots/dataset/armor_value_dataset_generator.py b/polystar_cv/research/robots/dataset/armor_value_dataset_generator.py
index a296d9c814af80914d5ba8e8ddc787221559fb41..f90fd6c4c32679b883034338931f3cd233cb49f7 100644
--- a/polystar_cv/research/robots/dataset/armor_value_dataset_generator.py
+++ b/polystar_cv/research/robots/dataset/armor_value_dataset_generator.py
@@ -1,14 +1,16 @@
from pathlib import Path
-from typing import Generic, Iterable, List
+from typing import Generic, Iterable, List, Tuple
from polystar.common.filters.exclude_filter import ExcludeFilter
from polystar.common.filters.filter_abc import FilterABC
from polystar.common.filters.pass_through_filter import PassThroughFilter
+from polystar.common.models.image import FileImage
from research.common.dataset.cleaning.dataset_changes import DatasetChanges
-from research.common.datasets.image_dataset import FileImageDataset
+from research.common.datasets.dataset import Dataset
from research.common.datasets.image_file_dataset_builder import DirectoryDatasetBuilder
from research.common.datasets.lazy_dataset import TargetT
from research.common.datasets.roco.roco_dataset_builder import ROCODatasetBuilder
+from research.common.datasets.roco.zoo.roco_dataset_zoo import ROCODatasetsZoo
from research.robots.dataset.armor_value_dataset_cache import ArmorValueDatasetCache
from research.robots.dataset.armor_value_target_factory import ArmorValueTargetFactory
@@ -29,15 +31,31 @@ class ArmorValueDatasetGenerator(Generic[TargetT]):
self.task_name = task_name
self.targets_filter = targets_filter or PassThroughFilter()
- # FIXME signature inconsistency across methods
- def from_roco_datasets(
- self, *roco_datasets_list: List[ROCODatasetBuilder]
- ) -> Iterable[List[FileImageDataset[TargetT]]]:
+ def default_datasets(
+ self, include_dji: bool = True
+ ) -> Tuple[List[Dataset[FileImage, TargetT]], List[Dataset[FileImage, TargetT]], List[Dataset[FileImage, TargetT]]]:
return (
- [self.from_roco_dataset(roco_dataset).to_file_images().build() for roco_dataset in roco_datasets]
- for roco_datasets in roco_datasets_list
+ self.default_train_datasets() if include_dji else self.twitch_train_datasets(),
+ self.default_validation_datasets(),
+ self.default_test_datasets(),
)
+ def default_test_datasets(self) -> List[Dataset[FileImage, TargetT]]:
+ return self.from_roco_datasets(ROCODatasetsZoo.DEFAULT_TEST_DATASETS)
+
+ def default_validation_datasets(self) -> List[Dataset[FileImage, TargetT]]:
+ return self.from_roco_datasets(ROCODatasetsZoo.DEFAULT_VALIDATION_DATASETS)
+
+ def default_train_datasets(self) -> List[Dataset[FileImage, TargetT]]:
+ return self.from_roco_datasets(ROCODatasetsZoo.DEFAULT_TRAIN_DATASETS)
+
+ def twitch_train_datasets(self) -> List[Dataset[FileImage, TargetT]]:
+ return self.from_roco_datasets(ROCODatasetsZoo.TWITCH_TRAIN_DATASETS)
+
+ # FIXME signature inconsistency across methods
+ def from_roco_datasets(self, roco_datasets: Iterable[ROCODatasetBuilder]) -> List[Dataset[FileImage, TargetT]]:
+ return [self.from_roco_dataset(roco_dataset).to_file_images().build() for roco_dataset in roco_datasets]
+
def from_roco_dataset(self, roco_dataset_builder: ROCODatasetBuilder) -> DirectoryDatasetBuilder[TargetT]:
cache_dir = roco_dataset_builder.main_dir / self.task_name
dataset_name = roco_dataset_builder.name
diff --git a/polystar_cv/research/robots/demos/demo_infer.py b/polystar_cv/research/robots/demos/demo_infer.py
deleted file mode 100644
index 51c1cc2258d9ec00b2e680ff2dd01dbb618f2ae0..0000000000000000000000000000000000000000
--- a/polystar_cv/research/robots/demos/demo_infer.py
+++ /dev/null
@@ -1,24 +0,0 @@
-from polystar.common.dependency_injection import make_injector
-from polystar.common.models.label_map import LabelMap
-from polystar.common.target_pipeline.detected_objects.detected_objects_factory import DetectedObjectFactory
-from polystar.common.target_pipeline.objects_detectors.tf_model_objects_detector import TFModelObjectsDetector
-from polystar.common.target_pipeline.objects_validators.confidence_object_validator import ConfidenceObjectValidator
-from polystar.common.utils.tensorflow import patch_tf_v2
-from polystar.common.view.plt_results_viewer import PltResultViewer
-from research.common.datasets.roco.zoo.roco_dataset_zoo import ROCODatasetsZoo
-from research.robots.demos.utils import load_tf_model
-
-if __name__ == "__main__":
- patch_tf_v2()
- injector = make_injector()
-
- objects_detector = TFModelObjectsDetector(DetectedObjectFactory(injector.get(LabelMap), []), load_tf_model())
- filters = [ConfidenceObjectValidator(confidence_threshold=0.5)]
-
- with PltResultViewer("Demo of tf model") as viewer:
- for image, _, _ in ROCODatasetsZoo.DJI.CENTRAL_CHINA.to_images().cap(5):
- objects = objects_detector.detect(image)
- for f in filters:
- objects = f.filter(objects, image)
-
- viewer.display_image_with_objects(image, objects)
diff --git a/polystar_cv/research/robots/demos/demo_pipeline.py b/polystar_cv/research/robots/demos/demo_pipeline.py
index ff5475b56724deabb5acb4d468ae92f763336e97..44fbbbebfd3460ac59408d7b92be95f994711692 100644
--- a/polystar_cv/research/robots/demos/demo_pipeline.py
+++ b/polystar_cv/research/robots/demos/demo_pipeline.py
@@ -1,51 +1,33 @@
-import cv2
+from injector import inject
-from polystar.common.communication.print_target_sender import PrintTargetSender
from polystar.common.dependency_injection import make_injector
-from polystar.common.models.camera import Camera
-from polystar.common.models.label_map import LabelMap
-from polystar.common.target_pipeline.armors_descriptors.armors_color_descriptor import ArmorsColorDescriptor
from polystar.common.target_pipeline.debug_pipeline import DebugTargetPipeline
-from polystar.common.target_pipeline.detected_objects.detected_objects_factory import DetectedObjectFactory
-from polystar.common.target_pipeline.object_selectors.closest_object_selector import ClosestObjectSelector
-from polystar.common.target_pipeline.objects_detectors.tf_model_objects_detector import TFModelObjectsDetector
-from polystar.common.target_pipeline.objects_linker.simple_objects_linker import SimpleObjectsLinker
-from polystar.common.target_pipeline.objects_validators.confidence_object_validator import ConfidenceObjectValidator
-from polystar.common.target_pipeline.target_factories.ratio_simple_target_factory import RatioSimpleTargetFactory
+from polystar.common.target_pipeline.objects_validators.armor_digit_validator import ArmorDigitValidator
from polystar.common.target_pipeline.target_pipeline import NoTargetFoundException
-from polystar.common.utils.tensorflow import patch_tf_v2
from polystar.common.view.plt_results_viewer import PltResultViewer
+from research.common.datasets.roco.roco_annotation_filters.roco_annotation_object_filter import (
+ ROCOAnnotationObjectFilter,
+)
from research.common.datasets.roco.zoo.roco_dataset_zoo import ROCODatasetsZoo
-from research.robots.armor_color.pipeline import ArmorColorPipeline
-from research.robots.armor_color.scripts.benchmark import MeanChannels, RedBlueComparisonClassifier
-from research.robots.demos.utils import load_tf_model
-if __name__ == "__main__":
- patch_tf_v2()
- injector = make_injector()
-
- pipeline = DebugTargetPipeline(
- objects_detector=TFModelObjectsDetector(
- DetectedObjectFactory(
- injector.get(LabelMap),
- [ArmorsColorDescriptor(ArmorColorPipeline.from_pipes([MeanChannels(), RedBlueComparisonClassifier()]))],
- ),
- load_tf_model(),
- ),
- objects_validators=[ConfidenceObjectValidator(0.6)],
- object_selector=ClosestObjectSelector(),
- target_factory=RatioSimpleTargetFactory(injector.get(Camera), 300, 100),
- target_sender=PrintTargetSender(),
- objects_linker=SimpleObjectsLinker(min_percentage_intersection=0.8),
- )
+@inject
+def demo_pipeline_on_images(pipeline: DebugTargetPipeline):
with PltResultViewer("Demo of tf model") as viewer:
- for builder in (ROCODatasetsZoo.TWITCH.T470150052, ROCODatasetsZoo.DJI.CENTRAL_CHINA):
- for image_path, _, _ in builder.cap(5):
+ for builder in ROCODatasetsZoo.DEFAULT_TEST_DATASETS:
+ for image in (
+ builder.to_images()
+ .filter_targets(ROCOAnnotationObjectFilter(ArmorDigitValidator((1, 3, 4))))
+ .shuffle()
+ .cap(15)
+ .build_examples()
+ ):
try:
- image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB)
- target = pipeline.predict_target(image)
+ pipeline.predict_target(image)
except NoTargetFoundException:
pass
- finally:
- viewer.display_debug_info(pipeline.debug_info_)
+ viewer.display_debug_info(pipeline.debug_info_)
+
+
+if __name__ == "__main__":
+ make_injector().call_with_injection(demo_pipeline_on_images)
diff --git a/polystar_cv/research/robots/demos/demo_pipeline_camera.py b/polystar_cv/research/robots/demos/demo_pipeline_camera.py
index 660e09e9620c79f37a6733b34c942da9b916531f..dbddb53a36b08b190e02168282228e06b1d14482 100644
--- a/polystar_cv/research/robots/demos/demo_pipeline_camera.py
+++ b/polystar_cv/research/robots/demos/demo_pipeline_camera.py
@@ -1,60 +1,31 @@
-import sys
-from time import time
+from injector import inject
-import pycuda.autoinit # This is needed for initializing CUDA driver
-
-from polystar.common.communication.file_descriptor_target_sender import FileDescriptorTargetSender
-from polystar.common.constants import MODELS_DIR
from polystar.common.dependency_injection import make_injector
-from polystar.common.frame_generators.camera_frame_generator import CameraFrameGenerator
-from polystar.common.models.camera import Camera
-from polystar.common.models.label_map import LabelMap
-from polystar.common.models.object import ObjectType
-from polystar.common.models.trt_model import TRTModel
+from polystar.common.frame_generators.frames_generator_abc import FrameGeneratorABC
from polystar.common.target_pipeline.debug_pipeline import DebugTargetPipeline
-from polystar.common.target_pipeline.object_selectors.closest_object_selector import ClosestObjectSelector
-from polystar.common.target_pipeline.objects_detectors.trt_model_object_detector import TRTModelObjectsDetector
-from polystar.common.target_pipeline.objects_validators.confidence_object_validator import ConfidenceObjectValidator
-from polystar.common.target_pipeline.objects_validators.type_object_validator import TypeObjectValidator
-from polystar.common.target_pipeline.target_factories.ratio_simple_target_factory import RatioSimpleTargetFactory
-from polystar.common.utils.tensorflow import patch_tf_v2
+from polystar.common.target_pipeline.target_pipeline import NoTargetFoundException
+from polystar.common.utils.fps import FPS
from polystar.common.view.cv2_results_viewer import CV2ResultViewer
-from polystar.robots_at_robots.globals import settings
-
-[pycuda.autoinit] # So pycharm won't remove the import
-
-if __name__ == "__main__":
- patch_tf_v2()
- injector = make_injector()
-
- objects_detector = TRTModelObjectsDetector(
- TRTModel(MODELS_DIR / settings.MODEL_NAME, (300, 300)), injector.get(LabelMap)
- )
- pipeline = DebugTargetPipeline(
- objects_detector=objects_detector,
- objects_validators=[ConfidenceObjectValidator(0.6), TypeObjectValidator(ObjectType.Armor)],
- object_selector=ClosestObjectSelector(),
- target_factory=RatioSimpleTargetFactory(injector.get(Camera), 300, 100),
- target_sender=FileDescriptorTargetSender(int(sys.argv[1])),
- )
- fps = 0
+@inject
+def demo_pipeline_on_camera(pipeline: DebugTargetPipeline, webcam: FrameGeneratorABC):
+ fps, pipeline_fps = FPS(), FPS()
with CV2ResultViewer("TensorRT demo") as viewer:
- for image in CameraFrameGenerator(1_280, 720).generate():
-
- previous_time = time()
-
- # inference
- pipeline.predict_target(image)
-
- # display
- fps = 0.9 * fps + 0.1 / (time() - previous_time)
- viewer.new(image)
- viewer.add_objects(pipeline.debug_info_.validated_objects, forced_color=(0.6, 0.6, 0.6))
- viewer.add_object(pipeline.debug_info_.selected_object)
- viewer.add_text(f"FPS: {fps:.1f}", 10, 10, (0, 0, 0))
+ for image in webcam.generate():
+ pipeline_fps.skip()
+ try:
+ pipeline.predict_target(image)
+ except NoTargetFoundException:
+ pass
+ pipeline_fps.tick(), fps.tick()
+ viewer.add_debug_info(pipeline.debug_info_)
+ viewer.add_text(f"FPS: {fps:.1f} / {pipeline_fps:.1f}", 10, 10, (0, 0, 0))
viewer.display()
-
+ fps.skip()
if viewer.finished:
- break
+ return
+
+
+if __name__ == "__main__":
+ make_injector().call_with_injection(demo_pipeline_on_camera)
diff --git a/polystar_cv/research/robots/demos/utils.py b/polystar_cv/research/robots/demos/utils.py
deleted file mode 100644
index 9b103e5e8f62e3d0619ae55e597499355bee559f..0000000000000000000000000000000000000000
--- a/polystar_cv/research/robots/demos/utils.py
+++ /dev/null
@@ -1,15 +0,0 @@
-import tensorflow as tf
-
-from polystar.common.constants import MODELS_DIR
-
-
-def load_tf_model():
- model = tf.saved_model.load(
- export_dir=str(
- MODELS_DIR
- / "robots"
- / "ssd_mobilenet_v2_roco_2018_03_29_20200314_015411_TWITCH_TEMP_733_IMGS_29595steps"
- / "saved_model"
- )
- )
- return model.signatures["serving_default"]
diff --git a/polystar_cv/research/robots/evaluation/benchmarker.py b/polystar_cv/research/robots/evaluation/benchmarker.py
index 0c3123faaefd550caf4097e28e8dbf543a7a9394..d9fc462e751c2fed6584dae2bd8f23412097f416 100644
--- a/polystar_cv/research/robots/evaluation/benchmarker.py
+++ b/polystar_cv/research/robots/evaluation/benchmarker.py
@@ -1,7 +1,8 @@
import logging
from dataclasses import dataclass
+from itertools import chain
from pathlib import Path
-from typing import List
+from typing import List, Sequence
from polystar.common.pipeline.classification.classification_pipeline import ClassificationPipeline
from research.common.datasets.image_dataset import FileImageDataset
@@ -40,10 +41,10 @@ class Benchmarker:
return pipeline_performances
def benchmark(
- self, pipelines: List[ClassificationPipeline], trained_pipelines: List[ClassificationPipeline] = None
+ self, pipelines: Sequence[ClassificationPipeline] = (), trained_pipelines: Sequence[ClassificationPipeline] = ()
):
self.trainer.train_pipelines(pipelines)
- self.performances += self.evaluator.evaluate_pipelines(pipelines + (trained_pipelines or []))
+ self.performances += self.evaluator.evaluate_pipelines(chain(pipelines, trained_pipelines))
self.make_report()
def make_report(self):
diff --git a/polystar_cv/research/robots/evaluation/evaluator.py b/polystar_cv/research/robots/evaluation/evaluator.py
index d03c1c9ec6b157c5c0c71b896d2367e6e15f4420..8bd591e7321026de8edad582ae6721d01704fe17 100644
--- a/polystar_cv/research/robots/evaluation/evaluator.py
+++ b/polystar_cv/research/robots/evaluation/evaluator.py
@@ -39,8 +39,9 @@ class ImageClassificationPipelineEvaluator(Generic[TargetT]):
self, pipeline: ClassificationPipeline, set_: Set
) -> Iterable[ContextualizedClassificationPerformance]:
for dataset in self.set2datasets[set_]:
+ images = file_images_to_images(dataset.examples)
t = time()
- proba, classes = pipeline.predict_proba_and_classes(file_images_to_images(dataset.examples))
+ proba, classes = pipeline.predict_proba_and_classes(images)
mean_time = (time() - t) / len(dataset)
yield ContextualizedClassificationPerformance(
examples=dataset.examples,
diff --git a/polystar_cv/research/robots/evaluation/trainer.py b/polystar_cv/research/robots/evaluation/trainer.py
index 6f5940bb0ab56abc2ce15c8b81f932e30971d1f2..2cc64f648c0a424149ee5a95c250dd2e5ffd555e 100644
--- a/polystar_cv/research/robots/evaluation/trainer.py
+++ b/polystar_cv/research/robots/evaluation/trainer.py
@@ -1,4 +1,4 @@
-from typing import Generic, List
+from typing import Generic, Iterable, List
from tqdm import tqdm
@@ -19,7 +19,7 @@ class ImageClassificationPipelineTrainer(Generic[TargetT]):
def train_pipeline(self, pipeline: ClassificationPipeline):
pipeline.fit(self.images, self.labels, validation_size=self.validation_size)
- def train_pipelines(self, pipelines: List[ClassificationPipeline]):
+ def train_pipelines(self, pipelines: Iterable[ClassificationPipeline]):
tqdm_pipelines = tqdm(pipelines, desc="Training Pipelines")
for pipeline in tqdm_pipelines:
tqdm_pipelines.set_postfix({"pipeline": pipeline.name}, True)
diff --git a/pyproject.toml b/pyproject.toml
index 82b631031aeb07f1dfbac085901ba6fa997fef3e..48c1e6876b7c863e42f3a90e4cdad54218ff33be 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,8 +35,9 @@ pyyaml = "^5.3.1"
six = "1.15.0" # https://github.com/googleapis/python-bigquery/issues/70
[tool.poetry.dev-dependencies]
-tensorflow = "2.3.x"
-tensorflow-estimator = "2.3.x"
+tensorflow = "1.14.x"
+tensorflow-estimator = "1.14.x"
+h5py = "<3.0.0"
kivy = "^1.11.1"
cloudml-hypertune = "^0.1.0-alpha.6"
google-api-python-client = "^1.12.8"