diff --git a/src/polystar/dependency_injection.py b/src/polystar/dependency_injection.py index cc9807f151497cd3747a3e2a7c7a011d66ccf388..0f5542f29cd2d83a3e8d6a9f674cc1ee06beb0a3 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 4a6c4337111622f6ae3d3b83d959a65e97097664..6f6033988c5a81a6b0cd742be391cf7ece64e8f5 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 0000000000000000000000000000000000000000..69c0bb3b64288ebc5901053b55b20f0b858aa406 --- /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 0000000000000000000000000000000000000000..169bead98eba1c2d29424b36b4caf3399c662b00 --- /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 4862070b1472b5b676e3ae1afaf73cd50d30b524..d116c8621f0d568beaa5045f02f7c2e698eab754 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")