setSLMImage() with np.arrays

pymmcore/__init__.pyi

@@ -1014,10 +1014,22 @@ class CMMCore:
         """Sets the current slm device."""
     def setSLMExposure(self, slmLabel: str, exposure_ms: float) -> None:
         """For SLM devices with build-in light source (such as projectors),
-
         this will set the exposure time, but not (yet) start the illumination"""
+    @overload
+    def setSLMImage(self, slmLabel: str, pixels: np.ndarray[np.uint8]) -> None:
+        """
+        Write a 8-bit grayscale image to the SLM. Pixels must be a 2D numpy array [h,w] of uint8s.
+        Warning: SLM might convert grayscale to binary internally.
+        """
+    @overload
+    def setSLMImage(self, slmLabel: str, pixels: np.ndarray[np.uint8, np.uint8, np.uint8]) -> None:
+        """
+        Write a color image to the SLM (imgRGB32). The pixels must be 3D numpy array [h,w,c] of uint8s with 3 color channels [R,G,B].
+        The dimensions of the array should match the width and height of the SLM.
+        """
+    @overload
     def setSLMImage(self, slmLabel: str, pixels: Any) -> None:
-        """Write a 32-bit color image to the SLM."""
+        """Write a list of chars to the SLM. Length of the list must match the number of pixels (or 4*number of pixels to write an imgRGB32.)"""
     @overload
     def setSLMPixelsTo(self, slmLabel: str, intensity: int) -> None:
         """Set all SLM pixels to a single 8-bit intensity."""

pymmcore/pymmcore_swig.i

@@ -183,31 +183,101 @@ import_array();
 }
 
 %rename(setSLMImage) setSLMImage_pywrap;
+%apply (PyObject *INPUT, int LENGTH) { (PyObject *pixels, int receivedLength) };
 %apply (char *STRING, int LENGTH) { (char *pixels, int receivedLength) };
 %extend CMMCore {
-void setSLMImage_pywrap(const char* slmLabel, char *pixels, int receivedLength) throw (CMMError)
-{
-    // TODO This size check is done here (instead of in MMCore) because the
-    // CMMCore::setSLMImage() interface is deficient: it does not include a
-    // length parameter. It will be better to change the CMMCore functions to
-    // require a length and move this check there.
+    // This is a wrapper for setSLMImage that accepts a list of chars
+    void setSLMImage_pywrap(const char* slmLabel, char *pixels, int receivedLength) throw (CMMError)
+    {
+        // TODO This size check is done here (instead of in MMCore) because the
+        // CMMCore::setSLMImage() interface is deficient: it does not include a
+        // length parameter. It will be better to change the CMMCore functions to
+        // require a length and move this check there.
 
-    long expectedLength = self->getSLMWidth(slmLabel) * self->getSLMHeight(slmLabel);
+        long expectedLength = self->getSLMWidth(slmLabel) * self->getSLMHeight(slmLabel);
 
-    if (receivedLength == expectedLength)
-    {
-        self->setSLMImage(slmLabel, (unsigned char *)pixels);
-    }
-    else if (receivedLength == 4*expectedLength)
-    {
-        self->setSLMImage(slmLabel, (imgRGB32)pixels);
+        if (receivedLength == expectedLength)
+        {
+            self->setSLMImage(slmLabel, (unsigned char *)pixels);
+        }
+        else if (receivedLength == 4*expectedLength)
+        {
+            self->setSLMImage(slmLabel, (imgRGB32)pixels);
+        }
+        else
+        {
+            throw CMMError("Pixels must be a 2D numpy array [h,w] of uint8, or a 3D numpy array [h,w,c] of uint8 with 3 color channels [R,G,B]");
+        }
     }
-    else
+
+    // This is a wrapper for setSLMImage that accepts a numpy array
+    void setSLMImage_pywrap(const char* slmLabel, PyObject *pixels) throw (CMMError)
     {
-        throw CMMError("Image dimensions are wrong for this SLM");
+        // Check if pixels is a numpy array
+        if (!PyArray_Check(pixels)) {
+            throw CMMError("Pixels must be a numpy array");
+        }
+
+        // Get the dimensions of the numpy array
+        PyArrayObject* np_pixels = reinterpret_cast<PyArrayObject*>(pixels);
+        int nd = PyArray_NDIM(np_pixels);
+        npy_intp* dims = PyArray_DIMS(np_pixels);
+
+        // Check if the array has the correct shape
+        long expectedWidth = self->getSLMWidth(slmLabel);
+        long expectedHeight = self->getSLMHeight(slmLabel);
+
+        if (dims[0] != expectedHeight || dims[1] != expectedWidth) {
+            std::ostringstream oss;
+            oss << "Image dimensions are wrong for this SLM. Expected (" << expectedHeight << ", " << expectedWidth << "), but received (" << dims[0] << ", " << dims[1] << ")";
+            throw CMMError(oss.str().c_str());
+        }
+
+        if (PyArray_TYPE(np_pixels) != NPY_UINT8) {
+            std::ostringstream oss;
+            oss << "Pixel array type is wrong. Expected uint8.";
+            throw CMMError(oss.str().c_str());
+        }
+
+        npy_intp num_bytes = PyArray_NBYTES(np_pixels);
+        long expectedBytes = expectedWidth * expectedHeight * self->getSLMBytesPerPixel(slmLabel);
+        if (num_bytes > expectedBytes) {
+            std::ostringstream oss;
+            oss << "Number of bytes per pixel in pixels is greater than expected. Received: " << num_bytes/(dims[0] * dims[1]) << ", Expected: " << self->getSLMBytesPerPixel(slmLabel)<< ". Does this SLM support RGB?";
+            throw CMMError(oss.str().c_str());
+        }
+
+        if (PyArray_TYPE(np_pixels) == NPY_UINT8 && nd == 2) {
+            // For 2D 8-bit array, cast integers directly to unsigned char
+            std::vector<unsigned char> vec_pixels(expectedWidth * expectedHeight);
+            for (npy_intp i = 0; i < expectedHeight; ++i) {
+                for (npy_intp j = 0; j < expectedWidth; ++j) {
+                    vec_pixels[i * expectedWidth + j] = static_cast<unsigned char>(*static_cast<uint8_t*>(PyArray_GETPTR2(np_pixels, i, j)));
+                }
+            }
+            self->setSLMImage(slmLabel, vec_pixels.data());
+
+        } else if (PyArray_TYPE(np_pixels) == NPY_UINT8 && nd == 3 && dims[2] == 3) {
+            // For 3D color array, convert to imgRGB32 and add a 4th byte for the alpha channel
+            std::vector<unsigned int> vec_pixels(expectedWidth * expectedHeight); // 1 imgRGB32 for RGBA
+            for (npy_intp i = 0; i < expectedHeight; ++i) {
+                for (npy_intp j = 0; j < expectedWidth; ++j) {
+                    unsigned int pixel = 0;
+                    for (npy_intp k = 0; k < 3; ++k) {
+                        uint8_t value = *static_cast<uint8_t*>(PyArray_GETPTR3(np_pixels, i, j, 2 - k)); // Reverse the order of RGB
+                        pixel |= static_cast<unsigned int>(value) << (8 * k);
+                    }
+                    // Set the alpha channel to 0
+                    vec_pixels[i * expectedWidth + j] = pixel;
+                }
+            }
+            self->setSLMImage(slmLabel, vec_pixels.data());
+        } else {
+        throw CMMError("Pixels must be a 2D numpy array [h,w] of uint8, or a 3D numpy array [h,w,c] of uint8 with 3 color channels [R,G,B]");
+        }
     }
 }
-}
+
 %ignore setSLMImage;
 
 %{