diff --git a/src/atldld/cli.py b/src/atldld/cli.py index fd2ef58..85d7a7a 100644 --- a/src/atldld/cli.py +++ b/src/atldld/cli.py @@ -15,12 +15,15 @@ # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . """The command line interface (CLI) for Atlas-Download-Tools.""" +import logging from typing import Any, Dict, Optional, Sequence import click import atldld +logger = logging.getLogger(__name__) + @click.group( help="""\ @@ -218,6 +221,252 @@ def dataset_preview(dataset_id, output_dir): click.secho(f"{img_path.resolve().as_uri()}", fg="yellow", bold=True) +@click.command( + "download-faithful", + help=""" + Download all section images of a given dataset and map them into the + reference space. + + More precisely, the section images are mapped into the 25µm version of the + reference space while exactly following the correct image-to-reference + synchronization. Hence the name "faithful", as opposed to a "parallel" + mapping where the section images end up parallel to one of the reference + space axes. + + Because of the faithful mapping and the fact that section images might + present a varying amount of tilt against the slicing axis the mapping into + the reference space will distribute the image data of one image across + different parallel slices of the reference space. + """, +) +@click.argument("dataset_id", type=int) +@click.option( + "--input-downsample", + "downsample_img", + type=int, + default=0, + help=""" + The downsampling factor for input section images. + + This equals the number of times the size of all section images is halved. + In other words, for the downsampling factor N the width and height of the + input images will be reduced by the factor 2^N. + + Setting this to a higher value will make the construction of the volume + faster at the cost of decreasing the image quality. + """, +) +@click.option( + "--output-scale", + "downsample_ref", + type=int, + default=25, + help=""" + The size of the voxels of the output reference volume in micrometres. + """, +) +@click.option( + "-o", + "--output-dir", + type=click.Path(file_okay=False, writable=True, resolve_path=True), + help=""" + The output directory for the volume. If not provided the current + working directory will be used. + """, +) +def download_faithful_dataset(dataset_id, output_dir, downsample_img, downsample_ref): + import pathlib + + from atldld.constants import REF_DIM_1UM + from atldld.utils import get_corners_in_ref_space, get_image + + # Download the dataset metadata + meta = get_dataset_meta_or_abort(dataset_id, include=["section_images"]) + section_thickness = meta["section_thickness"] + section_image_meta = meta.pop("section_images") + + # Download the section images + click.secho("Downloading the section images...", fg="green") + with click.progressbar(section_image_meta) as image_metas: + section_images = { + image_meta["id"]: get_image(image_meta["id"], downsample=downsample_img) + for image_meta in image_metas + } + click.secho(f"Successfully loaded {len(section_images)} section images", fg="green") + + # Map the section images into the reference volume + click.secho("Mapping section images into the reference space volume") + volume = np.zeros(tuple(dim // downsample_ref for dim in REF_DIM_1UM)) + with click.progressbar(section_image_meta) as image_metas: + for image_meta in image_metas: + corners = get_corners_in_ref_space( + image_meta["id"], + image_meta["image_width"], + image_meta["image_height"], + ) + image = section_images[image_meta["id"]] + out, out_bbox = get_true_ref_image( + image, + corners, + section_thickness_um=section_thickness, + downsample_img=downsample_img, + downsample_ref=downsample_ref, + ) + insert_subvolume(volume, out, out_bbox) + + # Save the volume to disk + click.secho("Saving...", fg="green") + volume_file_name = f"dataset-id-{dataset_id}-faithful-downsample-{downsample_img}-scale-{downsample_ref}.npy" + if output_dir is None: + volume_path = pathlib.Path.cwd() / volume_file_name + else: + volume_path = pathlib.Path(output_dir) / volume_file_name + volume_path.parent.mkdir(exist_ok=True, parents=True) + np.save(str(volume_path), volume) + click.secho("Volume was saved in ", fg="green", nl=False) + click.secho(f"{volume_path.resolve().as_uri()}", fg="yellow", bold=True) + + +import numpy as np +from scipy import ndimage +from skimage.color import rgb2gray + + +def get_bbox(points): + """Get the bounding box for a sequence of points. + + Parameters + ---------- + points : np.ndarray or sequence + The points in an d-dimensional space. Shape should be (n_points, d) + + Returns + ------- + np.ndarray + Array with semi-open intervals `[min, max)` for each axis representing + the bounding box for all points. The shape of the array is (2, d). + """ + mins = np.floor(np.min(points, axis=0)) + maxs = np.ceil(np.max(points, axis=0)) + 1 + + return np.stack([mins, maxs]).astype(int) + + +def bbox_meshgrid(bbox): + slices = tuple(slice(start, stop) for start, stop in bbox.T) + + return np.mgrid[slices] + + +def get_true_ref_image( + image, corners, section_thickness_um=25, downsample_img=0, downsample_ref=25 +): + from atldld.maths import find_shearless_3d_affine + + # skip image download and corner queries because it would take too long. + # instead pre-compute them and take them as parameters for now + # image = aibs.get_image(image_id) + # corners = get_ref_corners(image_id, image) + # map corners from the 1µm reference space scale to the given one + corners = corners / downsample_ref + + # compute the affine transformation from the first three corner coordinates + ny, nx = image.shape[:2] + p_to = np.array( + [ + (0, 0, 0), + (nx, 0, 0), + (nx, ny, 0), + ] + ) + p_from = corners[:3] + affine = find_shearless_3d_affine(p_from, p_to) + + # Create the meshgrid + out_bbox = get_bbox(corners) + meshgrid = bbox_meshgrid(out_bbox) + + # Convert the affine transformation to displacements + linear = affine[:3, :3] + translation = affine[:3, 3] + coords = np.tensordot(linear, meshgrid, axes=(1, 0)) + np.expand_dims( + translation, axis=(1, 2, 3) + ) + + # Use the grayscale version for mapping. Originals have white background, + # invert to have black one + # TODO: support RGB + input_img = 1 - rgb2gray(image) + + # Rotate to convert from "ij" coordinates to "xy" + input_img = np.rot90(input_img, 3) + + # Add the z-axis + input_img = np.expand_dims(input_img, axis=2) + + # Correctly take into account the scaling along the slicing axis. The affine + # transformation is computed from the dimensions of the full-resolution + # section images, so the scale along the section axis is the same as along + # the other axes, namely 1µm. We rescale the values for the section axis + # to take into account the true section thickness as well as the + # downsampling of the input images. + coords[2] = coords[2] * 2 ** downsample_img / section_thickness_um + + # Add padding to the z-coordinate to allow for interpolation + pad = 3 + input_img = np.pad(input_img, ((0, 0), (0, 0), (pad, pad))) + coords[2] += pad + + # Apply the transformation + out = ndimage.map_coordinates(input_img, coords, prefilter=False) + + # Rotate to convert from "xy" coordinates to "ij" + out = np.rot90(out) + + # Transpose to get from ipr to pir + out = out.transpose((1, 0, 2)) + + return out, out_bbox + + +def bbox_to_slices(bbox, subset_bbox): + starts = subset_bbox[0] - bbox[0] + ends = subset_bbox[1] - bbox[0] + slices = tuple(slice(start, end) for start, end in zip(starts, ends)) + + return slices + + +def insert_subvolume(volume, subvolume, subvolume_bbox): + # Bounding box of the volume, lower bounds are zeros by definition + volume_bbox = np.stack([(0, 0, 0), volume.shape]) + + # The data bounding box is the intersection of the volume and sub-volume + # bounding boxes + data_bbox = np.stack( + [ + np.max([subvolume_bbox[0], volume_bbox[0]], axis=0), + np.min([subvolume_bbox[1], volume_bbox[1]], axis=0), + ] + ) + if not np.all(data_bbox[1] - data_bbox[0] > 0): + logger.warning( + "The volume and sub-volume don't intersect!\n" + f"Volume bounding box:\n{volume_bbox}\n" + f"Sub-volume bounding box:\n{subvolume_bbox}" + ) + return + + # Convert bounding boxes to array slices + subvolume_slices = bbox_to_slices(subvolume_bbox, data_bbox) + volume_slices = bbox_to_slices(volume_bbox, data_bbox) + + volume[volume_slices] = np.max( + [volume[volume_slices], subvolume[subvolume_slices]], axis=0 + ) + + root.add_command(dataset) dataset.add_command(dataset_info) dataset.add_command(dataset_preview) +dataset.add_command(download_faithful_dataset) diff --git a/src/atldld/maths.py b/src/atldld/maths.py new file mode 100644 index 0000000..e7dff62 --- /dev/null +++ b/src/atldld/maths.py @@ -0,0 +1,71 @@ +"""Mathematical computations.""" +from typing import Sequence + +import numpy as np + + +def find_shearless_3d_affine( + p_from: Sequence[np.ndarray], p_to: Sequence[np.ndarray] +) -> np.ndarray: + """Find a 3D shearless affine transformation given the mapping of 3 points. + + Parameters + ---------- + p_from + Three input points in a 3-dimensional Euclidean space. + p_to + Three output point in a 3-dimensional Euclidean space. + + Returns + ------- + np.ndarray + The affine transform that maps ``p_from`` to ``p_to``. + + With the assumption of no shear we can find the 4th point and its mapping + by the cross product. For example ``p_from = (p1, p2, p3)`` gives two + vectors ``v1 = p2 - p1`` and ``v2 = p3 - p1``, giving ``v3 = v1 x v2``. The + fourth point is then given by ``p4 = p1 + v3``. + + The only caveat is the length of the vector ``v3``. We must make sure that + it scales in the same way between ``p_from`` and ``p_to`` as all other + points. The easiest way to ensure this is to give it the same length as the + ``v1`` vector: ``v3 = |v1| * (v1 x v2 / |v1| / |v2|) = v1 x v2 / |v2|``. + """ + # TODO + # Some sanity checks on the input data: make sure the two triangles + # formed by the input points have the same shape and orientation + + def add_fourth(points): + v1 = points[1] - points[0] + v2 = points[2] - points[0] + v3 = np.cross(v1, v2) / np.linalg.norm(v2) + p4 = points[0] + v3 + + return np.stack([*points, p4]) + + p_from = add_fourth(p_from) + p_to = add_fourth(p_to) + + # check uniform scaling + # from itertools import combinations + # scales = [] + # tolerance = 0.02 + # for i, j in combinations(range(3), 2): + # len_from = np.linalg.norm(p_from[i] - p_from[j]) + # len_to = np.linalg.norm(p_to[i] - p_to[j]) + # scales.append(len_from / len_to) + # print(scales) + # print(np.all([ + # abs(s1 - s2) / min(s1, s2) < tolerance + # for s1, s2 in combinations(scales, 2) + # ])) + + # Add homogenous coordinate and transpose so that + # - dim_0 = "xyz1" + # - dim_1 = points + p_from = np.concatenate([p_from.T, np.array([[1, 1, 1, 1]])]) + p_to = np.concatenate([p_to.T, np.array([[1, 1, 1, 1]])]) + + # Compute the affine transform so that p_to = A @ p_from + # => A = p_to @ p_from_inv + return p_to @ np.linalg.inv(p_from)