N-planes-invmap.py

#!/usr/bin/env python

# a derivative of https://github.com/google-research/sofima/blob/main/notebooks/em_alignment.ipynb

# this third part only uses the CPU (and also a lot of RAM).  it depends on the
# output of N-planes-mesh.py

import sys
import os
import argparse
import importlib

from connectomics.common import bounding_box
from sofima import map_utils
from datetime import datetime 

# Parse command line arguments
parser = argparse.ArgumentParser(
    description="inverts the coordinate map of the mesh"
)
parser.add_argument(
    "data_loader",
    help="Data loader module name, e.g., data-test-2-planes"
)
parser.add_argument(
    "basepath",
    help="filepath to stitched planes"
)
parser.add_argument(
    "min_z",
    type=int,
    help="lower bound on the planes to align"
)
parser.add_argument(
    "max_z",
    type=int,
    help="upper bound on the planes to align"
)
parser.add_argument(
    "patch_size",
    type=str,
    help="Side length of (square) patch for processing (in pixels, e.g., 32)",
)
parser.add_argument(
    "stride",
    type=str,
    help="Distance of adjacent patches (in pixels, e.g., 8)"
)
parser.add_argument(
    "scales",
    help="the spatial resolutions to use when computing the flow field"
)
parser.add_argument(
    "k0",
    type=float,
    help="spring constant for inter-section springs"
)
parser.add_argument(
    "k",
    type=float,
    help="spring constant for intra-section springs"
)
parser.add_argument(
    "parallelism",
    type=int,
    help="how many processes/threads to use"
)
parser.add_argument(
    "write_metadata",
    type=int,
    help="whether to write the zarr metadata for not"
)
parser.add_argument(
    "X",
    type=int,
    help=""
)

args = parser.parse_args()

data_loader = args.data_loader
basepath = args.basepath
min_z = args.min_z
max_z = args.max_z
patch_size = args.patch_size
stride = args.stride
k0 = args.k0
k = args.k
parallelism = args.parallelism
write_metadata = args.write_metadata
X = args.X

stride_int_min = [int(x) for x in args.stride.split(',')][-1]
scales_int = [int(x) for x in args.scales.split(',')]
X = X==1

print("data_loader =", data_loader)
print("basepath =", basepath)
print("min_z =", min_z)
print("max_z =", max_z)
print("patch_size =", patch_size)
print("stride =", stride)
print("scales =", scales_int)
print("k0 =", k0)
print("k =", k)
print("parallelism =", parallelism)
print("write_metadata =", write_metadata)
print("X =", X)

data = importlib.import_module(os.path.basename(data_loader))

params = 'patch'+patch_size+'.stride'+stride+'.scales'+args.scales.replace(",",'')+'.k0'+str(k0)+'.k'+str(k)

print(datetime.now(), 'loading mesh')
mesh = data.load_mesh(basepath, params, min_z+1, max_z+1, "Xfine" if X else "")

boxMx = bounding_box.BoundingBox(start=(0, 0, 0), size=(mesh.shape[-1], mesh.shape[-2], 1))

s_min = min(scales_int)
stride_scale_min = stride_int_min * (2**s_min)

print(datetime.now(), 'creating inverted map')
fid = data.create_invmap(mesh.shape, basepath, params, write_metadata, X)

print(datetime.now(), 'inverting map')
invmap = map_utils.invert_map(mesh, boxMx, boxMx, stride_scale_min,
                              parallelism=parallelism, verbose=True)

print(datetime.now(), 'saving inverted map')
for z in range(min_z+1, max_z+2):
    data.write_invmap_plane(fid, invmap[:, z-min_z-1 : z-min_z, ...], z)