From 4dc6ead115a4244b213655584ee3066e8e6c5c9f Mon Sep 17 00:00:00 2001
From: Mathieu Beligon <beligonmathieu@gmail.com>
Date: Sat, 20 Feb 2021 20:19:56 -0500
Subject: [PATCH] [armor-color] move the featurizers into code
---
src/polystar/dependency_injection.py | 3 ++-
.../cv2_frame_generator_abc.py | 1 +
.../red_blue_comparison_classifier.py | 11 ++++++++++
.../pipeline/featurizers/mean_channels.py | 11 ++++++++++
.../armors/armor_color/scripts/benchmark.py | 22 ++-----------------
5 files changed, 27 insertions(+), 21 deletions(-)
create mode 100644 src/polystar/pipeline/classification/red_blue_comparison_classifier.py
create mode 100644 src/polystar/pipeline/featurizers/mean_channels.py
diff --git a/src/polystar/dependency_injection.py b/src/polystar/dependency_injection.py
index cc9807f..0f5542f 100644
--- a/src/polystar/dependency_injection.py
+++ b/src/polystar/dependency_injection.py
@@ -27,7 +27,8 @@ from polystar.target_pipeline.target_factories.ratio_simple_target_factory impor
from polystar.target_pipeline.target_factories.target_factory_abc import TargetFactoryABC
from polystar.utils.serialization import pkl_load
from research.armors.armor_color.pipeline import ArmorColorPipeline
-from research.armors.armor_color.scripts.benchmark import MeanChannels, RedBlueComparisonClassifier
+from polystar.pipeline.classification.red_blue_comparison_classifier import RedBlueComparisonClassifier
+from polystar.pipeline.featurizers.mean_channels import MeanChannels
from research.common.constants import PIPELINES_DIR
diff --git a/src/polystar/frame_generators/cv2_frame_generator_abc.py b/src/polystar/frame_generators/cv2_frame_generator_abc.py
index 4a6c433..6f60339 100644
--- a/src/polystar/frame_generators/cv2_frame_generator_abc.py
+++ b/src/polystar/frame_generators/cv2_frame_generator_abc.py
@@ -15,6 +15,7 @@ class CV2FrameGeneratorABC(FrameGeneratorABC, ABC):
def __enter__(self):
self._cap = cv2.VideoCapture(*self._capture_params())
+ # self._cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
assert self._cap.isOpened()
self._post_opening_operation()
diff --git a/src/polystar/pipeline/classification/red_blue_comparison_classifier.py b/src/polystar/pipeline/classification/red_blue_comparison_classifier.py
new file mode 100644
index 0000000..69c0bb3
--- /dev/null
+++ b/src/polystar/pipeline/classification/red_blue_comparison_classifier.py
@@ -0,0 +1,11 @@
+from nptyping import Array
+
+from polystar.models.roco_object import ArmorColor
+from polystar.pipeline.classification.rule_based_classifier import RuleBasedClassifierABC
+
+
+class RedBlueComparisonClassifier(RuleBasedClassifierABC):
+ """A very simple model that compares the blue and red values obtained by the MeanChannels"""
+
+ def predict_single(self, features: Array[float, float, float]) -> ArmorColor:
+ return ArmorColor.RED if features[0] >= features[2] else ArmorColor.BLUE
\ No newline at end of file
diff --git a/src/polystar/pipeline/featurizers/mean_channels.py b/src/polystar/pipeline/featurizers/mean_channels.py
new file mode 100644
index 0000000..169bead
--- /dev/null
+++ b/src/polystar/pipeline/featurizers/mean_channels.py
@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+from nptyping import Array
+
+from polystar.models.image import Image
+from polystar.pipeline.pipe_abc import PipeABC
+
+
+@dataclass
+class MeanChannels(PipeABC):
+ def transform_single(self, image: Image) -> Array[float, float, float]:
+ return image.mean(axis=(0, 1))
\ No newline at end of file
diff --git a/src/research/armors/armor_color/scripts/benchmark.py b/src/research/armors/armor_color/scripts/benchmark.py
index 4862070..d116c86 100644
--- a/src/research/armors/armor_color/scripts/benchmark.py
+++ b/src/research/armors/armor_color/scripts/benchmark.py
@@ -1,36 +1,18 @@
import logging
-from dataclasses import dataclass
from pathlib import Path
-from nptyping import Array
from sklearn.linear_model import LogisticRegression
-from polystar.models.image import Image
-from polystar.models.roco_object import ArmorColor
from polystar.pipeline.classification.random_model import RandomClassifier
-from polystar.pipeline.classification.rule_based_classifier import RuleBasedClassifierABC
+from polystar.pipeline.classification.red_blue_comparison_classifier import RedBlueComparisonClassifier
from polystar.pipeline.featurizers.histogram_2d import Histogram2D
from polystar.pipeline.featurizers.histogram_blocs_2d import HistogramBlocs2D
-from polystar.pipeline.pipe_abc import PipeABC
+from polystar.pipeline.featurizers.mean_channels import MeanChannels
from polystar.pipeline.preprocessors.rgb_to_hsv import RGB2HSV
from research.armors.armor_color.benchmarker import make_armor_color_benchmarker
from research.armors.armor_color.pipeline import ArmorColorPipeline
from research.common.utils.experiment_dir import prompt_experiment_dir
-
-@dataclass
-class MeanChannels(PipeABC):
- def transform_single(self, image: Image) -> Array[float, float, float]:
- return image.mean(axis=(0, 1))
-
-
-class RedBlueComparisonClassifier(RuleBasedClassifierABC):
- """A very simple model that compares the blue and red values obtained by the MeanChannels"""
-
- def predict_single(self, features: Array[float, float, float]) -> ArmorColor:
- return ArmorColor.RED if features[0] >= features[2] else ArmorColor.BLUE
-
-
if __name__ == "__main__":
logging.getLogger().setLevel("INFO")
logging.info("Benchmarking")
--
GitLab