def __init__(self, model_path, arch='marvin', imposed_batch_size=1, channels_to_load=(0, 1, 2, 3), normal=False, nb_labels=2, normalize_data=False, normalize_func=None, init_gpu=None): self.imposed_batch_size = imposed_batch_size self.channels_to_load = channels_to_load self.arch = arch self._path = model_path self._fname = os.path.split(model_path)[1] self.nb_labels = nb_labels self.normal = normal self.normalize_data = normalize_data self.normalize_func = normalize_func if init_gpu is None: init_gpu = 'auto' if e2config.device is None: from elektronn2.utils.gpu import initgpu initgpu(init_gpu) elektronn2.logger.setLevel("ERROR") from elektronn2.neuromancer.model import modelload self.model = modelload(model_path, replace_bn='const', imposed_batch_size=imposed_batch_size) self.original_do_rates = self.model.dropout_rates self.model.dropout_rates = ([ 0.0, ] * len(self.original_do_rates))
# "xset q" will always succeed to run if an X server is currently running check_call(['xset', 'q'], stdout=devnull, stderr=devnull) print('X available') # Don't set backend explicitly, use system default... # if "xset q" fails, conclude that X is not running except (OSError, CalledProcessError): print('X unavailable') matplotlib.use('AGG') from elektronn2.config import config config.use_manual_cudnn_conv = True config.use_manual_cudnn_conv_not_w1 = False config.uuse_manual_cudnn_pool = True if config.device is None: from elektronn2.utils.gpu import initgpu initgpu('auto') from elektronn2.neuromancer.model import modelload, kernel_lists_from_node_descr if model_path in [None, 'None']: model_path = "~/axon/mkilling/investigation/MA-TEX/CNN-Timings/old_cnn_rec.mdl" no_gc = False mfp = False model_path = os.path.expanduser(model_path) model_dir, model_name = os.path.split(model_path) os.chdir(model_dir) f_name = model_name[:-4] + '-Speed.csv' # Benchmark Model Compiled with static Shape #
def predict_kzip(kzip_p, m_path, kd_path, clf_thresh=0.5, mfp_active=False, dest_path=None, overwrite=False, gpu_ix=0, imposed_patch_size=None): """ Predicts data contained in k.zip file (defined by bounding box in knossos) Parameters ---------- kzip_p : str path to kzip containing the raw data cube information m_path : str path to predictive model kd_path : str path to knossos dataset clf_thresh : float classification threshold overwrite : bool mfp_active : False imposed_patch_size : tuple dest_path : str path to destination folder, if None folder of k.zip is used. gpu_ix : int """ cube_name = os.path.splitext(os.path.basename(kzip_p))[0] if dest_path is None: dest_path = os.path.dirname(kzip_p) if not os.path.isfile(dest_path + "/%s_data.h5" % cube_name) or overwrite: raw, labels = load_gt_from_kzip(kzip_p, kd_p=kd_path, raw_data_offset=0) raw = xyz2zxy(raw) initgpu(gpu_ix) from elektronn2.neuromancer.model import modelload m = modelload( m_path, imposed_patch_size=list(imposed_patch_size) if isinstance( imposed_patch_size, tuple) else imposed_patch_size, override_mfp_to_active=mfp_active, imposed_batch_size=1) original_do_rates = m.dropout_rates m.dropout_rates = ([ 0.0, ] * len(original_do_rates)) pred = m.predict_dense(raw[None, ], pad_raw=True)[1] # remove area without sufficient FOV pred = zxy2xyz(pred) raw = zxy2xyz(raw) save_to_h5py([pred, raw], dest_path + "/%s_data.h5" % cube_name, ["pred", "raw"]) else: pred, raw = load_from_h5py(dest_path + "/%s_data.h5" % cube_name, hdf5_names=["pred", "raw"]) offset = parse_movement_area_from_zip(kzip_p)[0] overlaycubes2kzip(dest_path + "/%s_pred.k.zip" % cube_name, (pred >= clf_thresh).astype(np.uint32), offset, kd_path)
def _pred_dataset(kd_p, kd_pred_p, cd_p, model_p, imposed_patch_size=None, mfp_active=False, gpu_id=0, overwrite=False, i=None, n=None): """ Helper function for dataset prediction. Runs prediction on whole or partial knossos dataset. Imposed patch size has to be given in Z, X, Y! Parameters ---------- kd_p : str path to knossos dataset .conf file kd_pred_p : str path to the knossos dataset head folder which will contain the prediction cd_p : str destination folder for chunk dataset containing prediction model_p : str path tho ELEKTRONN2 model imposed_patch_size : tuple or None patch size (Z, X, Y) of the model mfp_active : bool activate max-fragment pooling (might be necessary to change patch_size) gpu_id : int the GPU used overwrite : bool True: fresh predictions ; False: earlier prediction continues Returns ------- """ initgpu(gpu_id) from elektronn2.neuromancer.model import modelload kd = KnossosDataset() kd.initialize_from_knossos_path(kd_p, fixed_mag=1) m = modelload(model_p, imposed_patch_size=list(imposed_patch_size) if isinstance( imposed_patch_size, tuple) else imposed_patch_size, override_mfp_to_active=mfp_active, imposed_batch_size=1) original_do_rates = m.dropout_rates m.dropout_rates = ([ 0.0, ] * len(original_do_rates)) offset = m.target_node.shape.offsets offset = np.array([offset[1], offset[2], offset[0]], dtype=np.int) cd = ChunkDataset() cd.initialize(kd, kd.boundary, [512, 512, 256], cd_p, overlap=offset, box_coords=np.zeros(3), fit_box_size=True) ch_dc = cd.chunk_dict print('Total number of chunks for GPU/GPUs:', len(ch_dc.keys())) if i is not None and n is not None: chunks = ch_dc.values()[i::n] else: chunks = ch_dc.values() print("Starting prediction of %d chunks in gpu %d\n" % (len(chunks), gpu_id)) if not overwrite: for chunk in chunks: try: _ = chunk.load_chunk("pred")[0] except Exception as e: chunk_pred(chunk, m) else: for chunk in chunks: try: chunk_pred(chunk, m) except KeyboardInterrupt as e: print("Exiting out from chunk prediction: ", str(e)) return save_dataset(cd) # single gpu processing also exports the cset to kd if n is None: kd_pred = KnossosDataset() kd_pred.initialize_without_conf(kd_pred_p, kd.boundary, kd.scale, kd.experiment_name, mags=[1, 2, 4, 8]) cd.export_cset_to_kd(kd_pred, "pred", ["pred"], [4, 4], as_raw=True, stride=[256, 256, 256])
def predict_h5(h5_path, m_path, clf_thresh=None, mfp_active=False, gpu_ix=0, imposed_patch_size=None, hdf5_data_key=None, data_is_zxy=True, dest_p=None, dest_hdf5_data_key="pred", as_uint8=True): """ Predicts data from h5 file. Assumes raw data is already float32. Parameters ---------- h5_path : str path to h5 containing the raw data m_path : str path to predictive model clf_thresh : float classification threshold, if None, no thresholding mfp_active : False imposed_patch_size : tuple gpu_ix : int hdf5_data_key: str if None, it uses the first entry in the list returned by 'load_from_h5py' data_is_zxy : bool if False, it will assumes data is [X, Y, Z] as_uint8: bool dest_p : str dest_hdf5_data_key : str """ if hdf5_data_key: raw = load_from_h5py(h5_path, hdf5_names=[hdf5_data_key])[0] else: raw = load_from_h5py(h5_path, hdf5_names=None) assert len(raw) == 1, "'hdf5_data_key' not given but multiple hdf5 " \ "elements found. Please define raw data key." raw = raw[0] if not data_is_zxy: raw = xyz2zxy(raw) initgpu(gpu_ix) if raw.dtype.kind in ('u', 'i'): raw = raw.astype(np.float32) / 255. from elektronn2.neuromancer.model import modelload m = modelload(m_path, imposed_patch_size=list(imposed_patch_size) if isinstance( imposed_patch_size, tuple) else imposed_patch_size, override_mfp_to_active=mfp_active, imposed_batch_size=1) original_do_rates = m.dropout_rates m.dropout_rates = ([ 0.0, ] * len(original_do_rates)) pred = m.predict_dense(raw[None, ], pad_raw=True)[1] pred = zxy2xyz(pred) raw = zxy2xyz(raw) if as_uint8: pred = (pred * 255).astype(np.uint8) raw = (raw * 255).astype(np.uint8) if clf_thresh: pred = (pred >= clf_thresh).astype(np.float32) if dest_p is None: dest_p = h5_path[:-3] + "_pred.h5" if hdf5_data_key is None: hdf5_data_key = "raw" save_to_h5py([raw, pred], dest_p, [hdf5_data_key, dest_hdf5_data_key])