Vision refactoring (#720)

Adds refactored vision

Author: cleWu03

bitbots_vision/bitbots_vision/vision.py

@@ -1,7 +1,7 @@
 #! /usr/bin/env python3
 from copy import deepcopy
+from typing import Optional
 
-import numpy as np
 import rclpy
 from ament_index_python.packages import get_package_share_directory
 from cv_bridge import CvBridge
@@ -10,7 +10,7 @@
 from rclpy.node import Node
 from sensor_msgs.msg import Image
 
-from bitbots_vision.vision_modules import ros_utils, yoeo
+from bitbots_vision.vision_modules import debug, ros_utils, yoeo
 
 from .params import gen
 
@@ -48,7 +48,8 @@ def __init__(self) -> None:
 
         self._sub_image = None
 
-        self._vision_components: list[yoeo.IVisionComponent] = []
+        self._vision_components: list[yoeo.AbstractVisionComponent] = []
+        self._debug_image: Optional[debug.DebugImage] = None
 
         # Setup reconfiguration
         gen.declare_params(self)
@@ -85,37 +86,44 @@ def _get_updated_config_with(self, params) -> dict:
     def _configure_vision(self, new_config: dict) -> None:
         yoeo.YOEOObjectManager.configure(new_config)
 
-        self._set_up_vision_components(new_config)
-        self._register_subscribers(new_config)
+        debug_image = debug.DebugImage(new_config["component_debug_image_active"])
+        self._debug_image = debug_image
 
-    def _set_up_vision_components(self, new_config: dict) -> None:
-        self._vision_components = []
-        self._vision_components.append(yoeo.YOEOComponent())
+        def make_vision_component(
+            component_class: type[yoeo.AbstractVisionComponent], **kwargs
+        ) -> yoeo.AbstractVisionComponent:
+            return component_class(
+                node=self,
+                yoeo_handler=yoeo.YOEOObjectManager.get(),
+                debug_image=debug_image,
+                config=new_config,
+                **kwargs,
+            )
+
+        self._vision_components = [make_vision_component(yoeo.YOEOComponent)]
 
         if new_config["component_ball_detection_active"]:
-            self._vision_components.append(yoeo.BallDetectionComponent(self))
+            self._vision_components.append(make_vision_component(yoeo.BallDetectionComponent))
         if new_config["component_robot_detection_active"]:
-            if yoeo.YOEOObjectManager.is_team_color_detection_supported():
-                self._vision_components.append(yoeo.RobotDetectionComponent(self))
-            else:
-                self._vision_components.append(yoeo.NoTeamColorRobotDetectionComponent(self))
+            self._vision_components.append(
+                make_vision_component(
+                    yoeo.RobotDetectionComponent,
+                    team_color_detection_supported=yoeo.YOEOObjectManager.is_team_color_detection_supported(),
+                )
+            )
         if new_config["component_goalpost_detection_active"]:
-            self._vision_components.append(yoeo.GoalpostDetectionComponent(self))
+            self._vision_components.append(make_vision_component(yoeo.GoalpostDetectionComponent))
         if new_config["component_line_detection_active"]:
-            self._vision_components.append(yoeo.LineDetectionComponent(self))
+            self._vision_components.append(make_vision_component(yoeo.LineDetectionComponent))
         if new_config["component_field_detection_active"]:
-            self._vision_components.append(yoeo.FieldDetectionComponent(self))
+            self._vision_components.append(make_vision_component(yoeo.FieldDetectionComponent))
         if new_config["component_debug_image_active"]:
-            self._vision_components.append(yoeo.DebugImageComponent(self))
-
-        for vision_component in self._vision_components:
-            vision_component.configure(new_config, new_config["component_debug_image_active"])
+            self._vision_components.append(make_vision_component(yoeo.DebugImageComponent))
 
-    def _register_subscribers(self, config: dict) -> None:
         self._sub_image = ros_utils.create_or_update_subscriber(
             self,
             self._config,
-            config,
+            new_config,
             self._sub_image,
             "ROS_img_msg_topic",
             Image,
@@ -124,24 +132,13 @@ def _register_subscribers(self, config: dict) -> None:
 
     @profile
     def _run_vision_pipeline(self, image_msg: Image) -> None:
-        image = self._extract_image_from_message(image_msg)
-        if image is None:
-            logger.error("Vision pipeline - Image content is None")
-            return
+        image = self._cv_bridge.imgmsg_to_cv2(image_msg, "bgr8")
 
-        self._forward_image_to_components(image)
-        self._run_components(image_msg)
-
-    def _extract_image_from_message(self, image_msg: Image) -> np.ndarray:
-        return self._cv_bridge.imgmsg_to_cv2(image_msg, "bgr8")
-
-    def _forward_image_to_components(self, image: np.ndarray) -> None:
-        for vision_component in self._vision_components:
-            vision_component.set_image(image)
+        assert self._debug_image is not None, "Debug image not initialized"
+        self._debug_image.set_image(image)
 
-    def _run_components(self, image_msg: Image) -> None:
         for vision_component in self._vision_components:
-            vision_component.run(image_msg)
+            vision_component.run(image, image_msg.header)
 
 
 def main(args=None):

bitbots_vision/bitbots_vision/vision_modules/candidate.py

@@ -3,9 +3,8 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    import numpy as np
+    pass
 
-import abc
 
 from rclpy import logging
 
@@ -229,61 +228,3 @@ def from_x1y1x2y2(cls, x1, y1, x2, y2, rating=None):
         width = abs(x1 - x2)
         height = abs(y1 - y2)
         return cls(min(x1, x2), min(y1, y2), width, height, rating)
-
-
-class CandidateFinder:
-    """
-    The abstract class :class:`.CandidateFinder` requires its subclasses to implement the methods
-    :meth:`.get_candidates`,  :meth:`.compute` and :meth:`set_image`.
-
-    CandidateFinder implementations produce a set of so called *Candidates* which are instances of the class :class:`bitbots_vision.vision_modules.candidate.Candidate`.
-    """
-
-    def __init__(self):
-        """
-        Initialization of :class:`.CandidateFinder`.
-        """
-        super().__init__()
-
-    def get_top_candidates(self, count: int = 1) -> list[Candidate]:
-        """
-        Returns the count highest rated candidates.
-
-        :param int count: Number of top-candidates to return
-        :return [Candidate]: The count top-candidates
-        """
-        candidates = self.get_candidates()
-        candidates = Candidate.sort_candidates(candidates)
-        return candidates[:count]
-
-    def get_top_candidate(self) -> Candidate:
-        """
-        Returns the highest rated candidate.
-
-        :return Candidate: Top candidate or None
-        """
-        return Candidate.select_top_candidate(self.get_candidates())
-
-    @abc.abstractmethod
-    def get_candidates(self) -> list[Candidate]:
-        """
-        Returns a list of all candidates.
-
-        :return [Candidate]: Candidates
-        """
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def compute(self) -> None:
-        """
-        Runs the most intense calculation without returning any output and caches the result.
-        """
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def set_image(self, image: np.ndarray) -> None:
-        """
-        Set the image which is going to be processed by calling :meth:`.compute.
-        """
-
-        raise NotImplementedError

bitbots_vision/bitbots_vision/vision_modules/debug.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Callable, Optional
+from typing import Optional
 
 import cv2
 import numpy as np
@@ -129,15 +129,14 @@ def draw_mask(self, mask, color, opacity=0.5):
         if not self.active:
             return
         assert self._debug_image is not None, "No image set"
-        # Make a colored image
-        colored_image = np.zeros_like(self._debug_image)
-        colored_image[:, :] = tuple(np.multiply(color, opacity).astype(np.uint8))
 
-        # Compose debug image with lines
-        self._debug_image = cv2.add(
-            cv2.bitwise_and(self._debug_image, self._debug_image, mask=255 - mask),
-            cv2.add(colored_image * opacity, self._debug_image * (1 - opacity), mask=mask).astype(np.uint8),
-        )
+        # Bring mask into canonical form
+        mask_bool = mask.astype(bool).squeeze()
+
+        # Blend color directly where mask is true
+        self._debug_image[mask_bool, :] = (
+            self._debug_image[mask_bool, :] * (1.0 - opacity) + np.array(color, dtype=np.float32) * opacity
+        ).astype(np.uint8)
 
     def get_image(self):
         """
@@ -146,41 +145,3 @@ def get_image(self):
         :return: image with debug stuff
         """
         return self._debug_image
-
-    def draw(self, debug_image_description, image=None):
-        """
-        Draws a debug image description, that contains the style and the data for each object/class that we debug
-        E.g.:
-        {
-            'type': 'field_boundary',
-            'thickness': 1,
-            'color': (255,255,255),
-            'data': #Some data
-        }
-
-        :param debug_image_description: List of dicts contains the style and the date for each object/class that we debug
-        In the dict 'type' refers to the type that we want to draw. Some types are ['robot', 'field_boundary', 'ball', 'line_point', 'line_segment'].
-        The key 'color' defines the color as BRG. For most types this is the border color.
-        The key 'thickness' refers to the border thickness.
-        The data, so the candidates we want to draw are defined with the 'data' key.
-        :return: Image with debug stuff
-        """
-        # Set Image if transmitted (optional). Otherwise take the manual set image.
-        if image:
-            self.set_image(image)
-        # Define the draw functions for each type
-        draw_functions: dict[str, Callable] = {
-            "robot": self.draw_robot_candidates,
-            "field_boundary": self.draw_field_boundary,
-            "ball": self.draw_ball_candidates,
-            "line_point": self.draw_points,
-            "line_segment": self.draw_line_segments,
-        }
-        # Draw all entries
-        for draw_type in debug_image_description:
-            # Get drawing function from dict
-            draw_function = draw_functions[draw_type["type"]]
-            # Call drawing function
-            draw_function(draw_type["data"], draw_type["color"], draw_type["thickness"])
-        # Return the image
-        return self.get_image()