Add shot_select support with set_unpack_params() method

rawpy/__init__.py

@@ -5,17 +5,19 @@
 import rawpy._rawpy
 globals().update({k:v for k,v in rawpy._rawpy.__dict__.items() if not k.startswith('_')})
 
-def imread(pathOrFile):
+def imread(pathOrFile, shot_select=0):
     """
     Convenience function that creates a :class:`rawpy.RawPy` instance, opens the given file,
     and returns the :class:`rawpy.RawPy` instance for further processing.
 
     :param str|file pathOrFile: path or file object of RAW image that will be read
+    :param int shot_select: select which image to extract from RAW files that contain multiple images
+                            (e.g., Dual Pixel RAW). Default is 0 for the first/main image.
     :rtype: :class:`rawpy.RawPy`
     """
     d = RawPy()
     if hasattr(pathOrFile, 'read'):
-        d.open_buffer(pathOrFile)
+        d.open_buffer(pathOrFile, shot_select=shot_select)
     else:
-        d.open_file(pathOrFile)
+        d.open_file(pathOrFile, shot_select=shot_select)
     return d

rawpy/_rawpy.pyx

@@ -102,7 +102,6 @@ cdef extern from "libraw.h":
         double      aber[4]        # -C 
         double      gamm[6]        # -g 
         float       user_mul[4]    # -r mul0 mul1 mul2 mul3 
-        unsigned    shot_select    # -s 
         float       bright         # -b 
         float       threshold      #  -n 
         int         half_size      # -h 
@@ -165,6 +164,18 @@ cdef extern from "libraw.h":
         # Force use x3f data decoding either if demosaic pack GPL2 enabled 
         int force_foveon_x3f
 
+    ctypedef struct libraw_raw_unpack_params_t:
+        int use_rawspeed
+        int use_dngsdk
+        unsigned options
+        unsigned shot_select
+        unsigned specials
+        unsigned max_raw_memory_mb
+        int sony_arw2_posterization_thr
+        float coolscan_nef_gamma
+        char p4shot_order[5]
+        char **custom_camera_strings
+
     ctypedef struct libraw_iparams_t:
         char        make[64]
         char        model[64]
@@ -183,6 +194,7 @@ cdef extern from "libraw.h":
         libraw_image_sizes_t        sizes
         libraw_iparams_t            idata
         libraw_output_params_t        params
+        libraw_raw_unpack_params_t  rawparams
 #         unsigned int                progress_flags
 #         unsigned int                process_warnings
         libraw_colordata_t          color
@@ -399,13 +411,15 @@ cdef class RawPy:
         with nogil:
             self.p.recycle()
 
-    def open_file(self, path):
+    def open_file(self, path, shot_select=0):
         """
         Opens the given RAW image file. Should be followed by a call to :meth:`~rawpy.RawPy.unpack`.
 
         .. NOTE:: This is a low-level method, consider using :func:`rawpy.imread` instead.
 
         :param str path: The path to the RAW image.
+        :param int shot_select: select which image to extract from RAW files that contain multiple images
+                                (e.g., Dual Pixel RAW). Default is 0 for the first/main image.
         """
         cdef wchar_t *wchars
         cdef Py_ssize_t wchars_len
@@ -421,14 +435,19 @@ cdef class RawPy:
         ELSE:
             res = self.p.open_file(path.encode('UTF-8'))
         self.handle_error(res)
+        # Set shot_select after opening file
+        cdef libraw_raw_unpack_params_t* rp = &self.p.imgdata.rawparams
+        rp.shot_select = shot_select
 
-    def open_buffer(self, fileobj):
+    def open_buffer(self, fileobj, shot_select=0):
         """
         Opens the given RAW image file-like object. Should be followed by a call to :meth:`~rawpy.RawPy.unpack`.
 
         .. NOTE:: This is a low-level method, consider using :func:`rawpy.imread` instead.
 
         :param file fileobj: The file-like object.
+        :param int shot_select: select which image to extract from RAW files that contain multiple images
+                                (e.g., Dual Pixel RAW). Default is 0 for the first/main image.
         """
         self.unpack_called = False
         self.unpack_thumb_called = False
@@ -439,6 +458,9 @@ cdef class RawPy:
         with nogil:
             e = self.p.open_buffer(buf, buf_len)
         self.handle_error(e)
+        # Set shot_select after opening buffer
+        cdef libraw_raw_unpack_params_t* rp = &self.p.imgdata.rawparams
+        rp.shot_select = shot_select
 
     def unpack(self):
         """
@@ -1208,7 +1230,7 @@ class Params(object):
             self.aber = (chromatic_aberration[0], chromatic_aberration[1])
         else:
             self.aber = (1, 1)
-        self.bad_pixels = bad_pixels_path            
+        self.bad_pixels = bad_pixels_path
 
 cdef class processed_image_wrapper:
     cdef RawPy raw

test/test_shot_select.py

@@ -0,0 +1,45 @@
+"""
+Test for shot_select parameter functionality
+"""
+import os
+import pytest
+import rawpy
+
+thisDir = os.path.dirname(__file__)
+rawTestPath = os.path.join(thisDir, 'iss030e122639.NEF')
+
+
+def test_shot_select_parameter_via_imread():
+    """Test that shot_select parameter can be passed to imread"""
+    # Test default shot_select=0
+    with rawpy.imread(rawTestPath, shot_select=0) as raw:
+        rgb = raw.postprocess(no_auto_bright=True)
+        assert rgb is not None
+        assert rgb.shape[2] == 3  # RGB image
+
+
+def test_shot_select_nonexistent_image():
+    """Test that shot_select=1 raises error for single-image files"""
+    # Test shot_select=1 on a file with only one image should raise an error
+    with pytest.raises(rawpy.LibRawRequestForNonexistentImageError):
+        with rawpy.imread(rawTestPath, shot_select=1) as raw:
+            rgb = raw.postprocess(no_auto_bright=True)
+
+
+def test_shot_select_via_open_file():
+    """Test that shot_select can be passed via open_file"""
+    raw = rawpy.RawPy()
+    raw.open_file(rawTestPath, shot_select=0)
+    raw.unpack()
+    rgb = raw.postprocess()
+    assert rgb is not None
+    assert rgb.shape[2] == 3
+    raw.close()
+
+
+def test_shot_select_default_value():
+    """Test that default shot_select value is 0"""
+    with rawpy.imread(rawTestPath) as raw:
+        rgb = raw.postprocess()
+        assert rgb is not None
+        assert rgb.shape[2] == 3