registration in starfish, Part 2: expose image translation

[berl]

If the registration interface is getting reworked to define a reference image #373, and will continue to be used, we should expose the shift method as functionality within the registration module. This would be used so we can (for example) move around images that arrive with non-zero shifts in the metadata, unless this is already handled in slicedimage.

[dganguli]

I think I have a version of this implemented. It looks something like below. Does this serve your needs? If so I can update the registration module to look like this, and potentially incorporate #408 to zero out values as well.

from starfish.pipeline.registration.fourier_shift import compute_shift, shift_im

def compute_shifts(images: ImageStack, reference_image_index: int, upsampling=1000):
    mp = images.max_proj(Indices.CH, Indices.Z)
    reference_image = mp[reference_image_index, :]
    res = dict()
    
    for h in range(images.num_rounds):
        # compute shift between maximum projection (across channels) and dots, for each hyb round
        # TODO: make the max projection array ignorant of axes ordering.
        shift, error = compute_shift(mp[h, :, :], reference_image, upsampling)
        print("For hyb: {}, Shift: {}, Error: {}".format(h, shift, error))
        res[h] = shift

    return res
    
def register(images: ImageStack, shifts: dict):
            
    registered_images = np.zeros(images.raw_shape)

    for h in range(images.num_rounds):
        for c in range(images.num_chs):
            for z in range(images.num_zlayers):
                # apply shift to all zlayers, channels, and hyb rounds
                indices = {Indices.ROUND: h, Indices.CH: c, Indices.Z: z}
                data, axes = images.get_slice(indices=indices)
                assert len(axes) == 0
                result = shift_im(data, shifts[h])      
                result[result<0] = 0
                registered_images[h, c, z, :, :] = result

    return ImageStack.from_numpy_array(registered_images)

def crop_stack(images, shifts):
    shift_vals = [v for k, v in shifts.items()]
    shift_vals = np.array(shift_vals)
    mins = np.round(np.min(shift_vals, axis=0))
    maxs = np.round(np.max(shift_vals, axis=0))

    crop_y_top = int(np.abs(mins[0]))
    crop_y_bot = int(np.abs(maxs[0]))
    crop_x_left = int(np.abs(mins[1]))
    crop_x_right = int(np.abs(maxs[1]))
    
    new_shape = np.array(images.tile_shape) - np.array([crop_y_top+crop_y_bot, crop_x_left+crop_x_right])
    cropped_images = np.zeros((images.num_rounds, images.num_chs, images.num_zlayers, new_shape[0], new_shape[1]))
    
    for h in range(images.num_rounds):
        for c in range(images.num_chs):
            for z in range(images.num_zlayers):
                indices = {Indices.ROUND: h, Indices.CH: c, Indices.Z: z}
                data, axes = images.get_slice(indices=indices)
                assert len(axes) == 0
                result = data[crop_y_top:-crop_y_bot, crop_x_left:-crop_x_right]
                cropped_images[h, c, z, :, :] = result
                
    return ImageStack.from_numpy_array(cropped_images)

# learn shifts and crops on bright field images
shifts = compute_shifts(bright_field_stack, 1)

# apply shifts to bright field images, and crop, to inspect that registraton and cropping works
registered_bstack = register(bright_field_stack, shifts)
registered_cropped_bstack = crop_stack(registered_bstack, shifts)

# register and crop hybridization stack
registered_hstack = register(mip_stack, shifts)
registered_cropped_hstack = crop_stack(registered_hstack, shifts)

[dganguli]

as implemented above, compute_shifts takes as input an index to an image in the input stack to use as the 'reference' image. But this can be easily generalized to compute a max projection, or accept a reference image.

[berl]

yeah, this is great. maybe without the crop per #408 discussion.