SOFIMA-scripts/data-aphid-N-planes-pass2+.py at main · JaneliaSciComp/SOFIMA-scripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import os
import re
import skimage.io as skio
import numpy as np
import tensorstore as ts

def load_data(basepath, z, s):
    za = ts.open({
        'driver': 'zarr',
        'kvstore': {"driver":"file", "path":basepath+'.zarr'},
        'open': True,
        }).result()
    return za[z,:,:].read().result()

def save_flow(flow, basepath, params):
    np.save(os.path.join(basepath, 'flow.'+params+'.npy'), flow)

def load_flow(basepath, params):
    return np.load(os.path.join(basepath, 'flow.'+params+'.npy'))

def save_mesh(mesh, basepath, params):
    np.save(os.path.join(basepath, 'mesh.'+params+'.npy'), mesh)

def load_mesh(basepath, params):
    return np.load(os.path.join(basepath, 'mesh.'+params+'.npy'))

def save_invmap(invmap, basepath, params):
    np.save(os.path.join(basepath, 'invmap.'+params+'.npy'), invmap)

def load_invmap(basepath, params):
    invmap = np.load(os.path.join(basepath, 'invmap.'+params+'.npy'))
    return invmap

def open_warp(shape, chunk_size, basepath, params):
    return ts.open({
        'driver': 'zarr',
        'kvstore': {"driver":"file", "path":os.path.join(basepath, 'warped.'+params+'.zarr')},
        'metadata': {
            "compressor":{"id":"zstd","level":3},
            "shape":shape,
            "chunks":[chunk_size,chunk_size,chunk_size],
            "fill_value":0,
            'dtype': '|u1',
            'dimension_separator': '/',
        },
        'create': True,
        'delete_existing': True,
        }).result()

def write_warp_planes(fid, planes, z0, z1):
    return fid[z0:z1,:,:].write(planes).result()