Further JPEG capture optimisations using simplejpeg

Using simplejpeg we can do zero-copy JPEG encode. This does rely on
some extra features pushed to the simplejpeg repo.

Signed-off-by: David Plowman <[email protected]>

Author: davidplowman

picamera2/picamera2.py

@@ -1148,7 +1148,7 @@ def configure_(self, camera_config="preview"):
         self.controls = Controls(self, controls=self.camera_config['controls'])
         self.configure_count += 1
 
-        if "ScalerCrops" in self.camera_controls:
+        if False and "ScalerCrops" in self.camera_controls:
             par_crop = self.camera_controls["ScalerCrops"]
             full_fov = self.camera_controls["ScalerCrop"][1]
             scaler_crops = [par_crop[0] if camera_config["main"]["preserve_ar"] else full_fov]

picamera2/request.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 import piexif
+import simplejpeg
 from pidng.camdefs import Picamera2Camera
 from pidng.core import PICAM2DNG
 from PIL import Image
@@ -59,8 +60,8 @@ def __enter__(self) -> "MappedArray":
             else:
                 config = self.__stream.configuration
 
-            # helpers.make_array never makes a copy.
-            array = self.__request.picam2.helpers.make_array(array, config)
+            # helpers._make_array_shared never makes a copy.
+            array = self.__request.picam2.helpers._make_array_shared(array, config)
 
         self.__array = array
         return self
@@ -76,6 +77,8 @@ def array(self) -> Optional[np.ndarray]:
 
 
 class CompletedRequest:
+    FASTER_JPEG = True  # set to False to use the older JPEG encode method
+
     def __init__(self, request: Any, picam2: "Picamera2") -> None:
         self.request = request
         self.ref_count: int = 1
@@ -142,7 +145,7 @@ def make_array(self, name: str) -> np.ndarray:
         config = self.config.get(name, None)
         if config is None:
             raise RuntimeError(f'Stream {name!r} is not defined')
-        elif config['format'] == 'MJPEG': 
+        elif config['format'] == 'MJPEG':
             return np.array(Image.open(io.BytesIO(self.make_buffer(name))))
 
         # We don't want to send out an exported handle to the camera buffer, so we're going to have
@@ -196,8 +199,47 @@ def save(self, name: str, file_output: Any, format: Optional[str] = None,
         exif_data - dictionary containing user defined exif data (based on `piexif`). This will
             overwrite existing exif information generated by picamera2.
         """
-        return self.picam2.helpers.save(self.make_image(name), self.get_metadata(), file_output,
-                                        format, exif_data)
+        # We have a more optimised path for writing JPEGs using simplejpeg.
+        config = self.config.get(name, None)
+        if config is None:
+            raise RuntimeError(f'Stream {name!r} is not defined')
+        if self.FASTER_JPEG and config['format'] != "MJPEG" and \
+           self.picam2.helpers._get_format_str(file_output, format) in ('jpg', 'jpeg'):
+            quality = self.picam2.options.get("quality", 90)
+            with MappedArray(self, 'main') as m:
+                format = self.config[name]["format"]
+                if format == 'YUV420':
+                    width, height = self.config[name]['size']
+                    Y = m.array[:height, :width]
+                    reshaped = m.array.reshape((m.array.shape[0] * 2, m.array.strides[0] // 2))
+                    U = reshaped[2 * height: 2 * height + height // 2, :width // 2]
+                    V = reshaped[2 * height + height // 2:, :width // 2]
+                    output_bytes = simplejpeg.encode_jpeg_yuv_planes(Y, U, V, quality)
+                    Y = reshaped = U = V = None
+                else:
+                    FORMAT_TABLE = {"XBGR8888": "RGBX", "XRGB8888": "BGRX", "BGR888": "RGB", "RGB888": "BGR"}
+                    output_bytes = simplejpeg.encode_jpeg(m.array, quality, FORMAT_TABLE[format], '420')
+
+            exif = self.picam2.helpers._prepare_exif(self.get_metadata(), exif_data)
+
+            if isinstance(file_output, io.BytesIO):
+                f = file_output
+            else:
+                f = open(file_output, 'wb')
+            try:
+                if exif:
+                    # Splice in the exif data as we write it out.
+                    f.write(output_bytes[:2] + bytes.fromhex('ffe1') + (len(exif) + 2).to_bytes(2, 'big'))
+                    f.write(exif)
+                    f.write(output_bytes[2:])
+                else:
+                    f.write(output_bytes)
+            except Exception:
+                if f is not file_output:
+                    f.close()
+        else:
+            return self.picam2.helpers.save(self.make_image(name), self.get_metadata(), file_output,
+                                            format, exif_data)
 
     def save_dng(self, file_output: Any, name: str = "raw") -> None:
         """Save a DNG RAW image of the raw stream's buffer."""
@@ -218,17 +260,19 @@ class Helpers:
     def __init__(self, picam2: "Picamera2"):
         self.picam2 = picam2
 
-    def make_array(self, buffer: np.ndarray, config: Dict[str, Any]) -> np.ndarray:
-        """Make a 2D numpy array from the named stream's buffer."""
+    def _make_array_shared(self, buffer: np.ndarray, config: Dict[str, Any]) -> np.ndarray:
+        """Makes a 2d numpy array from the named stream's buffer without copying memory.
+
+        This method makes an array that is guaranteed to be shared with the underlying
+        buffer, that is, no copy of the pixel data is made.
+        """
         array = buffer
         fmt = config["format"]
         w, h = config["size"]
         stride = config["stride"]
 
-        # Turning the 1d array into a 2d image-like array only works if the
-        # image stride (which is in bytes) is a whole number of pixels. Even
-        # then, if they don't match exactly you will get "padding" down the RHS.
-        # Working around this requires another expensive copy of all the data.
+        # Reshape the 1d array into an image, and "slice" off any padding bytes on the
+        # right hand edge (which doesn't copy the pixel data).
         if fmt in ("BGR888", "RGB888"):
             if stride != w * 3:
                 array = array.reshape((h, stride))
@@ -263,6 +307,18 @@ def make_array(self, buffer: np.ndarray, config: Dict[str, Any]) -> np.ndarray:
             raise RuntimeError("Format " + fmt + " not supported")
         return image
 
+    def make_array(self, buffer, config):
+        """Makes a 2d numpy array for the named stream's buffer.
+
+        This method makes a copy of the underlying camera buffer, so that it can be
+        safely returned to the camera system.
+        """
+        array = self._make_array_shared(buffer, config)
+        if array.data.c_contiguous:
+            return np.copy(array)
+        else:
+            return np.ascontiguousarray(array)
+
     def _get_pil_mode(self, fmt):
         mode_lookup = {"RGB888": "BGR", "BGR888": "RGB", "XBGR8888": "RGBX", "XRGB8888": "BGRX"}
         mode = mode_lookup.get(fmt, None)
@@ -277,7 +333,7 @@ def make_image(self, buffer: np.ndarray, config: Dict[str, Any], width: Optional
         if fmt == "MJPEG":
             return Image.open(io.BytesIO(buffer))  # type: ignore
         else:
-            rgb = self.make_array(buffer, config)
+            rgb = self._make_array_shared(buffer, config)
 
         # buffer was already a copy, so don't need to worry about an extra copy for the "RGBX" mode.
         buf = rgb
@@ -294,6 +350,34 @@ def make_image(self, buffer: np.ndarray, config: Dict[str, Any], width: Optional
             pil_img = pil_img.resize((width, height))  # type: ignore
         return pil_img
 
+    def _prepare_exif(self, metadata, exif_data):
+        exif = b''
+        if "AnalogueGain" in metadata and "DigitalGain" in metadata:
+            datetime_now = datetime.now().strftime("%Y:%m:%d %H:%M:%S")
+            zero_ifd = {piexif.ImageIFD.Make: "Raspberry Pi",
+                        piexif.ImageIFD.Model: self.picam2.camera.id,
+                        piexif.ImageIFD.Software: "Picamera2",
+                        piexif.ImageIFD.DateTime: datetime_now}
+            total_gain = metadata["AnalogueGain"] * metadata["DigitalGain"]
+            exif_ifd = {piexif.ExifIFD.DateTimeOriginal: datetime_now,
+                        piexif.ExifIFD.ExposureTime: (metadata["ExposureTime"], 1000000),
+                        piexif.ExifIFD.ISOSpeedRatings: int(total_gain * 100)}
+            exif_dict = {"0th": zero_ifd, "Exif": exif_ifd}
+            # merge user provided exif data, overwriting the defaults
+            exif_dict = exif_dict | (exif_data or {})
+            exif = piexif.dump(exif_dict)
+        return exif
+
+    def _get_format_str(self, file_output, format):
+        if isinstance(format, str):
+            return format.lower()
+        elif isinstance(file_output, str):
+            return file_output.split('.')[-1].lower()
+        elif isinstance(file_output, Path):
+            return file_output.suffix.lower()
+        else:
+            raise RuntimeError("Cannot determine format to save")
+
     def save(self, img: Image.Image, metadata: Dict[str, Any], file_output: Union[str, Path], format: Optional[str] = None,
              exif_data: Optional[Dict] = None) -> None:
         """Save a JPEG or PNG image of the named stream's buffer.
@@ -305,36 +389,15 @@ def save(self, img: Image.Image, metadata: Dict[str, Any], file_output: Union[st
             exif_data = {}
         # This is probably a hideously expensive way to do a capture.
         start_time = time.monotonic()
-        exif = b''
-        if isinstance(format, str):
-            format_str = format.lower()
-        elif isinstance(file_output, str):
-            format_str = file_output.split('.')[-1].lower()
-        elif isinstance(file_output, Path):
-            format_str = file_output.suffix.lower()
-        else:
-            raise RuntimeError("Cannot determine format to save")
+        format_str = self._get_format_str(file_output, format)
         if format_str in ('png') and img.mode == 'RGBX':
             # It seems we can't save an RGBX png file, so make it RGBA instead. We can't use RGBA
             # everywhere, because we can only save an RGBX jpeg, not an RGBA one.
             img = img.convert(mode='RGBA')
+        exif = b''
         if format_str in ('jpg', 'jpeg'):
             # Make up some extra EXIF data.
-            if "AnalogueGain" in metadata and "DigitalGain" in metadata:
-                datetime_now = datetime.now().strftime("%Y:%m:%d %H:%M:%S")
-                assert self.picam2.camera is not None
-                zero_ifd = {piexif.ImageIFD.Make: "Raspberry Pi",
-                            piexif.ImageIFD.Model: self.picam2.camera.id,
-                            piexif.ImageIFD.Software: "Picamera2",
-                            piexif.ImageIFD.DateTime: datetime_now}
-                total_gain = metadata["AnalogueGain"] * metadata["DigitalGain"]
-                exif_ifd = {piexif.ExifIFD.DateTimeOriginal: datetime_now,
-                            piexif.ExifIFD.ExposureTime: (metadata["ExposureTime"], 1000000),
-                            piexif.ExifIFD.ISOSpeedRatings: int(total_gain * 100)}
-                exif_dict = {"0th": zero_ifd, "Exif": exif_ifd}
-                # merge user provided exif data, overwriting the defaults
-                exif_dict = exif_dict | exif_data
-                exif = piexif.dump(exif_dict)
+            exif = self._prepare_exif(metadata, exif_data)
         # compress_level=1 saves pngs much faster, and still gets most of the compression.
         png_compress_level = self.picam2.options.get("compress_level", 1)
         jpeg_quality = self.picam2.options.get("quality", 90)