215 lines
7.6 KiB
Python
215 lines
7.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Created on Fri Mar 25 08:39:28 2022
|
|
|
|
@author: firo
|
|
"""
|
|
|
|
import xarray as xr
|
|
import os
|
|
from skimage import measure
|
|
from skimage import morphology
|
|
from skimage.morphology import ball
|
|
import trimesh
|
|
import numpy as np
|
|
from joblib import delayed, Parallel
|
|
import socket
|
|
from scipy import ndimage
|
|
host = socket.gethostname()
|
|
|
|
|
|
|
|
# TODO: make this more elegant
|
|
try:
|
|
import cupy as cp
|
|
import cucim
|
|
GPU_avail = True
|
|
except:
|
|
print('cupy and cucim not found')
|
|
GPU_avail = False
|
|
|
|
|
|
# TODO: add GPU support once cucim/cupy allow GPU marching cubes
|
|
# TODO: already allow cupy for matrix operations and loading, image cleaning already possible with GPU
|
|
temp_folder = None
|
|
if host == 'DDM06609':
|
|
temp_folder = r"Z:\users\firo\joblib_tmp"
|
|
|
|
num_GPU = 1
|
|
if host == 'mpc2959.psi.ch':
|
|
num_GPU = 5
|
|
|
|
|
|
|
|
def clean_binary_image(im, clean = True, remove_small=True, minsize = 200, fp_radius = 1, i=0, GPU = True, GPU_avail=GPU_avail):
|
|
# GPU?
|
|
if clean or remove_small:
|
|
if GPU and GPU_avail:
|
|
gpu_id = i%num_GPU #use gpus 1 through 4, leaving the big A40 (0) alone or i%5 to use all 5
|
|
|
|
with cp.cuda.Device(gpu_id):
|
|
im = cp.array(im)
|
|
if clean: im = cucim.skimage.morphology.binary_opening(im, footprint=cucim.skimage.morphology.ball(fp_radius))
|
|
if remove_small: im = cucim.skimage.morphology.remove_small_objects(im, min_size=minsize)
|
|
im = cp.asnumpy(im)
|
|
mempool = cp.get_default_memory_pool()
|
|
mempool.free_all_blocks()
|
|
|
|
else:
|
|
if clean: im = ndimage.binary_opening(im, structure=ball(fp_radius))
|
|
if remove_small: im = morphology.remove_small_objects(im, min_size=minsize)
|
|
|
|
return im
|
|
|
|
|
|
class mesh_maker:
|
|
def __init__(self,
|
|
nc_path = None,
|
|
out_path = None,
|
|
ref_timestep = 0,
|
|
timesteps = None, #'all', or array
|
|
x_boundaries = (0,-1),
|
|
y_boundaries = (0,-1),
|
|
z_boundaries = (0,-1),
|
|
use_gpu = False, #stub to call either cupy/cucim or not,
|
|
cpu_parallel = True,
|
|
njobs = 16,
|
|
GPU = True
|
|
):
|
|
self.nc_path = nc_path
|
|
self.out_path = out_path
|
|
|
|
self.xbot = x_boundaries[0]
|
|
self.xtop = x_boundaries[1]
|
|
self.ybot = y_boundaries[0]
|
|
self.ytop = y_boundaries[1]
|
|
self.zbot = z_boundaries[0]
|
|
self.ztop = z_boundaries[1]
|
|
|
|
self.ref_ts = ref_timestep
|
|
self.timesteps = timesteps
|
|
self.use_gpu = use_gpu
|
|
self.parallel = cpu_parallel
|
|
self.njobs = njobs
|
|
self.GPU = GPU
|
|
|
|
def load_data(self):
|
|
self.dyn_data = xr.open_dataset(self.nc_path)
|
|
if 'transition_matrix' in list(self.dyn_data.keys()):
|
|
self.watermask = (self.dyn_data['transition_matrix']>0).data[self.xbot:self.xtop,self.ybot:self.ytop,self.zbot:self.ztop]
|
|
if not os.path.exists(self.out_path):
|
|
os.mkdir(self.out_path)
|
|
|
|
|
|
def time_step_stl(self, ts):
|
|
|
|
water = self.dyn_data['transition_matrix']<ts
|
|
water = np.bitwise_and(water[self.xbot:self.xtop,self.ybot:self.ytop,self.zbot:self.ztop], self.watermask)
|
|
|
|
if np.any(water):
|
|
water[:,:,:2] = 0
|
|
water[:,:,-3:] = 0
|
|
verts, faces, _, _ = measure.marching_cubes(water.data) #_lewiner
|
|
watermesh = trimesh.Trimesh(vertices = verts, faces = faces)
|
|
stl = trimesh.exchange.stl.export_stl_ascii(watermesh)
|
|
stlpath = os.path.join(self.out_path, ''.join(['water_ts_',f'{ts:05}','.stl']))
|
|
|
|
with open(stlpath, 'w+') as file:
|
|
file.write(stl)
|
|
|
|
def time_4D_stl(self, ts, phase=1, phasename = 'phase_1', clean=False, remove_small= False, fp_radius = 1, minsize=20):
|
|
|
|
# im = self.dyn_data['segmented'].sel(time=ts)[self.xbot:self.xtop,self.ybot:self.ytop,self.zbot:self.ztop]==phase
|
|
im = self.dyn_data['segmented'].sel(timestep=ts)[self.xbot:self.xtop,self.ybot:self.ytop,self.zbot:self.ztop]==phase
|
|
|
|
if np.any(im):
|
|
im = im.data
|
|
# TODO: use GPU
|
|
# if clean: im = ndimage.binary_opening(im, structure=fp)
|
|
# if remove_small: im = morphology.remove_small_objects(im, min_size=minsize)
|
|
|
|
#crop boundaries
|
|
im[:,:2,:] = 0
|
|
im[:,-3:,:] = 0
|
|
im[:2,:,:] = 0
|
|
im[-3:,:,:] = 0
|
|
im[:,:,:2] = 0
|
|
im[:,:,-3:] = 0
|
|
|
|
im = clean_binary_image(im, clean, remove_small, minsize, fp_radius, i=ts, GPU = self.GPU)
|
|
|
|
verts, faces, _, _ = measure.marching_cubes(im) #_lewiner
|
|
mesh = trimesh.Trimesh(vertices = verts, faces = faces)
|
|
stl = trimesh.exchange.stl.export_stl_ascii(mesh)
|
|
stlpath = os.path.join(self.out_path, ''.join([phasename, f'{ts:05}','.stl']))
|
|
with open(stlpath, 'w+') as file:
|
|
file.write(stl)
|
|
|
|
|
|
def run(self):
|
|
# laod the data
|
|
self.load_data()
|
|
# TODO:close ncfile if error
|
|
# make reference time step stl
|
|
try:
|
|
self.time_step_stl(self.ref_ts)
|
|
|
|
# make all the other stls if called
|
|
if self.timesteps is not None:
|
|
if self.timesteps == 'all':
|
|
steps = np.arange(len(self.dyn_data['time']))
|
|
else:
|
|
steps = self.timesteps
|
|
|
|
if self.parallel:
|
|
Parallel(n_jobs=self.njobs, temp_folder=temp_folder)(delayed(self.time_step_stl)(ts) for ts in steps)
|
|
else:
|
|
for ts in steps:
|
|
if not ts == self.ref_ts:
|
|
self.time_step_stl(ts)
|
|
except Exception as e:
|
|
print(e)
|
|
self.dyn_data.close()
|
|
self.dyn_data.close()
|
|
|
|
def run2(self, phase, name, clean=False, remove_small=False, fp_radius=1, minsize=20):
|
|
# laod the data
|
|
self.load_data()
|
|
# TODO:close ncfile if error
|
|
# make reference time step stl
|
|
try:
|
|
self.time_4D_stl(self.ref_ts, phase, name, clean)
|
|
|
|
# make all the other stls if called
|
|
if self.timesteps is not None:
|
|
if self.timesteps == 'all':
|
|
# steps = np.arange(len(self.dyn_data['time']))
|
|
steps = self.dyn_data['time'].data
|
|
steps = self.dyn_data['time_step'].data
|
|
else:
|
|
steps = self.timesteps
|
|
|
|
if self.parallel:
|
|
Parallel(n_jobs=self.njobs, temp_folder=temp_folder)(delayed(self.time_4D_stl)(ts, phase, name, clean,remove_small, fp_radius, minsize) for ts in steps)
|
|
else:
|
|
for ts in steps:
|
|
if not ts == self.ref_ts:
|
|
self.time_4D_stl(ts, phase, name, clean)
|
|
except Exception as e:
|
|
print(e)
|
|
self.dyn_data.close()
|
|
self.dyn_data.close()
|
|
|
|
|
|
|
|
ncpath = '/das/home/fische_r/DASCOELY/processing/04_membrane_ML/3II/3IImembrane_segmentation.nc'
|
|
outpath = '/das/home/fische_r/DASCOELY/processing/04_membrane_ML/3II/crack'
|
|
MM = mesh_maker(nc_path=ncpath, out_path=outpath)
|
|
MM.load_data()
|
|
MM.run2(2, 'crack', clean=True, remove_small=True)
|
|
|
|
|
|
|
|
|
|
|