126 lines
4.9 KiB
Python
126 lines
4.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Created on Tue Jun 20 13:37:40 2023
|
|
|
|
extracts numpy boolean array from nc and stores as npy
|
|
strips all metadata, but allows to process volume data within blender using pyopenvdb
|
|
|
|
@author: fische_r
|
|
"""
|
|
|
|
import xarray as xr
|
|
import os
|
|
from skimage import morphology
|
|
from skimage.morphology import ball
|
|
import numpy as np
|
|
from joblib import delayed, Parallel
|
|
import socket
|
|
from scipy import ndimage
|
|
host = socket.gethostname()
|
|
import argparse
|
|
|
|
# TODO: make this more elegant
|
|
try:
|
|
import cupy as cp
|
|
import cucim
|
|
GPU_avail = True
|
|
except:
|
|
print("cupy or cucim not found, can't use GPU")
|
|
GPU_avail = False
|
|
|
|
# default parameters
|
|
|
|
# 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
|
|
|
|
num_GPU = 1
|
|
n_jobs = 4
|
|
if host == 'mpc2959.psi.ch':
|
|
num_GPU = 5
|
|
n_jobs = 16
|
|
|
|
|
|
|
|
class volume_maker:
|
|
def __init__(self,
|
|
args = None):
|
|
self.clean = args.clean_image
|
|
self.remove_small = args.remove_small
|
|
self.footprint = args.footprint
|
|
self.ts = args.time_step
|
|
self.minsize = args.minsize
|
|
|
|
self.topoutfolder = args.output_path
|
|
self.array_name = args.segmented_name
|
|
self.ph = args.phase
|
|
|
|
|
|
|
|
def clean_binary_image(self, im, clean = True, remove_small=True, minsize = 200, fp_radius = 1, i=0, GPU = True, GPU_avail=GPU_avail):
|
|
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
|
|
|
|
def xarray_to_npy(self, im, ts, i=0):
|
|
im = im == self.ph
|
|
im = self.clean_binary_image(im, i=i, clean = self.clean , remove_small=self.remove_small, minsize = self.minsize, fp_radius = self.footprint)
|
|
|
|
outpath = os.path.join(self.topoutfolder, self.array_name+'_phase_'+str(self.ph)+'_ts_'+f'{ts:04d}'+'.npy')
|
|
np.save(outpath, im)
|
|
|
|
def nc_to_set_of_npy(self):
|
|
imdata = self.data[self.array_name]
|
|
if self.ts<0:
|
|
print('processing all time steps')
|
|
timesteps = imdata.timestep.data
|
|
length = len(timesteps)
|
|
|
|
Parallel(n_jobs=n_jobs, temp_folder=temp_folder)(delayed(self.xarray_to_npy)(imdata.sel(timestep=timesteps[i]).data, timesteps[i], i) for i in range(length))
|
|
|
|
else:
|
|
print('processing time step ',str(self.ts))
|
|
self.xarray_to_npy(imdata.sel(timestep=self.ts).data, self.ts)
|
|
|
|
if __name__ == '__main__':
|
|
### Parse arguments
|
|
parser = argparse.ArgumentParser(description='preprocessing parameters')
|
|
parser.add_argument('-o', '--output_path', type = str, default = '', help = 'path to folder that will contain the .npy files')
|
|
parser.add_argument('-i', '--input_path', type=str, default='', help = 'path to the .nc-file containing the segmented data')
|
|
parser.add_argument('-cl', '--clean_image', type = bool, default = True, help = 'wheter to remove spurious pixel by binaray opening')
|
|
parser.add_argument('-fp', '--footprint', type = int, default = 1, help = 'radius of the sphere stencel used for the binary opening')
|
|
parser.add_argument('-rs', '--remove_small', type = bool, default=True, help = 'wheter to remove objects smaller than minsize')
|
|
parser.add_argument('-ms', '--minsize', type = int, default=20, help = 'minimum size of connected objects to keep')
|
|
parser.add_argument('-ts', '--time_step', type = int, default=0, help = 'time step that is processed, -1 for all')
|
|
parser.add_argument('-sn', '--segmented_name', type = str, default = 'segmented', help = 'name of the data array in the .nc')
|
|
parser.add_argument('-ph', '--phase', type = int, default = 1, help='which phase to extract, -1 for all except 0 (not yet implemented)')
|
|
|
|
args = parser.parse_args()
|
|
|
|
VM = volume_maker(args)
|
|
|
|
# lazy load the data
|
|
VM.data = xr.open_dataset(args.input_path)
|
|
|
|
#create outfolder if necessary
|
|
if not os.path.exists(VM.topoutfolder):
|
|
os.mkdir(VM.topoutfolder)
|
|
|
|
# process the data with the given parameters
|
|
VM.nc_to_set_of_npy()
|
|
|