# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python (phathom)
#     language: python
#     name: phathom
# ---

# +
import neuroglancer
# Use this in IPython to allow external viewing
# neuroglancer.set_server_bind_address(bind_address='192.168.158.128',
#                                      bind_port=80)
from nuggt.utils import ngutils

viewer = neuroglancer.Viewer()
viewer

# +
import numpy as np
import zarr
import os

# working_dir = '/media/jswaney/Drive/Justin/coregistration/whole_brain_tde'
working_dir = '/home/jswaney/coregistration'

fixed_path = 'fixed/zarr_stack/8_8_8'
moving_path = 'moving/zarr_stack/8_8_8'

fixed_store = zarr.NestedDirectoryStore(os.path.join(working_dir, fixed_path))
moving_store = zarr.NestedDirectoryStore(os.path.join(working_dir, moving_path))

fixed_img = zarr.open(fixed_store, mode='r')
moving_img = zarr.open(moving_store, mode='r')

print(fixed_img.shape)
print(moving_img.shape)

# +
normalization = 2000

def plot_image(img, viewer, layer, shader):
    with viewer.txn() as txn:
        source = neuroglancer.LocalVolume(img.astype(np.float32))
        txn.layers[layer] = neuroglancer.ImageLayer(source=source, shader=shader)

def plot_fixed(fixed_img, viewer):
    fixed_shader = ngutils.red_shader % (1 / normalization)
    plot_image(fixed_img, viewer, 'fixed', fixed_shader)
    
def plot_moving(moving_img, viewer):
    moving_shader = ngutils.green_shader % (1 / normalization)
    plot_image(moving_img, viewer, 'moving', moving_shader)
    
def plot_both(fixed_img, moving_img, viewer):
    plot_fixed(fixed_img, viewer)
    plot_moving(moving_img, viewer)
    
plot_both(fixed_img, moving_img, viewer)
# -

# # Downsampling Zarr

from phathom.io.zarr import downsample_zarr

# +
factors = (8, 8, 8)
output_path = os.path.join(working_dir, 'fixed/zarr_down8')
nb_workers = 1

downsample_zarr(fixed_img, factors, output_path, nb_workers)

# +
factors = (8, 8, 8)
output_path = os.path.join(working_dir, 'moving/zarr_down8')
nb_workers = 1

downsample_zarr(moving_img, factors, output_path, nb_workers)
# -

# # Downsampling ndarray

# +
from skimage.measure import block_reduce

factors = (16, 16, 16)

def downsample_mean(img, factors):
    return block_reduce(img, factors, np.mean, 0)

def downsample_max(img, factors):
    return block_reduce(img, factors, np.max, 0)

fixed_downsample = downsample_mean(fixed_img, factors)
moving_downsample = downsample_mean(moving_img, factors)

plot_both(fixed_downsample, moving_downsample, viewer)
# -

# # Gaussian smoothing

# +
from skimage.filters import gaussian

sigma = 1.0

fixed_smooth = gaussian(fixed_downsample, sigma, preserve_range=True)
moving_smooth = gaussian(moving_downsample, sigma, preserve_range=True)

plot_both(fixed_smooth, moving_smooth, viewer)
# -

# # Destriping

# +
import pystripe
import multiprocessing
import tqdm

bandwidth = [64, 64]
wavelet = 'db4'

def _filter_streaks(img):
    return pystripe.filter_streaks(img, sigma=bandwidth, wavelet=wavelet)

with multiprocessing.Pool(12) as pool:
    rf = list(tqdm.tqdm(pool.imap(_filter_streaks, fixed_smooth), total=len(fixed_smooth)))
    rm = list(tqdm.tqdm(pool.imap(_filter_streaks, moving_smooth), total=len(moving_smooth)))

fixed_destripe = np.array(rf).T
moving_destripe = np.array(rm).T

with multiprocessing.Pool(12) as pool:
    rf = list(tqdm.tqdm(pool.imap(_filter_streaks, fixed_destripe), total=len(fixed_smooth)))
    rm = list(tqdm.tqdm(pool.imap(_filter_streaks, moving_destripe), total=len(moving_smooth)))

fixed_destripe = np.array(rf).T
moving_destripe = np.array(rm).T
    
plot_both(fixed_destripe, moving_destripe, viewer)
# -

# # Rigid transformation

# +
from phathom.registration import coarse, pcloud
from phathom import utils
from scipy.ndimage import map_coordinates

t = np.array([0, 0, 0])
thetas = np.array([np.pi/4, 0, 0])

def rigid_warp(img, t, thetas, center, output_shape):
    r = pcloud.rotation_matrix(thetas)
    idx = np.indices(output_shape)
    pts = np.reshape(idx, (idx.shape[0], idx.size//idx.shape[0])).T
    warped_pts = coarse.rigid_transformation(t, r, pts, center)
    interp_values = map_coordinates(img, warped_pts.T)
    transformed = np.reshape(interp_values, output_shape)
    return transformed

transformed = rigid_warp(fixed_downsample, 
                         t, 
                         thetas, 
                         np.zeros(3), 
                         moving_downsample.shape)

plot_fixed(transformed, viewer)

# +
from scipy.ndimage.measurements import center_of_mass

def center_mass(img):
    return np.asarray(center_of_mass(img))

fixed_com = center_mass(fixed_downsample)
moving_com = center_mass(moving_downsample)

print(fixed_com)
print(moving_com)


# -

# # Optimization

# +
def ncc(fixed, transformed, nonzero=False):
    if nonzero:
        idx = np.where(transformed)
        a = fixed[idx]
        b = transformed[idx]
    else:
        a = fixed
        b = transformed
    return np.sum((a-a.mean())*(b-b.mean())/((a.size-1)*a.std()*b.std()))

def ssd(fixed, transformed):
    return np.mean((fixed-transformed)**2)

def registration_objective(x, fixed_img, moving_img, t):
    transformed_img = rigid_warp(moving_img, 
                                 t=t, 
                                 thetas=x,
                                 center=fixed_com,
                                 output_shape=moving_img.shape)
    return ssd(moving_img, transformed_img)

def callback(x, f, accept):
    pass


# +
from scipy.optimize import basinhopping

niter = 4

t_star = moving_com-fixed_com
bounds = [(0, np.pi/2) for _ in range(3)]

res = basinhopping(registration_objective,
                       x0=np.zeros(3),
                       niter=niter,
                       T=1.0,
                       stepsize=1.0,
                       interval=5,
                       minimizer_kwargs={
                           'method': 'L-BFGS-B',
#                            'method': 'Nelder-Mead',
                           'args': (fixed_smooth, 
                                    moving_smooth, 
                                    t_star),
                           'bounds': bounds,
                           'tol': 0.01,
                           'options': {'disp': False}
                       },
                       disp=True)

theta_star = res.x
print(res)

# +
registered = rigid_warp(fixed_smooth, t_star, theta_star, fixed_com, moving_destripe.shape)

plot_fixed(registered, viewer)
# -

# # Contour

# +
import matplotlib.pyplot as plt

plt.hist(fixed_downsample.ravel(), bins=100)
plt.xlim([0, 1000])
plt.ylim([0, 100000])
plt.show()

plt.hist(moving_downsample.ravel(), bins=100)
plt.xlim([0, 1000])
plt.ylim([0, 100000])
plt.show()

# +
threshold = 150

fixed_mask = fixed_downsample > threshold
moving_mask = moving_downsample > threshold

plot_both(1000*fixed_mask, 1000*moving_mask, viewer)
# -

# # Convex hull

# +
from skimage.morphology import convex_hull_image
import tqdm

fixed_hull = np.zeros_like(fixed_mask)

for i, f in enumerate(tqdm.tqdm(fixed_mask)):
    if not np.all(f == 0):
        fixed_hull[i] = convex_hull_image(f)

moving_hull = np.zeros_like(moving_mask)
        
for i, m in enumerate(tqdm.tqdm(moving_mask)):
    if not np.all(m == 0):
        moving_hull[i] = convex_hull_image(m)

plot_both(1000*fixed_hull, 1000*moving_hull, viewer)

# +
from scipy.ndimage.morphology import distance_transform_edt

fixed_distance = distance_transform_edt(fixed_mask)
moving_distance = distance_transform_edt(moving_mask)

plot_both(100*fixed_distance, 100*moving_distance, viewer)

# +
niter = 3

from scipy.optimize import basinhopping

fixed_com = center_mass(fixed_mask)
moving_com = center_mass(moving_mask)

t0 = moving_com-fixed_com
bounds = [(-s, s) for s in moving_distance.shape]+[(-np.pi, np.pi) for _ in range(3)]
# bounds = [(-np.pi, np.pi) for _ in range(3)]

def absolute_difference(img1, img2):
    return np.mean(np.abs(img1-img2))

def registration_objective(x, fixed_img, moving_img):
    transformed_img = rigid_warp(moving_img, 
                                 t= x[:3], 
                                 thetas= x[3:],
                                 center=fixed_com,
                                 output_shape=fixed_img.shape)
    return absolute_difference(fixed_img, transformed_img)
#     return ssd(fixed_img, transformed_img)


def callback(x, f, accept):
    print(x)

res = basinhopping(registration_objective,
                       x0=np.concatenate((t0, np.zeros(3))),
                       niter=niter,
                       T=0.5,
                       stepsize=0.5,
                       interval=5,
                       minimizer_kwargs={
                           'method': 'L-BFGS-B',
                           'args': (fixed_distance, 
                                    moving_distance),
                           'bounds': bounds,
                           'tol': 0.001,
                           'options': {'disp': False}
                       },
                       callback=callback,
                       disp=True)

t_star = res.x[:3]
theta_star = res.x[3:]
print(res)

reg_distance = rigid_warp(fixed_distance, 
                          t_star, 
                          theta_star, 
                          fixed_com, 
                          moving_distance.shape)

plot_fixed(100*reg_distance, viewer)
# -

# Sum of squared differences seems to provide slightly better registration than Normalized cross-correlation in the case of distance transformed convex hulls. This might be because NCC is indifferent to intensity difference and only considers correlations in the intensities, whereas SSD will penalize for any difference in intensity. In a multi-modal setting, this is usually not desired, but since we are dealing with the same brain in both images, the overall shape (and therefore distance transforms) should take similar values (not just correlated).
#
# Also, it was necessary to include the translation component in the optimization procedure because our center of mass estimate for the center of rotation is not accurate. This causes the optimization for our rigid transformation to be partially constrained to inaccurate values, making it hard to converge to a rotation

# # Coarse Registration

# +
registered = rigid_warp(fixed_downsample, 
                        t_star, 
                        theta_star, 
                        fixed_com, 
                        moving_downsample.shape)

plot_both(registered, moving_downsample, viewer)
# -

# We need to convert the downsampled transformation into an approprate transformation for the original resolution image. The rotation matrix is scale invariant, but we need to make sure the center of rotation and translation are upsampled by the same amount that we downsampled

print('Converged rigid transformation for downsampled image')
print('Rotation (deg):', theta_star*180/np.pi)
print('Center (px):', fixed_com)
print('Translation (px):', t_star)

# +
fixed_fullres_path = os.path.join(working_dir, 'fixed/zarr_stack/1_1_1')
fixed_fullres_store = zarr.NestedDirectoryStore(fixed_fullres_path)
fixed_fullres = zarr.open(fixed_fullres_store, mode='r')

theta = theta_star
true_factors = np.array(fixed_fullres.shape) / np.array(fixed_downsample.shape)
t, center = coarse._scale_rigid_params(t_star, 
                                       fixed_com, 
                                       true_factors)

print('Converged rigid transformation for original image')
print('Rotation (deg):', theta*180/np.pi)
print('Center (px):', center)
print('Translation (px):', t)
# -

plot_both(fixed_img, moving_img, viewer)

# +
registered = rigid_warp(fixed_img, 
                        t, 
                        theta, 
                        center, 
                        moving_img.shape)

plot_fixed(registered, viewer)
# -

# # Save the transformation

# +
from phathom.utils import pickle_save


transformation_dict = {'t': t,
                       'center': center,
                       'theta': theta}

pickle_save(os.path.join(working_dir, 'rigid_transformation.pkl'), 
            transformation_dict)

# +
from phathom.utils import pickle_load

transformation_dict = pickle_load(os.path.join(working_dir, 'rigid_transformation.pkl'))
# -

transformation_dict



StarcoderdataJupyterScriptsDedupFiltered