Read every image once for all channels in step 3

ledsa/core/image_reading.py

@@ -101,20 +101,64 @@ def read_channel_data_mask_and_noise_from_img(
     :raises CorruptedImageError: If the image cannot be read or has an unsupported file extension.
     :raises ValueError: If rotation is neither 'auto' nor one of 0, 90, 180 and 270.
     """
-    def read_img_file(img_filename, img_channel, img_orientation=None):
-        channel_array = _read_channel_data_from_img_file(img_filename, img_channel, img_orientation)
-        return channel_array, np.ones(channel_array.shape, dtype=bool), None
-
-    def read_raw_file(raw_filename, raw_channel, raw_orientation=None):
-        img_array, filter_array, noise = _read_img_array_and_noise_from_raw_file(
-            raw_filename, raw_channel, raw_orientation)
-        channel_mask = get_bayer_channel_mask(filter_array, raw_channel)
-        return (np.where(channel_mask, img_array, 0), channel_mask,
-                None if noise is None else np.where(channel_mask, noise, 0.0))
-
-    orientation = _orientation_from_rotation(rotation)
-    reader = _get_reader(filename, read_img_file, read_raw_file)
-    return _read_or_raise(filename, reader, channel, orientation)
+    return ImageChannels(filename, rotation).channel_data_mask_and_noise(channel)
+
+
+class ImageChannels:
+    """
+    An image read once for the analysis of several color channels.
+
+    Decoding a RAW file takes most of the time of reading it, and nothing the decoding yields depends on the
+    channel, because the black level and the read noise are taken per Bayer color. The file is therefore read
+    when the object is created, and channel_data_mask_and_noise builds the arrays of a channel from it. They
+    are built for one channel at a time, so that besides the decoded image only the channel in work is held in
+    memory.
+    """
+
+    def __init__(self, filename: str, rotation: str = AUTO_ROTATION):
+        """
+        :param filename: The path of the image file to read.
+        :type filename: str
+        :param rotation: Degrees the stored array is turned clockwise to be upright, 'auto' to take the rotation
+                         from the metadata of the image.
+        :type rotation: str or int
+        :raises CorruptedImageError: If the image cannot be read or has an unsupported file extension.
+        :raises ValueError: If rotation is neither 'auto' nor one of 0, 90, 180 and 270.
+        """
+        def read_img_file(img_filename, _, img_orientation=None):
+            # every pixel of a JPG or PNG belongs to every channel, and there is no optical black to take the
+            # read noise from
+            return _read_img_array_from_img_file(img_filename, img_orientation), None, None
+
+        self.filename = filename
+        orientation = _orientation_from_rotation(rotation)
+        reader = _get_reader(filename, read_img_file, _read_img_array_and_noise_from_raw_file)
+        # no channel is passed, the read is the same for all of them
+        self._img_array, self._filter_array, self._pixel_noise = _read_or_raise(filename, reader, None,
+                                                                                orientation)
+
+    def channel_data_mask_and_noise(self, channel: int) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        The channel array, the boolean channel mask and the read noise of every pixel of one color channel, as
+        read_channel_data_mask_and_noise_from_img returns them.
+
+        :param channel: The color channel to process.
+        :type channel: int
+        :return: The channel array, the boolean channel mask and the read noise of every pixel, all of the
+            same shape. The noise is None for JPG and PNG images and for sensors without an optical black
+            margin.
+        :rtype: tuple[np.ndarray, np.ndarray, np.ndarray or None]
+        :raises CorruptedImageError: If the image does not have the channel.
+        """
+        try:
+            if self._filter_array is None:
+                channel_array = self._img_array[:, :, channel]
+                return channel_array, np.ones(channel_array.shape, dtype=bool), None
+            channel_mask = get_bayer_channel_mask(self._filter_array, channel)
+            return (np.where(channel_mask, self._img_array, 0), channel_mask,
+                    None if self._pixel_noise is None else np.where(channel_mask, self._pixel_noise, 0.0))
+        except Exception as error:
+            raise CorruptedImageError(self.filename, error) from error
 
 
 def get_bayer_channel_mask(filter_array: np.ndarray, channel: int) -> np.ndarray:

ledsa/data_extraction/DataExtractor.py

@@ -279,12 +279,20 @@ def process_img_file(self, img_filename: str) -> Optional[CorruptedImageError]:
         :rtype: Optional[CorruptedImageError]
         """
         img_id = ledsa.core.image_handling.get_img_id(img_filename)
+        in_file_path = os.path.join(self.config['DEFAULT']['img_directory'], img_filename)
+        try:
+            # read once for all channels, decoding a RAW file takes most of the time of reading it
+            img_channels = ledsa.core.image_reading.ImageChannels(in_file_path,
+                                                                  rotation=self.config.get('DEFAULT', 'img_rotation'))
+        except CorruptedImageError as error:
+            return error
         for channel in self.channels:
             try:
                 img_data = ledsa.data_extraction.step_3_functions.generate_analysis_data(img_filename, channel,
                                                                                          self.search_areas,
                                                                                          self.line_indices,
-                                                                                         self.config, self.fit_leds)
+                                                                                         self.config, self.fit_leds,
+                                                                                         img_channels=img_channels)
             except CorruptedImageError as error:
                 return error
             ledsa.data_extraction.step_3_functions.create_fit_result_file(img_data, img_id, channel)

ledsa/data_extraction/step_3_functions.py

@@ -10,7 +10,7 @@
 from ledsa.core.ConfigData import ConfigData
 from ledsa.core.file_handling import read_table
 from ledsa.core.image_handling import get_img_name
-from ledsa.core.image_reading import read_channel_data_mask_and_noise_from_img
+from ledsa.core.image_reading import ImageChannels
 from ledsa.data_extraction.LEDAnalysisData import LEDAnalysisData
 from ledsa.data_extraction.model import target_function
 from ledsa.data_extraction.sideband_background import (DEFAULT_SIDEBAND_DISTANCES, DEFAULT_BACKGROUND_FLOOR,
@@ -19,7 +19,8 @@
 
 
 def generate_analysis_data(img_filename: str, channel: int, search_areas: np.ndarray, line_indices: List[List[int]],
-                           conf: ConfigData, fit_leds=True, debug=False, debug_led=None) -> List[LEDAnalysisData]:
+                           conf: ConfigData, fit_leds=True, debug=False, debug_led=None,
+                           img_channels: Optional[ImageChannels] = None) -> List[LEDAnalysisData]:
     """
     Generate LED analysis data for the given image.
 
@@ -39,12 +40,16 @@ def generate_analysis_data(img_filename: str, channel: int, search_areas: np.nda
     :type debug: bool
     :param debug_led: The specific LED to debug. Default is None.
     :type debug_led: Optional[int]
+    :param img_channels: The image, already read for the analysis of several channels. If None, the image is
+        read from the image directory of the config.
+    :type img_channels: Optional[ImageChannels]
     :return: A list of LEDAnalysisData objects containing analysis results.
     :rtype: List[LEDAnalysisData]
     """
-    file_path = os.path.join(conf['DEFAULT']['img_directory'], img_filename)
-    data, channel_mask, pixel_noise = read_channel_data_mask_and_noise_from_img(
-        file_path, channel=channel, rotation=conf.get('DEFAULT', 'img_rotation'))
+    if img_channels is None:
+        file_path = os.path.join(conf['DEFAULT']['img_directory'], img_filename)
+        img_channels = ImageChannels(file_path, rotation=conf.get('DEFAULT', 'img_rotation'))
+    data, channel_mask, pixel_noise = img_channels.channel_data_mask_and_noise(channel)
     search_area_radius = int(conf['find_search_areas']['search_area_radius'])
     img_analysis_data = []
 

ledsa/tests/UnitTests/test_image_channels.py

@@ -0,0 +1,144 @@
+"""Unit tests for reading an image once for the analysis of several color channels.
+
+Decoding a RAW file takes most of the time of reading it, so step 3 reads every image once and builds
+the arrays of all channels it analyses from that one read. The arrays have to be the same as those of
+a read for each channel.
+"""
+import configparser
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+from PIL import Image
+
+from ledsa.core import image_handling, image_reading
+from ledsa.core.image_integrity import CorruptedImageError
+from ledsa.core.image_reading import ImageChannels, RawImage
+from ledsa.data_extraction import step_3_functions
+from ledsa.data_extraction.DataExtractor import DataExtractor
+
+# RGGB pattern as in rawpy raw_colors_visible: 0 red, 1 and 3 green, 2 blue
+FILTER_ARRAY = np.tile(np.array([[0, 1], [3, 2]]), (2, 3))
+CHANNEL_COLORS = {0: [0], 1: [1, 3], 2: [2]}
+# different per Bayer color, so that a test can tell the noise of one channel from that of another
+READ_NOISE = np.array([1.0, 2.0, 3.0, 4.0])
+SIDEWAYS = 6
+
+
+@pytest.fixture
+def stored():
+    return np.arange(FILTER_ARRAY.size, dtype=np.float32).reshape(FILTER_ARRAY.shape)
+
+
+@pytest.fixture
+def raw_reads(monkeypatch, stored):
+    """
+    Let the one pass over a RAW file return a prepared image and count how often it is made. A black level
+    of 0 and a white level of 255 leave the stored values as they are.
+    """
+    def install(orientation=image_reading.UPRIGHT_ORIENTATION):
+        reads = []
+
+        def read_raw_image(filename):
+            reads.append(filename)
+            return RawImage(stored, FILTER_ARRAY, np.zeros(4), 255.0, orientation, READ_NOISE)
+
+        monkeypatch.setattr(image_reading, '_read_raw_image', read_raw_image)
+        return reads
+    return install
+
+
+class TestImageChannels:
+    @pytest.mark.parametrize('orientation', [image_reading.UPRIGHT_ORIENTATION, SIDEWAYS])
+    def test_raw_channels_are_built_from_one_read(self, raw_reads, stored, orientation):
+        reads = raw_reads(orientation)
+        img_channels = ImageChannels('img.CR3')
+        upright = image_reading.ORIENTATION_TRANSFORMATIONS[orientation]
+        for channel, colors in CHANNEL_COLORS.items():
+            channel_array, channel_mask, pixel_noise = img_channels.channel_data_mask_and_noise(channel)
+            mask = upright(np.isin(FILTER_ARRAY, colors))
+            np.testing.assert_array_equal(channel_mask, mask)
+            np.testing.assert_array_equal(channel_array, np.where(mask, upright(stored), 0))
+            np.testing.assert_array_equal(pixel_noise, np.where(mask, upright(READ_NOISE[FILTER_ARRAY]), 0))
+        assert reads == ['img.CR3']
+
+    def test_raw_channels_are_those_of_a_read_per_channel(self, raw_reads):
+        raw_reads(SIDEWAYS)
+        img_channels = ImageChannels('img.CR3')
+        for channel in CHANNEL_COLORS:
+            for once, per_channel in zip(img_channels.channel_data_mask_and_noise(channel),
+                                         image_reading.read_channel_data_mask_and_noise_from_img('img.CR3', channel)):
+                np.testing.assert_array_equal(once, per_channel)
+            np.testing.assert_array_equal(img_channels.channel_data_mask_and_noise(channel)[0],
+                                          image_reading.read_channel_data_from_img('img.CR3', channel))
+
+    def test_rotation_option_overrides_the_raw_metadata(self, raw_reads, stored):
+        raw_reads(SIDEWAYS)
+        channel_array, _, _ = ImageChannels('img.CR3', rotation=0).channel_data_mask_and_noise(0)
+        np.testing.assert_array_equal(channel_array, np.where(FILTER_ARRAY == 0, stored, 0))
+
+    def test_png_channels_are_the_color_planes(self, tmp_path):
+        path = str(tmp_path / 'img.png')
+        plt.imsave(path, np.random.default_rng(0).random((6, 8, 3)))
+        img_channels = ImageChannels(path)
+        for channel in CHANNEL_COLORS:
+            channel_array, channel_mask, pixel_noise = img_channels.channel_data_mask_and_noise(channel)
+            np.testing.assert_array_equal(channel_array, image_reading.read_channel_data_from_img(path, channel))
+            assert channel_mask.shape == channel_array.shape and channel_mask.all()
+            assert pixel_noise is None  # a PNG has no optical black
+
+    def test_missing_channel_is_reported_as_corrupted_image(self, tmp_path):
+        path = str(tmp_path / 'gray.png')
+        Image.fromarray(np.zeros((6, 8), dtype=np.uint8)).save(path)  # read as a 2D array without color planes
+        img_channels = ImageChannels(path)
+        with pytest.raises(CorruptedImageError, match='gray.png'):
+            img_channels.channel_data_mask_and_noise(1)
+
+    def test_unreadable_image_raises_when_it_is_read(self, tmp_path):
+        with pytest.raises(CorruptedImageError, match='missing.CR3'):
+            ImageChannels(str(tmp_path / 'missing.CR3'))
+
+    def test_unsupported_format_raises(self):
+        with pytest.raises(CorruptedImageError, match='unsupported file extension'):
+            ImageChannels('img.xyz')
+
+
+class TestStep3ReadsEveryImageOnce:
+    @pytest.fixture
+    def extractor(self, tmp_path, monkeypatch):
+        """A DataExtractor for three channels, without the config file and the results of steps 1 and 2."""
+        monkeypatch.chdir(tmp_path)
+        monkeypatch.setattr(image_handling, 'get_img_id', lambda img_filename: '42')
+        config = configparser.ConfigParser()
+        config['DEFAULT'] = {'img_directory': str(tmp_path), 'img_rotation': 'auto'}
+        extractor = DataExtractor.__new__(DataExtractor)
+        extractor.channels = [0, 1, 2]
+        extractor.config = config
+        extractor.search_areas = np.zeros((1, 3))
+        extractor.line_indices = [[0]]
+        extractor.fit_leds = False
+        return extractor
+
+    def test_all_channels_are_analysed_from_one_read(self, extractor, raw_reads, monkeypatch):
+        reads = raw_reads()
+        analysed, written = [], []
+
+        def generate_analysis_data(img_filename, channel, *args, img_channels=None, **kwargs):
+            analysed.append(channel)
+            return img_channels.channel_data_mask_and_noise(channel)
+
+        monkeypatch.setattr(step_3_functions, 'generate_analysis_data', generate_analysis_data)
+        monkeypatch.setattr(step_3_functions, 'create_fit_result_file',
+                            lambda img_data, img_id, channel: written.append((img_id, channel)))
+        assert extractor.process_img_file('V001_Cam01_0042.CR3') is None
+        assert len(reads) == 1
+        assert analysed == [0, 1, 2]
+        assert written == [('42', 0), ('42', 1), ('42', 2)]
+
+    def test_unreadable_image_is_skipped_for_all_channels(self, extractor, monkeypatch):
+        written = []
+        monkeypatch.setattr(step_3_functions, 'create_fit_result_file',
+                            lambda img_data, img_id, channel: written.append(channel))
+        error = extractor.process_img_file('V001_Cam01_0042.CR3')  # there is no such file
+        assert isinstance(error, CorruptedImageError)
+        assert written == []