Returning frames as PIL images

src/framegrab/grabber.py

@@ -12,6 +12,7 @@
 import cv2
 import numpy as np
 import yaml
+from PIL import Image
 
 from .unavailable_module import UnavailableModule
 
@@ -316,13 +317,18 @@ def autodiscover(warmup_delay: float = 1.0) -> dict:
         return grabbers
 
     @abstractmethod
-    def grab(self) -> np.ndarray:
+    def _grab_impl(self) -> np.ndarray:
         """Read a frame from the camera, zoom and crop if called for, and then perform any camera-specific
         postprocessing operations.
-        Returns a frame.
+        Returns a numpy array.
         """
         pass
 
+    def grab(self) -> Image:
+        """Executes the camera-specific grab implementation and returns the frame as a PIL Image."""
+        frame = self._grab_impl()
+        return Image.fromarray(frame)
+
     def _autogenerate_name(self) -> None:
         """For generating and assigning unique names for unnamed FrameGrabber objects.
 
@@ -558,7 +564,7 @@ def __init__(self, config: dict):
         self.idx = idx
         GenericUSBFrameGrabber.indices_in_use.add(idx)
 
-    def grab(self) -> np.ndarray:
+    def _grab_impl(self) -> np.ndarray:
         # OpenCV VideoCapture buffers frames by default. It's usually not possible to turn buffering off.
         # Buffer can be set as low as 1, but even still, if we simply read once, we will get the buffered (stale) frame.
         # Assuming buffer size of 1, we need to read twice to get the current frame.
@@ -708,7 +714,7 @@ def _close_connection(self):
             if self.capture is not None:
                 self.capture.release()
 
-    def grab(self) -> np.ndarray:
+    def _grab_impl(self) -> np.ndarray:
         if not self.keep_connection_open:
             self._open_connection()
             try:
@@ -793,7 +799,7 @@ def __init__(self, config: dict):
         self.camera = camera
         BaslerFrameGrabber.serial_numbers_in_use.add(self.config["id"]["serial_number"])
 
-    def grab(self) -> np.ndarray:
+    def _grab_impl(self) -> np.ndarray:
         with self.camera.GrabOne(2000) as result:
             if result.GrabSucceeded():
                 # Convert the image to BGR for OpenCV
@@ -884,7 +890,7 @@ def __init__(self, config: dict):
         # In case the serial_number wasn't provided by the user, add it to the config
         self.config["id"] = {"serial_number": curr_serial_number}
 
-    def grab(self) -> np.ndarray:
+    def _grab_impl(self) -> np.ndarray:
         frames = self.pipeline.wait_for_frames()
 
         # Convert color images to numpy arrays and convert from RGB to BGR
@@ -977,7 +983,7 @@ def __init__(self, config: dict):
         # In case the serial_number wasn't provided by the user, add it to the config
         self.config["id"] = {"serial_number": curr_serial_number}
 
-    def grab(self) -> np.ndarray:
+    def _grab_impl(self) -> np.ndarray:
         width = self.config.get("options", {}).get("resolution", {}).get("width", 640)
         height = self.config.get("options", {}).get("resolution", {}).get("height", 480)
 

test/test_framegrab_with_mock_camera.py

@@ -34,7 +34,8 @@ def test_crop_pixels(self):
 
         grabber.release()
 
-        assert frame.shape == (400, 400, 3)
+        assert frame.size == (400, 400)
+        assert frame.mode == 'RGB'
 
     def test_crop_relative(self):
         """Grab a frame, crop a frame in an relative manner (0 to 1), and make sure the shape is correct.
@@ -63,7 +64,8 @@ def test_crop_relative(self):
 
         grabber.release()
 
-        assert frame.shape == (384, 512, 3)
+        assert frame.size == (512, 384)
+        assert frame.mode == 'RGB'
 
     def test_zoom(self):
         """Grab a frame, zoom a frame, and make sure the shape is correct.
@@ -87,7 +89,8 @@ def test_zoom(self):
 
         grabber.release()
 
-        assert frame.shape == (240, 320, 3)
+        assert frame.size == (320, 240)
+        assert frame.mode == 'RGB'
 
     def test_attempt_create_grabber_with_invalid_input_type(self):
         config = {
@@ -157,11 +160,10 @@ def test_attempt_create_more_grabbers_than_exist(self):
         # Try to connect to another grabber, this should raise an exception because there are only 3 mock cameras available
         try:
             FrameGrabber.create_grabber({'input_type': 'mock'})
-            self.fail()
+            self.fail() # we shouldn't get here
         except ValueError:
             pass
         finally:
-            # release all the grabbers
             for grabber in grabbers.values():
                 grabber.release()