Merge pull request #36 from WillB97/colourspace

Add support for defining how to convert a frame source's frame to grey
WillB97 · Jun 30, 2024 · eb2fd46 · eb2fd46
2 parents c518732 + 6464c09
commit eb2fd46
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/april_vision/frame_sources.py b/april_vision/frame_sources.py
@@ -16,6 +16,9 @@ class FrameSource:
     Allows april_vision.Processor to be created prior to frames being available.
     """
 
+    # The conversion to apply to the frame to get it to grayscale
+    COLOURSPACE = cv2.COLOR_BGR2GRAY
+
     def read(self, fresh: bool = True) -> NDArray:
         """
         The method for getting a new frame.

diff --git a/april_vision/utils.py b/april_vision/utils.py
@@ -2,7 +2,7 @@
 from collections import deque
 from math import hypot
 from pathlib import Path
-from typing import Deque, NamedTuple, Tuple, Union
+from typing import Deque, NamedTuple, Optional, Tuple, Union
 
 import cv2
 import numpy as np
@@ -21,9 +21,16 @@ class Frame(NamedTuple):
     colour_frame: NDArray
 
     @classmethod
-    def from_colour_frame(cls, colour_frame: NDArray) -> 'Frame':
+    def from_colour_frame(
+        cls,
+        colour_frame: NDArray,
+        colourspace: Optional[int] = cv2.COLOR_BGR2GRAY,
+    ) -> 'Frame':
         """Load frame from a colour image in a numpy array."""
-        grey_frame = cv2.cvtColor(colour_frame, cv2.COLOR_BGR2GRAY)
+        if colourspace is not None:
+            grey_frame = cv2.cvtColor(colour_frame, colourspace)
+        else:
+            grey_frame = colour_frame.copy()
 
         return cls(
             grey_frame=grey_frame,

diff --git a/april_vision/vision.py b/april_vision/vision.py
@@ -67,7 +67,10 @@ def _capture(self, fresh: bool = True) -> Frame:
         # hook to allow modification of the captured frame
         colour_frame = self.capture_filter(colour_frame)
 
-        return Frame.from_colour_frame(colour_frame)
+        return Frame.from_colour_frame(
+            colour_frame,
+            colourspace=self._frame_source.COLOURSPACE,
+        )
 
     def _detect(self, frame: Frame) -> List[Marker]:
         """Locate fiducial markers in frame using pyapriltags."""
@@ -193,15 +196,15 @@ def see(self, *, frame: Optional[NDArray] = None) -> List[Marker]:
         if frame is None:
             frames = self._capture()
         else:
-            frames = Frame.from_colour_frame(frame)
+            frames = Frame.from_colour_frame(frame, colourspace=self._frame_source.COLOURSPACE)
         return self._detect(frames)
 
     def see_ids(self, *, frame: Optional[NDArray] = None) -> List[int]:
         """Get a list of visible marker IDs."""
         if frame is None:
             frames = self._capture()
         else:
-            frames = Frame.from_colour_frame(frame)
+            frames = Frame.from_colour_frame(frame, colourspace=self._frame_source.COLOURSPACE)
         markers = self._detect(frames)
         return [marker.id for marker in markers]
 
@@ -215,7 +218,7 @@ def save(
         if frame is None:
             frames = self._capture()
         else:
-            frames = Frame.from_colour_frame(frame)
+            frames = Frame.from_colour_frame(frame, colourspace=self._frame_source.COLOURSPACE)
         markers = self._detect(frames)
         frames = self._annotate(
             frames,