Merge branch 'new_optimised_xfm' into 'master'

New optimised xfm

See merge request !3

slider/chart_reg.py

@@ -34,8 +34,8 @@ from skimage import (
     transform,
 )
 from typing import Optional
-
 from slider import util
+from scipy.spatial import ConvexHull
 
 
 def register_chart_to_slide(
@@ -43,11 +43,23 @@ def register_chart_to_slide(
     slide: str,
     slide_res: float,
     outdir: str,
-    boundary_key: Optional[str] = None,
     config: Optional[str] = None,
     do_plots: Optional[bool] = None,
     justify: Optional[str] = None,
+    field: Optional[str] = None,
 ):
+    '''
+    It takes a chart and a slide, and aligns the chart to the slide.
+    
+    Args:
+      chart (str): The path to the chart file.
+      slide (str): The slide to be registered.
+      slide_res (float): The resolution of the slide.
+      outdir (str): The output directory for the results.
+      config (Optional[str]): Path to the configuration file.
+      do_plots (Optional[bool]): Whether to plot the results of the alignment.
+      justify (Optional[str]): str
+    '''
 
     if config is None:
         config = util.get_resource("chart.yaml")
@@ -57,8 +69,8 @@ def register_chart_to_slide(
 
     rlevel = config["slide"]["resolution_level"]
 
-    if boundary_key is None:
-        boundary_key = config["chart"]["boundary_key"]
+    # if boundary_key is None:
+    #     boundary_key = config["chart"]["boundary_key"]
 
     if do_plots is None:
         do_plots = config["general"]["plots"]
@@ -71,15 +83,14 @@ def register_chart_to_slide(
 
     chart = util.Chart.from_neurolucida_ascii(chart)
 
-    outline = chart.get_contours(name=boundary_key, closed=True)
-    if len(outline) < 1:
-        raise ValueError(f"Boundary key {boundary_key} not found in chart")
-    if len(outline) > 1:
-        warnings.warn(f"Repeated entries of boundary key in chart: {boundary_key}")
+    outline = chart.get_contours(closed=True)
 
     edge_crds = [x.points[:, :2] * [1, -1] for x in outline]
     edge_crds_cat = np.concatenate(edge_crds)
 
+    # hull = ConvexHull(edge_crds_cat)
+    # edge_crds_cat = np.concatenate([edge_crds_cat[simplex,:] for simplex in hull.simplices], axis=0)
+
     # load slide, convert to grayscale, and invert if light-field
 
     slide_res = slide_res * (2 ** rlevel)
@@ -89,7 +100,16 @@ def register_chart_to_slide(
 
     img = color.rgb2gray(img_orig)
 
-    field = estimate_field(img)
+    # if exclude_zeros:
+    #     nnz_mask = img != 0
+    # else:
+    #     nnz_mask = None
+
+    # r, c = img.shape
+
+    if field is None:
+        field = estimate_field(img)
+        print(f'Estimated field is: {field}')
 
     if field == "light":
         img = 1 - img
@@ -110,11 +130,13 @@ def register_chart_to_slide(
 
     np.savetxt(f"{outdir}/chart-to-image-init.xfm", init_xfm.params)
 
-    # refine alignment (to mask edges)
-
-    opt_xfm = optimise_xfm(img, slide_res, edge_crds_cat, init_xfm)
-
-    # print(opt_xfm)
+    # optimise alignment
+    # The `optimise_xfm` function takes a mask and a set of coordinates, and returns the optimal
+    # transformation to map the coordinates to the mask.
+    mask = img_props["mask"]
+    opt_xfm, mdist, iter_props = optimise_xfm(
+        mask, slide_res, edge_crds_cat[::3, :], init_xfm
+    )
     tr_x, tr_y = opt_xfm.translation
     print(
         f"Optimised XFM - Rotation: {opt_xfm.rotation:0.5f}, Translation: [{tr_x:0.5f} {tr_y:0.5f}], Scale: {opt_xfm.scale:0.5f}"
@@ -152,7 +174,7 @@ def register_chart_to_slide(
 
     if do_plots:
 
-        fig, ax = plt.subplots(2, 2, figsize=(40, 40))
+        fig, ax = plt.subplots(4, 2, figsize=(20, 30))
         ax = np.ravel(ax)
 
         # chart bounding box
@@ -175,31 +197,100 @@ def register_chart_to_slide(
         )
         plot_box(ax[1], img_props["bbox"])
 
+        # normals (first iteration)
+
+        plot_slide(
+            ax[2], img_props["mask"], slide_res, "Brainmask + normals: first iteration"
+        )
+
+        line_x, line_y = iter_props[0]["normals"]
+        contour_x, contour_y = iter_props[0]["contour_init"]
+        single_crossing_idx = iter_props[0]["single_crossing_idx"]
+
+        plot_normals(
+            ax[2],
+            line_x[~single_crossing_idx, :],
+            line_y[~single_crossing_idx, :],
+            contour_x,
+            contour_y,
+            color="y",
+        )
+        plot_normals(
+            ax[2],
+            line_x[single_crossing_idx, :],
+            line_y[single_crossing_idx, :],
+            contour_x,
+            contour_y,
+            color="b",
+        )
+
+        # normals (last iteration)
+
+        plot_slide(
+            ax[3], img_props["mask"], slide_res, "Brainmask + normals: last iteration"
+        )
+
+        line_x, line_y = iter_props[-1]["normals"]
+        contour_x, contour_y = iter_props[-1]["contour_init"]
+        single_crossing_idx = iter_props[-1]["single_crossing_idx"]
+
+        plot_normals(
+            ax[3],
+            line_x[~single_crossing_idx, :],
+            line_y[~single_crossing_idx, :],
+            contour_x,
+            contour_y,
+            color="y",
+        )
+        plot_normals(
+            ax[3],
+            line_x[single_crossing_idx, :],
+            line_y[single_crossing_idx, :],
+            contour_x,
+            contour_y,
+            color="b",
+        )
+
+        # distance plot
+
+        ax[4].plot(mdist)
+        ax[4].set_xlabel("iteration")
+        ax[4].set_ylabel("mean residual distance")
+
+        ax[5].set_axis_off()
+
         # boundary plots
 
-        plot_slide(ax[2], img, slide_res, title="Boundary alignment")
+        plot_slide(ax[6], img, slide_res, title="Boundary alignment")
 
         for contour in edge_crds:
             plot_contour(
-                ax[2], apply_xfm(init_xfm, contour), color=(1, 0, 0), marker="."
+                ax[6], apply_xfm(init_xfm, contour), color=(1, 0, 0), marker="."
             )
             plot_contour(
-                ax[2], apply_xfm(opt_xfm, contour), color=(0, 1, 0), marker="."
+                ax[6], apply_xfm(opt_xfm, contour), color=(0, 1, 0), marker="."
             )
 
-        ax[2].legend(["boundary_init", "boundary_optimised"])
+        ax[6].legend(["boundary_init", "boundary_optimised"])
 
         # aligned chart
 
-        plot_slide(ax[3], img, slide_res, title="Aligned Chart")
+        plot_slide(ax[7], img, slide_res, title="Aligned Chart")
         cmap = plt.get_cmap("tab10")
 
         for idx, (name, coords, closed) in enumerate(contour_xfm):
             if not closed:
                 continue
-            plot_contour(ax[3], coords, name, color=cmap(idx))
+            plot_contour(ax[7], coords, name, color=cmap(idx))
 
         fig.savefig(f"{outdir}/alignment.png", bbox_inches="tight", dpi=300)
+        # plt.close(fig)
+
+
+def plot_normals(ax, line_x, line_y, contour_x, contour_y, color="b"):
+    for line_x0, line_y0 in zip(line_x, line_y):
+        ax.plot(line_x0, line_y0, ".", color=color)
+    ax.plot(contour_x, contour_y, "r.")
 
 
 def plot_box(ax, bbox, edgecolor="r", facecolor="none", linewidth=1):
@@ -242,13 +333,16 @@ def plot_contour(ax, coords, name=None, color="r", title=None, linewidth=1, **kw
         ax.invert_yaxis()
 
 
-def estimate_field(img):
+def estimate_field(img, mask=None):
     """Estimate field of slide as being light or dark"""
 
-    # calculate mean intensity of the four corners of the image
-    corners = [img[:20, :20], img[:20, -20:], img[-20:, :20], img[-20:, -20:]]
+    if mask is None:
+        edges = np.zeros(img.shape, dtype=bool)
+        edges[:, 0:2], edges[:, -2:], edges[0:2, :], edges[-2:, :] = True, True, True, True
+    else:
+        edges = mask ^ morphology.binary_erosion(mask, selem=morphology.disk(10))
 
-    corners = np.median(np.concatenate(corners, axis=0))
+    edges_m = np.median(img[edges])
 
     # calculate mean intensity of the centre of the image
     h, w = img.shape
@@ -256,10 +350,12 @@ def estimate_field(img):
     w = w // 2
 
     centre = img[h - 100 : h + 100, w - 100 : w + 100]
-    centre = np.mean(centre)
+    centre_m = np.median(centre)
+
+    # print(edges_m, centre_m)
 
     # estimate field (light or dark)
-    if corners > centre:
+    if edges_m > centre_m:
         return "light"
     else:
         return "dark"
@@ -316,7 +412,42 @@ def segment_foreground(
     return brainmask
 
 
-def optimise_xfm(image, image_resolution, contour, xfm_init):
+def optimise_xfm(mask, mask_resolution, contour, xfm, n_iter=100):
+    '''
+    Given a mask, a contour, and an initial transformation, optimise_xfm() iteratively optimises the
+    transformation by finding the best fit between the contour and the mask, and then returns the best 
+    fit transformation.
+    
+    Args:
+      mask: the binary mask of the structure of interest
+      mask_resolution: The resolution of the mask image.
+      contour: the contour to be transformed
+      xfm: the initial transform
+      n_iter: number of iterations to run the optimisation for. Defaults to 100
+    
+    Returns:
+      The optimise_xfm function returns the optimised xfm, the mean distance between the optimised xfm
+      and the contour, and the optimised xfm and the contour at each iteration.
+    '''
+
+    mdist = np.zeros(n_iter)
+    iter_props = [None] * n_iter
+
+    # TODO: move line extent values to config
+    max_line_extent = 1.5
+    min_line_extent = 0.05
+
+    for idx in range(n_iter):
+
+        xfm, mdist[idx], iter_props[idx] = optimise_xfm_worker(
+            mask, mask_resolution, contour, xfm, line_extent=max_line_extent
+        )
+        max_line_extent = np.maximum(max_line_extent * 0.9, min_line_extent)
+
+    return xfm, mdist, iter_props
+
+
+def optimise_xfm_worker(image, image_resolution, contour, xfm_init, line_extent=1.5):
     """
     Refine contour by sampling image along normal (to each edge point) and looking for big step change.
     Calculate a new transform (xfm) using refined contour points.
@@ -332,30 +463,13 @@ def optimise_xfm(image, image_resolution, contour, xfm_init):
 
     # TODO: move these line extents to the config file
     # line_smpl = np.linspace(-0.03, 0.15, 20)
-    line_smpl = np.linspace(-0.2, 0.2, 20)
+    line_smpl = np.linspace(-line_extent, line_extent, 1000)
 
     line_x = edge_init_x[:, np.newaxis] + nrml_x[:, np.newaxis] * line_smpl
     line_y = edge_init_y[:, np.newaxis] + nrml_y[:, np.newaxis] * line_smpl
 
-    brainmask = segment_foreground(image)
-
-    # if do_plots:
-    #     fig, ax = plt.subplots(figsize=(20, 30))
-
-    #     extent = np.array([0, img.shape[1], img.shape[0], 0]) * img_res
-    #     ax.imshow(brainmask, extent=extent, cmap='gray')
-
-    #     for line_x0, line_y0 in zip(line_x, line_y):
-    #         ax.plot(line_x0, line_y0, "b.")
-
-    #     ax.plot(edge_init_x, edge_init_y, "r.")
-
-    #     ax.set_xlabel("mm")
-    #     ax.set_ylabel("mm")
-    #     ax.set_title("Brainmask + normals")
-
-    #     plt.show()
-    #     # fig.savefig(f"{OUTDIR}/normals.png", bbox_inches="tight", dpi=300)
+    # brainmask = segment_foreground(image)
+    brainmask = image
 
     # sample image along normal line
 
@@ -365,8 +479,8 @@ def optimise_xfm(image, image_resolution, contour, xfm_init):
     y, x = brainmask.shape
     line_int = brainmask[np.clip(line_y, 0, y - 1), np.clip(line_x, 0, x - 1)]
 
-    # exlude constant rows
-    constant_idx = np.all(line_int == line_int[:, 0][:, np.newaxis], axis=1)
+    # exclude normals that cross the mask edge twice
+    single_crossing_idx = np.sum(np.abs(np.diff(line_int, axis=1)), axis=1) == 1
 
     min_idx = np.argmax(np.abs(np.diff(line_int, axis=-1)), axis=-1)
 
@@ -380,10 +494,32 @@ def optimise_xfm(image, image_resolution, contour, xfm_init):
 
     opt_xfm = transform.SimilarityTransform()
     opt_xfm.estimate(
-        contour[~constant_idx, :], refined_edge_coords[~constant_idx, :]
+        contour[single_crossing_idx, :], refined_edge_coords[single_crossing_idx, :]
+    )
+
+    # calculate error (mean distance)
+
+    opt_edge_coords = apply_xfm(opt_xfm, contour)
+
+    dist = np.sqrt(
+        np.sum(
+            (
+                opt_edge_coords[single_crossing_idx, :]
+                - refined_edge_coords[single_crossing_idx, :]
+            )
+            ** 2,
+            axis=1,
+        )
     )
+    mdist = np.mean(dist)
+
+    iter_props = {
+        "single_crossing_idx": single_crossing_idx,
+        "normals": (line_x * image_resolution, line_y * image_resolution),
+        "contour_init": (edge_init_x, edge_init_y),
+    }
 
-    return opt_xfm
+    return opt_xfm, mdist, iter_props
 
 
 def image_props(img, img_resolution):
@@ -442,46 +578,104 @@ def justify_bbox(bbox, aspect_ratio, justify):
 
     new_w = w * scale_factor
 
+    new_bbox = bbox.copy()
+
     if justify == "right":
-        bbox[1] += w - new_w
+        new_bbox[1] += w - new_w
     elif justify == "left":
-        bbox[3] -= w - new_w
+        new_bbox[3] -= w - new_w
     else:
         raise ValueError(
             f'Unknown value for justify ("{justify}"). Can only be "left" or "right".'
         )
 
-    return bbox
+    return new_bbox
 
 
 def initialise_xfm(image, image_resolution, contour, tol=0.05, justify=None):
-    '''Calculate image and contour bounding boxes, then estimate the transform (xfm) based on bounding box corners.'''
+    '''
+    Calculate the transform (xfm) based on bounding box corners
+    
+    Args:
+      image: the image to be transformed
+      image_resolution: The resolution of the image.
+      contour: the contour to be transformed
+      tol: tolerance for aspect ratio mismatch between slide and contour
+      justify: str
+    
+    Returns:
+      the transform (xfm), the image properties (image_p) and the contour properties (contour_p).
+    '''
+
 
     brainmask = segment_foreground(image)
 
     image_p = image_props(brainmask, image_resolution)
     contour_p = point_props(contour)
 
+    # If the aspect-ratio of the image and contour bounding boxes are not equal it typically means that the image is bilateral whilst the contour is lateralised.
+    # Thust the contour bbox should left or right justified within the image bbox.
+
+    def aspect_is_equal(x, y):
+        return np.abs(x["aspect_ratio"] - y["aspect_ratio"]) < tol
+
+    if (not aspect_is_equal(image_p, contour_p)) & (justify is None):
+
+        # The justification is not set by the user, will attempt left and right and see which has best fit.
+
+        warnings.warn(
+            "Slide and chart aspect ratios appear to be different, and direction of justification is not set.  Will attempt to estimate direction."
+        )
+
+        bbox_idx = np.array([[1, 2], [1, 0], [3, 2], [3, 0]])
+        src = np.array(contour_p["bbox"])[bbox_idx]
+
+        goodness_fit = np.zeros(2)
+
+        directions = ("left", "right")
+
+        for dir_idx, dir in enumerate(directions):
+            bbox = justify_bbox(image_p["bbox"], contour_p["aspect_ratio"], dir)
+            dest = np.array(bbox)[bbox_idx]
+
+            xfm = transform.SimilarityTransform()
+            xfm.estimate(src, dest)
+
+            # calc goodness of fit (number of single-edge-crossing normals)
+            # fit is determined by the number of contour normals that cross the image mask boundary only once
+
+            single_crossing_idx = optimise_xfm_worker(
+                brainmask, image_resolution, contour, xfm, line_extent=1
+            )[2]["single_crossing_idx"]
+            goodness_fit[dir_idx] = np.sum(single_crossing_idx)
+
+        # maximise goodness-of-fit
+        justify = directions[np.argmax(goodness_fit)]
+
+        print(
+            f"Estimated that the contour should be {justify} justified to the image: {goodness_fit}."
+        )
+
     if justify is not None:
 
         # force image bbox aspect ratio to be the same as contour bbox, and justify left or right.
 
         bbox = justify_bbox(image_p["bbox"], contour_p["aspect_ratio"], justify)
 
-        # exclude voxels from brainmask that are outside the bounding box
+        # # exclude voxels from brainmask that are outside the bounding box
 
-        bbox = np.round(np.array(bbox) / image_resolution).astype(int)
+        # bbox = np.round(np.array(bbox) / image_resolution).astype(int)
 
-        cropped_brainmask = brainmask.copy()
-        cropped_brainmask[:, : bbox[1]] = 0
-        cropped_brainmask[:, bbox[3] :] = 0
+        # cropped_brainmask = brainmask.copy()
+        # cropped_brainmask[:, : bbox[1]] = 0
+        # cropped_brainmask[:, bbox[3] :] = 0
 
-        image_p = image_props(cropped_brainmask, image_resolution)
+        # image_p = image_props(cropped_brainmask, image_resolution)
 
-        # recalculate justified bounding box on cropped brainmask
-        # force image bbox aspect ratio to be the same as contour bbox, and justify left or right.
+        # # recalculate justified bounding box on cropped brainmask
+        # # force image bbox aspect ratio to be the same as contour bbox, and justify left or right.
 
-        bbox = justify_bbox(image_p["bbox"], contour_p["aspect_ratio"], justify)
+        # bbox = justify_bbox(image_p["bbox"], contour_p["aspect_ratio"], justify)
 
         centroid = (
             (bbox[0] + bbox[2]) / 2,
@@ -496,8 +690,10 @@ def initialise_xfm(image, image_resolution, contour, tol=0.05, justify=None):
             "mask": brainmask,
         }
 
-    if np.abs(image_p["aspect_ratio"] - contour_p["aspect_ratio"]) > tol:
-        raise RuntimeError("Slide and chart aspect ratios appear to be different.  Consider correcting aspect-ratio with left/right justification.")
+    if not aspect_is_equal(image_p, contour_p):
+        raise RuntimeError(
+            "Slide and chart aspect ratios appear to be different.  Consider correcting aspect-ratio with left/right justification."
+        )
 
     idx = np.array([[1, 2], [1, 0], [3, 2], [3, 0]])
 

slider/util.py

@@ -21,6 +21,11 @@ from typing import Optional, List
 import numpy as np
 from slider.external import neurolucida
 import glymur
+import collections
+import six
+
+def is_nonstring_iterable(arg):
+    return (isinstance(arg, collections.Iterable) and not isinstance(arg, six.string_types))
 
 def get_slider_dir() -> str:
     rpath = op.realpath(op.dirname(op.abspath(__file__)))