Source code for owl.analysis.z_stack_functions

from threading import Event, Thread

import cv2
import numpy as np
import xarray as xr

from owl.analysis.util import create_full_dataset

from ..util.index_tricks import ndrange


def _tqdm(*_, **__):
    return _[0]


def measure_slice_array(mcam_dataarray, *,
                        best_images=None,
                        measured_list_array=None,
                        downsample=2,
                        start_pixel=(0, 1),
                        regions_of_interest=None):
    """Measure focus of individual images of given image set.

    Given a set of images from mcam_data measure the focus of the entire set of
    data considering each image individually, appends to or creates a list of
    focus score array and height pairs and creates and array of the highest
    individually scored images between the images given to measure and the
    previously best scoring images.

    Parameters
    ----------
    mcam_dataarray: xarray DataArray
        DataArray of mcam_data.
    best_images: mcam_data, optional
        Image taken from mcam_data as the current best scoring images to compare
        to the current_images. These images are replaced by current_images if the
        given an individual image in the image data has the highest focus score in
        the measured list for its camera position.
    measured_list_array: xarray.DataArray
        Array of focus scores. New scores will be appended
    downsample : int
        integer amount to down sample the image in both x and y direction. Should
        be a multiple of 2 to avoid issues from bayer pattern.
    start_pixel : tuple of int, optional
        The position of the starting pixel in the image to begin the down sampling.
    regions_of_interest : array of bool or None
        Array to index the areas of the individual images that will be used for
        measuring the focus of. If None, the full image is used.

    Returns
    -------
    measured_list_array: xarray.DataArray
        Array of focus scores.

    best_images: mcam_data
        The current highest focus scoring image set.

    """
    # Ensure that we have the entire data in memory so we can repeatidly access it
    mcam_dataarray = mcam_dataarray.compute()

    score_array = calc_laplacian_variance(
        mcam_dataarray,
        ksize=3,
        downsample=downsample,
        start_pixel=start_pixel,
        regions_of_interest=regions_of_interest,
    )

    frame_number = int(np.asarray(mcam_dataarray.frame_number))
    dims = tuple(mcam_dataarray.dims)

    new_measured_list_array = xr.DataArray(score_array[None, ...],
                                           dims=('frame_number',) + mcam_dataarray.dims[:2],
                                           coords=[[frame_number],
                                                   np.asarray(mcam_dataarray[dims[0]]),
                                                   np.asarray(mcam_dataarray[dims[1]])])

    if measured_list_array is None:
        measured_list_array = new_measured_list_array
    else:
        if not _measure_list_concatable(measured_list_array, new_measured_list_array):
            raise ValueError('Acquired `measured_list_array` cannot be concatenated '
                             'to given `measured_list_array`')
        measured_list_array = xr.concat([measured_list_array, new_measured_list_array],
                                        'frame_number')
    if best_images is not None:
        best_frame_indices = np.asarray(measured_list_array.argmax('frame_number'))
        new_best_images_indices = frame_number == best_frame_indices
        best_images.data[new_best_images_indices] = mcam_dataarray.data[new_best_images_indices]
    else:
        # allocate new space in memory to hold the current mcam_data to be used
        # as best_images location
        best_images = mcam_dataarray.copy()

    best_images['frame_number'] = (
        ('image_y', 'image_x'),
        np.asarray(measured_list_array['frame_number'][measured_list_array.argmax('frame_number')]),
    )

    return measured_list_array, best_images


def _measure_list_concatable(measured_list_array, new_measured_list_array):
    image_y = measured_list_array.image_y.data
    image_x = measured_list_array.image_x.data
    new_image_y = new_measured_list_array.image_y.data
    new_image_x = new_measured_list_array.image_x.data
    if len(image_x) != len(new_image_x):
        return False
    if len(image_y) != len(new_image_y):
        return False
    if (image_x != new_image_x).any():
        return False
    return not (image_y != new_image_y).any()


[docs] def find_best_z_position_array(mcam_dataset, tqdm=_tqdm, stop_event=None): """Find array of most in focused images considering images individually Given a list of z heights and a directory holding the z-stack images measure the focus of individual images and give an array of the most in focus individual images. Measures the focus of the images and organizes data with ``measure_slice_array`` Parameters ---------- mcam_dataset: xarray Dataset MCAM data. tqdm: tqdm progress indicator. stop_event: An event that can be set to stop the analysis early. Returns ------- measured_list_array: xarray.DataArray Array of focus scores. best_images: xarray.DataArray DataArray of most in focused image for each sensor. """ if stop_event is None: stop_event = Event() if 'z_stage' not in mcam_dataset: raise RuntimeError("Dataset must contain z stage data") if 'frame_number' not in mcam_dataset: raise RuntimeError("Dataset must contain frame_number coordinate") z_positions = np.asarray(mcam_dataset['z_stage'][:, 0, 0]) frame_numbers = np.asarray(mcam_dataset['frame_number']) # must sort to control order of output data z_frame = sorted(zip(z_positions, frame_numbers)) measured_list_array = None best_images = None for z, frame_number in tqdm(z_frame): if stop_event.is_set(): return dataarray = mcam_dataset.isel(frame_number=frame_number).images measured_list_array, best_images = measure_slice_array( dataarray, best_images=best_images, measured_list_array=measured_list_array) return measured_list_array, best_images
def calc_laplacian_variance(mcam_data, ksize=3, downsample=2, start_pixel=(0, 1), regions_of_interest=None): """ given a mcam_data find the variance of the laplacian's of a 2d array of image(s) This method is experimental and may change without notice due to a change in our standard method of saving z-stacks """ if downsample % 2 == 1: raise ValueError('downsample should be a multiple of 2') if np.dtype(mcam_data.dtype) != np.uint8: raise ValueError("We only support images with dtype uint8. " "To use this function, convert your data to uint8 " "before calling this function.") if mcam_data.ndim != 4: raise ValueError( 'Must be an mcam_data array with 4 dimensions: (image_y, image_x, y, x)' ) start_y, start_x = start_pixel # preallocating memory is ~15% faster for 54 images (3120x4096) edges = np.zeros( mcam_data[0, 0, start_y::downsample, start_x::downsample].shape, dtype='int16', ) images_var = np.empty(mcam_data.shape[:-2], dtype=float) for i in ndrange(mcam_data.shape[:-2]): # figure for selection are in google drive # https://drive.google.com/drive/folders/1o50Dfpi65NSU6nZ1mAAQ9PBt6V-oS1Qn image = np.asarray(mcam_data[i]) if regions_of_interest is None: region_of_interest = None else: region_of_interest = regions_of_interest[i] image_var = calc_image_laplacian_variance( image, ksize=ksize, downsample=downsample, start_pixel=start_pixel, region_of_interest=region_of_interest, dst=edges) images_var[i] = image_var return images_var def calc_image_laplacian_variance(image, *, ksize=3, downsample=2, start_pixel=(0, 1), region_of_interest=None, dst=None): start_y, start_x = start_pixel image = image[start_y::downsample, start_x::downsample] dst = cv2.Laplacian(image, cv2.CV_16S, ksize=ksize, dst=dst) if region_of_interest is not None: roi = region_of_interest[start_y::downsample, start_x::downsample] else: roi = None image_var = np.var(dst[roi]) return image_var def area_normalized_calc_image_laplacian_variance( image, *, ksize=3, downsample=2, start_pixel=(0, 1), region_of_interest=None, dst=None ): """Weight laplacian variance by region area.""" image_var = calc_image_laplacian_variance( image, ksize=ksize, downsample=downsample, start_pixel=start_pixel, region_of_interest=region_of_interest, dst=dst ) region = region_of_interest if region_of_interest else image area = np.prod(region.shape[:2]) area_normalized = image_var * area return area_normalized def set_stage_position(stage, position, speed, *, start_event, ready_event, timeout=None): if stage is None: return stage_max_speed_original = stage.maximum_speed try: stage.maximum_speed = speed ready_event.set() if start_event is not None: start_event.wait(timeout=timeout) stage.set_position(position, wait=True) finally: stage.maximum_speed = stage_max_speed_original def _validate_z_stack_parameters( mcam, z_start, z_step_size, z_num_steps, ): if mcam is None: raise ValueError("MCAM cannot be None.") if mcam.z_stage is None: raise RuntimeError("MCAM does not have a z stage connected.") minimum_position = mcam.z_stage.minimum_position maximum_position = mcam.z_stage.maximum_position if z_start < minimum_position: raise ValueError("z_start is below minimum z stage position.") if z_start > maximum_position: raise ValueError("z_start is above maximum z stage position.") if z_num_steps > 1: z_end = z_start + z_step_size * (z_num_steps - 1) else: z_end = z_start if z_end < minimum_position: raise ValueError("z_end is below minimum z stage position.") if z_end > maximum_position: raise ValueError("z_end is above maximum z stage position.") # ... what else should we check? def create_z_positions(z_start, z_step_size, z_num_steps): if z_num_steps < 2: z_positions = [z_start] else: z_positions = z_start + np.arange(0, z_num_steps) * z_step_size return z_positions def acquire_z_stack_video( mcam, *, z_start, z_step_size, z_num_steps, final_mcam_state=None, stop_event=None, tqdm=None, data_buffers=None, ): if stop_event is None: stop_event = Event() start_event = Event() # We limit the maximum frame rate to 20 fps # to avoid any complications with the z stage # being able to keep up with the frame rate # Prior t oJanuary 6, 2024, this was set to 10 fps # But we are now more confident that we can run the system # faster, especially in brigthfrield illumination (high brightness settings) maximum_frame_rate_system = 20 # 20250724 - John - We are limiting the frame rate of these two systems because # the streaming is filling up frame buffers on older direct connect systems serial_number = mcam.serial_number if serial_number in ['Kestrel0010', 'Kestrel0010R', 'Kestrel0013', 'Kestrel0013R']: maximum_frame_rate_system = 10 acceleration_time = 0.3 z_positions = create_z_positions(z_start, z_step_size, z_num_steps) if z_positions[0] > mcam.z_stage.maximum_position: raise ValueError( f"z starting position ({z_positions[0]}) is above maximum z stage position.") if z_positions[-1] > mcam.z_stage.maximum_position: raise ValueError( f"z ending position ({z_positions[-1]}) is above maximum z stage position.") if z_positions[0] < mcam.z_stage.minimum_position: raise ValueError( f"z starting position ({z_positions[0]}) is below minimum z stage position.") if z_positions[-1] < mcam.z_stage.minimum_position: raise ValueError( f"z ending position ({z_positions[-1]} is below minimum z stage position.") mcam.z_stage.set_position(z_start, wait=False) number_of_frames = len(z_positions) moving_z = number_of_frames > 1 and z_step_size != 0 starting_frame_rate = mcam.frame_rate_setpoint exposure = mcam.exposure # Only use up to 95% of the maximum datarate to avoid qunatization errors # in setting different frame rates and velocities maximum_system_fps = 0.95 * mcam.maximum_datarate / ( mcam.image_shape[0] * mcam.image_shape[1] * mcam.N_cameras[0] * mcam.N_cameras[1] ) maximum_exposure_fps = 1 / mcam.exposure if moving_z: mcam_z_stage_max_speed = mcam.z_stage.maximum_speed maximum_speed = min([ mcam_z_stage_max_speed, # We don't want to have the z stage accelerate indefinitely acceleration_time * mcam.z_stage.acceleration ]) z_step_time = abs(z_step_size) / maximum_speed z_step_fps = 1 / z_step_time maximum_frame_rate_setpoint = mcam.maximum_frame_rate_setpoint minimum_frame_rate_setpoint = mcam.minimum_frame_rate_setpoint desired_frame_rate_setpoint = max(min([ maximum_frame_rate_setpoint, maximum_frame_rate_system, maximum_system_fps, maximum_exposure_fps, z_step_fps ]), minimum_frame_rate_setpoint) maximum_exposure = 1 / desired_frame_rate_setpoint # need to account for some level of rounding if (exposure - maximum_exposure) / exposure > 0.01: raise ValueError( f"Exposure time is too long {exposure * 1E3} ms " f"for desired z stack settings. " f"Estimated maximum exposure {maximum_exposure * 1E3:.3f} ms." f"({desired_frame_rate_setpoint=:}, " f"{maximum_frame_rate_setpoint=:}, " f"{maximum_frame_rate_system=:}, " f"{maximum_system_fps=:}, " f"{maximum_exposure_fps=:}, " f"{z_step_fps=:})" ) else: desired_frame_rate_setpoint = min(maximum_system_fps, maximum_exposure_fps) # Clay / 2023/12/04 # We want to create metadata of the z positions # In the same format as our non video acquisitions # So we will create it here and backfill it z_stage_metadata = np.zeros((number_of_frames,) + mcam.N_cameras) for i, z in enumerate(z_positions): z_stage_metadata[i] = z # using create_full_dataset to create the full dataset here is a little bit of a hack, # but it's the easiest way to get the metadata in the right format and ensure that it matches # the all other stack acquisition methods full_dataset = create_full_dataset( mcam, z_num_steps=number_of_frames, well_plate_labels=False, data_buffers=data_buffers, ) mcam.frame_rate_setpoint = desired_frame_rate_setpoint mcam.exposure = exposure # ^^^^^ These are all desired settings. We don't really know what the system # supports when it rounds. if abs(mcam.exposure - exposure) / exposure > 0.01: raise RuntimeError( f"MCAM exposure ({mcam.exposure * 1E3} ms) " f"cannot be set to desired value ({exposure * 1E3} ms)." ) frame_rate = mcam.frame_rate # before we start the stage thread, we assure all stages are stopped mcam.wait_until_idle() if moving_z: ready_event = Event() stage_thread = Thread( target=set_stage_position, kwargs=dict( stage=mcam.z_stage, position=z_positions[-1], speed=abs(z_step_size) * frame_rate, start_event=start_event, ready_event=ready_event, timeout=15, ) ) stage_thread.start() ready_event.wait() dataset = mcam.acquire_high_speed_video( N_frames=number_of_frames, start_event=start_event, stop_event=stop_event, tqdm=tqdm, data_buffers=full_dataset.images.data, ) mcam.frame_rate_setpoint = starting_frame_rate mcam.exposure = exposure # If we don't want to wait till idle, we should at least stop the motion # if the user clicks stop mcam.z_stage.wait_until_idle() if stop_event.is_set(): if moving_z: stage_thread.join() return None dataset = dataset.assign({ 'z_stage': (('frame_number', 'image_y', 'image_x'), z_stage_metadata), }) for key, value in dataset.drop_vars(['images']).items(): if key in full_dataset: new_axis = tuple( ax for ax, d in enumerate(full_dataset[key].dims) if d not in value.dims ) # The ellipsis is used to ensure that the data # is copied to the correct location # Consider the following code # >>> a = np.zeros((2)) # >>> np.copyto(a[0], 2) # will not work -- TypeError: copyto() argument 1 must be numpy.ndarray, not numpy.float64 # noqa # instead # >>> np.copyto(a[0][...], 2) # will work np.copyto( # use np.newaxis so we can copy the data into "scalars" full_dataset[key].data[...], np.expand_dims(value.data, new_axis) ) if moving_z: stage_thread.join() # for 0 dimension collapse the frame_numer dimension allowing # people to take a single "snapshot" with this function if z_num_steps == 0: dataset = dataset.isel(frame_number=0) return dataset