Repository: bluesky/bluesky
Pull request #1816 (2024-10-15 20:32:32)

Add JPEG consolidator for JPEG image stacks
<!--- Provide a general summary of your changes in the Title above -->

## Description
Makes a superclass `MultipartRelatedConsolidator` that contains the logic for templating based on file number (not attached to the name, perhaps someone could suggest a better one?). Adds `JPEGConsolidator` that goes along with `TIFFConsolidator` to allow for collecting JPEG images.

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here. -->
While not used for scientific data that will be analyzed, beamlines frequently would like the ability to save jpeg snapshots or image series of monitoring cameras or the like.

## How Has This Been Tested?
<!--- Please describe in detail how you tested your changes. -->
<!--- Include details of your testing environment, and the tests you ran to -->
<!--- see how your change affects other areas of the code, etc. -->
Tested with jpeg image stack collected with tomo demo setup in the lab. Requires a source install of a specific branch of ophyd async. Data readback requires https://github.com/bluesky/tiled/pull/796

<!--
## Screenshots (if appropriate):
-->



Patch:
diff --git a/src/bluesky/consolidators.py b/src/bluesky/consolidators.py
index c12a8875..35880e7d 100644
--- a/src/bluesky/consolidators.py
+++ b/src/bluesky/consolidators.py
@@ -1,6 +1,7 @@
 import collections
 import dataclasses
 import enum
+import os
 from typing import Any, Optional
 
 import numpy as np
@@ -271,21 +272,24 @@ class HDF5Consolidator(ConsolidatorBase):
         return {"dataset": self._sres_parameters["dataset"].strip("/").split("/"), "swmr": self.swmr}
 
 
-class TIFFConsolidator(ConsolidatorBase):
-    supported_mimetypes = {"multipart/related;type=image/tiff"}
-
-    def __init__(self, stream_resource: StreamResource, descriptor: EventDescriptor):
+class MultipartRelatedConsolidator(ConsolidatorBase):
+    def __init__(
+        self, permitted_extensions: set[str], stream_resource: StreamResource, descriptor: EventDescriptor
+    ):
         super().__init__(stream_resource, descriptor)
+        self.permitted_extensions: set[str] = permitted_extensions
         self.data_uris: list[str] = []
+        self.template = self._sres_parameters["template"]
 
     def get_datum_uri(self, indx: int):
-        """Return a full uri for a datum (an individual TIFF file) based on its index in the sequence.
+        """Return a full uri for a datum (an individual image file) based on its index in the sequence.
 
         This relies on the `template` parameter passed in the StreamResource, which is a string either in the "new"
         Python formatting style that can be evaluated to a file name using the `.format(indx)` method given an
-        integer index, e.g. "{:05d}.tif".
+        integer index, e.g. "{:05d}.ext".
         """
-        return self.uri + self._sres_parameters["template"].format(indx)
+        assert os.path.splitext(self.template)[1] in self.permitted_extensions
+        return self.uri + self.template.format(indx)
 
     def consume_stream_datum(self, doc: StreamDatum):
         # Determine the indices in the names of tiff files from indices of frames and number of frames per file
@@ -305,9 +309,24 @@ class TIFFConsolidator(ConsolidatorBase):
         super().consume_stream_datum(doc)
 
 
+class TIFFConsolidator(MultipartRelatedConsolidator):
+    supported_mimetypes = {"multipart/related;type=image/tiff"}
+
+    def __init__(self, stream_resource: StreamResource, descriptor: EventDescriptor):
+        super().__init__({".tif", ".tiff"}, stream_resource, descriptor)
+
+
+class JPEGConsolidator(MultipartRelatedConsolidator):
+    supported_mimetypes = {"multipart/related;type=image/jpeg"}
+
+    def __init__(self, stream_resource: StreamResource, descriptor: EventDescriptor):
+        super().__init__({".jpeg", ".jpg"}, stream_resource, descriptor)
+
+
 CONSOLIDATOR_REGISTRY = {
     "application/x-hdf5": HDF5Consolidator,
     "multipart/related;type=image/tiff": TIFFConsolidator,
+    "multipart/related;type=image/jpeg": JPEGConsolidator,
 }
 
 


Test patch:
diff --git a/src/bluesky/tests/test_consolidators.py b/src/bluesky/tests/test_consolidators.py
index a78d1fab..b73e6f0c 100644
--- a/src/bluesky/tests/test_consolidators.py
+++ b/src/bluesky/tests/test_consolidators.py
@@ -2,7 +2,7 @@ from math import ceil
 
 import pytest
 
-from bluesky.consolidators import HDF5Consolidator, TIFFConsolidator
+from bluesky.consolidators import HDF5Consolidator, consolidator_factory
 
 
 @pytest.fixture
@@ -59,12 +59,12 @@ def hdf5_stream_resource_factory():
 
 
 @pytest.fixture
-def tiff_stream_resource_factory():
-    return lambda data_key, chunk_shape: {
+def image_seq_stream_resource_factory():
+    return lambda image_format, data_key, chunk_shape: {
         "data_key": data_key,
-        "mimetype": "multipart/related;type=image/tiff",
+        "mimetype": f"multipart/related;type=image/{image_format}",
         "uri": "file://localhost/test/file/path",
-        "parameters": {"chunk_shape": chunk_shape, "template": "img_{:06d}.tiff"},
+        "parameters": {"chunk_shape": chunk_shape, "template": "img_{:06d}." + image_format},
         "uid": f"stream-resource-uid-{data_key}",
     }
 
@@ -99,13 +99,25 @@ def test_hdf5_shape(descriptor, hdf5_stream_resource_factory, stream_datum_facto
     assert cons.shape == expected
 
 
+supported_image_seq_formats = ["jpeg", "tiff"]
+
+
 @pytest.mark.parametrize("data_key, expected", shape_testdata)
+@pytest.mark.parametrize("image_format", supported_image_seq_formats)
 @pytest.mark.parametrize("files_per_stream_datum", [1, 2, 3, 5])
 def test_tiff_shape(
-    descriptor, tiff_stream_resource_factory, stream_datum_factory, data_key, expected, files_per_stream_datum
+    descriptor,
+    image_seq_stream_resource_factory,
+    stream_datum_factory,
+    image_format,
+    data_key,
+    expected,
+    files_per_stream_datum,
 ):
-    stream_resource = tiff_stream_resource_factory(data_key=data_key, chunk_shape=(1,))
-    cons = TIFFConsolidator(stream_resource, descriptor)
+    stream_resource = image_seq_stream_resource_factory(
+        image_format=image_format, data_key=data_key, chunk_shape=(1,)
+    )
+    cons = consolidator_factory(stream_resource, descriptor)
     assert cons.shape == (0, *expected[1:])
     for i in range(ceil(5 / files_per_stream_datum)):
         doc = stream_datum_factory(