예제 #1
0
def main(pred_file,
         result_folder='.',
         pred_keys=['volumes/pred_affs'],
         result_key='volumes/instances',
         patchshape=[7, 7, 7],
         chunksize=[92, 92, 92],
         num_workers=8,
         remove_small_comps=0,
         save_mip=False,
         **kwargs):
    sample = os.path.basename(pred_file).split('.')[0]
    result_file = os.path.join(result_folder, sample + '.zarr')
    tmp_key = 'volumes/tmp_worker'

    if num_workers == 1:
        single_worker(pred_file, result_file, num_workers, **kwargs)

    else:
        non_overlapping_chessboard_worker(pred_file, result_file, pred_keys,
                                          tmp_key, num_workers, chunksize,
                                          patchshape, **kwargs,
                                          **{'result_folder': result_folder})
        # stitch blocks
        instances = stitch_vote_instances(result_file, tmp_key, result_key,
                                          chunksize, patchshape)
        # convert result to hdf
        if kwargs['output_format'] == 'hdf':
            zarr2hdf(result_file,
                     hdf_file=os.path.join(result_folder, sample + '.hdf'),
                     zarr_key=['volumes/instances'],
                     hdf_key=['volumes/instances'])

        if remove_small_comps > 0:

            instances = remove_small_components(instances, remove_small_comps)
            instances = relabel(instances)
            io.imsave(os.path.join(result_folder, sample + '.tif'),
                      instances.astype(np.uint16),
                      plugin='tifffile')

        if save_mip:
            colored = color(np.max(instances, axis=0))
            io.imsave(os.path.join(result_folder, sample + '.png'),
                      colored.astype(np.uint8))
예제 #2
0
            single_worker(pred_file, out_file, args.num_workers)

        else:
            non_overlapping_chessboard_worker(pred_file, out_file, args.in_key,
                                              tmp_key, args.num_workers,
                                              args.chunksize, args.patchshape,
                                              **kwargs)

            instances = stitch_vote_instances(out_file, tmp_key, args.out_key,
                                              args.chunksize, args.patchshape)

            if args.remove_small_comps > 0:
                cleaned = remove_small_components(instances,
                                                  args.remove_small_comps)
                relabeled = relabel(cleaned)
                io.imsave(out_file.replace('.zarr', '.tif'),
                          relabeled.astype(np.uint16),
                          plugin='tifffile')

            if args.save_mip:
                colored = color(np.max(relabeled, axis=0))
                io.imsave(out_file.replace('.zarr', '_mip.png'),
                          colored.astype(np.uint8))

            zarr2hdf(out_file,
                     hdf_file=out_file.replace('.zarr', '.hdf'),
                     zarr_key=['volumes/instances'],
                     hdf_key=['volumes/instances'])

        logger.info("%s", datetime.now() - t1)
def main(pred_file, result_folder='.', **kwargs):
    """Calls vote_instances blockwise and stitches them afterwards.

    Args:

        pred_file (``string``):

            Filename of prediction. Should be zarr.

        result_folder (``string``):

            Path to result folder.

        **kwargs (``dict``):

            All arguments needed for vote_instances and stitching. In the
            following only stitching arguments are listed: (maybe they should be
            renamed?)

            aff_key
            fg_key
            res_key
            overlapping_inst
            only_bb
            patchshape
            patch_threshold
            ignore_small_comps
            skeletonize_foreground
            chunksize
            num_parallel_blocks

            OUTPUT / POSTPROCESSING:
            save_mip
            remove_small_comps
            dilate_instances

    """
    assert os.path.exists(pred_file), \
        'Prediction file {} does not exist. Please check!'.format(pred_file)

    sample = os.path.basename(pred_file).split('.')[0]
    result_file = os.path.join(result_folder, sample + '.zarr')
    kwargs['result_folder'] = result_folder

    kwargs['return_intermediates'] = True

    # dataset keys
    aff_key = kwargs['aff_key']
    fg_key = kwargs.get('fg_key')
    res_key = kwargs.get('res_key', 'vote_instances')
    cleaned_mask_key = 'volumes/foreground_cleaned'
    tmp_key = 'volumes/blocks'

    # read input shape
    if pred_file.endswith('.zarr'):
        in_f = zarr.open(pred_file, mode='a')
    else:
        raise NotImplementedError
    aff_shape = in_f[aff_key].shape
    channel_order = [slice(0, aff_shape[0])]
    pred_keys = [aff_key]

    if kwargs['overlapping_inst']:
        numinst_shape = in_f[fg_key].shape
        channel_order.append(slice(0, numinst_shape[0]))
        pred_keys += [fg_key]
        assert aff_shape[1:] == numinst_shape[1:], \
            'Please check: affinity and numinst shape do not match!'
    input_shape = aff_shape[1:]  # without first channel dimension

    # get bounding box for foreground
    if kwargs.get('only_bb'):
        mid = np.prod(kwargs['patchshape']) // 2
        if cleaned_mask_key in in_f:
            pred_keys += [cleaned_mask_key]
            channel_order.append(slice(0, 1))
            shape = np.array(in_f[cleaned_mask_key].attrs['fg_shape'])
            bb_offset = np.array(in_f[cleaned_mask_key].attrs['offset'])
        else:
            mask = np.array(in_f[aff_key][mid])
            mask = mask > kwargs['patch_threshold']
            if np.sum(mask) == 0:
                logger.info('Volume has no foreground voxel, returning...')
                return
            # remove small components
            if kwargs.get('ignore_small_comps', 0) > 0:
                mask = clean_mask(mask, np.ones([3] * mask.ndim),
                                  kwargs.get('ignore_small_comps'))
                mask = mask.astype(np.uint8)
            # skeletonize mask (flylight specific)
            if kwargs.get('skeletonize_foreground'):
                mask = skeletonize_3d(mask) > 0
                mask = mask.astype(np.uint8)
            # save processed mask to input file
            if kwargs.get('ignore_small_comps',
                          0) > 0 or kwargs.get('skeletonize_foreground'):
                in_f.create_dataset(cleaned_mask_key,
                                    data=np.reshape(mask, (1, ) + mask.shape),
                                    shape=(1, ) + mask.shape,
                                    compressor=compressor,
                                    dtype=np.uint8,
                                    overwrite=True)
            min = np.min(np.transpose(np.nonzero(mask)), axis=0)
            max = np.max(np.transpose(np.nonzero(mask)), axis=0)
            shape = max - min + 1
            bb_offset = min
            if kwargs.get('ignore_small_comps',
                          0) > 0 or kwargs.get('skeletonize_foreground'):
                in_f[cleaned_mask_key].attrs['offset'] = [
                    int(off) for off in bb_offset
                ]
                in_f[cleaned_mask_key].attrs['fg_shape'] = [
                    int(dim) for dim in shape
                ]
    else:
        shape = input_shape
        bb_offset = [0] * len(shape)
    if len(shape) == 2:
        shape = (1, ) + tuple(shape)
    logger.info("input shape: {}, bb cropped shape: {}, offset: {}".format(
        input_shape, shape, bb_offset))

    # create offset lists
    offsets = get_offsets(shape, kwargs['chunksize'])
    # offsets = [offset + bb_offset for offset in offsets]
    logger.info("processing {} blocks".format(len(offsets)))
    logger.debug("blocks: {}".format(offsets))

    def init(l):
        global mutex
        mutex = l

    mutex = Lock()
    if kwargs['num_parallel_blocks'] > 1:
        pool = Pool(processes=kwargs['num_parallel_blocks'],
                    initializer=init,
                    initargs=(mutex, ))
        pool.map(
            functools.partial(blockwise_vote_instances, pred_file, pred_keys,
                              result_file, tmp_key, shape, channel_order,
                              bb_offset, kwargs), offsets)
        pool.close()
        pool.join()
    else:
        kwargs['mutex'] = mutex
        for idx, offset in enumerate(offsets):
            logger.info("start block idx: %s/%s (file %s)", idx, len(offsets),
                        sample)
            blockwise_vote_instances(pred_file, pred_keys, result_file,
                                     tmp_key, shape, channel_order, bb_offset,
                                     kwargs, offset)

    # stitch blocks
    #child_pid = os.fork()
    #if child_pid == 0:
    # child process
    instances = stitch_vote_instances(result_file, tmp_key, res_key,
                                      input_shape, bb_offset, shape, pred_file,
                                      pred_keys, channel_order, **kwargs)
    # save mip
    if kwargs.get('save_mip', False):
        colored = color(np.max(instances, axis=0))
        io.imsave(os.path.join(result_folder, sample + '.png'),
                  colored.astype(np.uint8))

    # remove small components
    remove_small_comps = kwargs.get('remove_small_comps', 0)
    if remove_small_comps > 0:

        instances = remove_small_components(instances, remove_small_comps)
        instances = relabel(instances)
        io.imsave(os.path.join(result_folder, sample + '.tif'),
                  instances.astype(np.uint16),
                  plugin='tifffile')
        if kwargs.get('save_mip', False):
            colored = color(np.max(instances, axis=0))
            io.imsave(os.path.join(result_folder, sample + '_cleaned.png'),
                      colored.astype(np.uint8))

    if kwargs['output_format'] == 'hdf':
        hf = h5py.File(os.path.join(result_folder, sample + '.hdf'), 'w')
        hf.create_dataset(res_key,
                          data=instances.astype(np.uint16),
                          dtype=np.uint16,
                          compression='gzip')
        if kwargs.get("dilate_instances", False):
            logger.info("dilating")
            instdil = np.copy(instances)
            for lbl in np.unique(instances):
                if lbl == 0:
                    continue
                label_mask = instdil == lbl
                dilated_label_mask = ndimage.binary_dilation(label_mask,
                                                             iterations=1)
                instdil[dilated_label_mask] = lbl
            hf.create_dataset(res_key + "_dil_1",
                              data=instdil.astype(np.uint16),
                              dtype=np.uint16,
                              compression='gzip')
예제 #4
0
def visualize_instances(label_fn,
                        label_key,
                        output_file,
                        max_axis=None,
                        show_outline=False,
                        raw_file=None,
                        raw_key=None):
    # read labeling
    if label_fn.endswith('.zarr'):
        inf = zarr.open(label_fn, mode='r')
    elif label_fn.endswith('.hdf'):
        inf = h5py.File(label_fn, 'r')
    else:
        raise NotImplementedError
    label = np.squeeze(np.array(inf[label_key]))
    print(label.shape, label.dtype)
    if label_fn.endswith('.hdf'):
        inf.close()
    if max_axis is not None:
        label = np.max(label, axis=max_axis)

    # read raw if given
    raw = None
    if raw_file is not None and raw_key is not None:
        if raw_file.endswith('.zarr'):
            inf = zarr.open(raw_file, mode='r')
        elif raw_file.endswith('.hdf'):
            inf = h5py.File(raw_file, 'r')
        else:
            raise NotImplementedError
        raw = np.squeeze(np.array(inf[raw_key]))
        if raw_file.endswith('.hdf'):
            raw_file.close()

    if show_outline:
        # label = label.astype(np.uint16)
        # for i in range(label.shape[0]):
        #     label[i] *= (i+1)

        labels, locations = np.unique(label, return_index=True)
        print(labels)
        locations = np.delete(locations, np.where(labels == 0))
        labels = np.delete(labels, np.where(labels == 0))
        # for different colormaps, see https://colorcet.pyviz.org/
        colormap = cc.glasbey_light
        # uncomment to choose randomly from colormap
        # colormap = np.random.choice(colormap, size=len(labels),
        #                             replace=(len(labels)>len(colormap)))

        if raw is None:
            colored = np.zeros(label.shape[1:] + (3, ), dtype=np.uint8)
        else:
            # convert raw to np.uint8
            # heads up: assuming normalized raw between [0, 1]
            raw = (raw * 255).astype(np.uint8)
            colored = np.stack([raw] * 3, axis=-1)
        for i, (lbl, loc) in enumerate(zip(labels, locations)):
            if lbl == 0:
                continue
            # heads up: assuming one instance per channel
            c = np.unravel_index(loc, label.shape)[0]
            outline = ndimage.distance_transform_cdt(label[c] == lbl) == 1
            #colored[outline, :] = np.random.randint(0, 255, 3)
            colored[outline, :] = hex_to_rgb(colormap[i])
    else:
        colored = color(label)
    io.imsave(output_file, colored.astype(np.uint8))
예제 #5
0
def main(pred_file, result_folder='.', **kwargs):
    assert os.path.exists(pred_file), \
        'Prediction file {} does not exists. Please check!'.format(pred_file)

    sample = os.path.basename(pred_file).split('.')[0]
    if sample == 'GMR_38F04_AE_01-20181005_63_G5':
        return
    if sample == 'BJD_127B01_AE_01-20171124_64_H5':
        return
    result_file = os.path.join(result_folder, sample + '.zarr')
    aff_key = kwargs['aff_key']
    fg_key = kwargs.get('fg_key')

    # read input shape
    if pred_file.endswith('.zarr'):
        in_f = zarr.open(pred_file, mode='r')
    else:
        raise NotImplementedError
    aff_shape = in_f[aff_key].shape
    channel_order = [slice(0, aff_shape[0])]
    pred_keys = [aff_key]

    if kwargs['overlapping_inst']:
        numinst_shape = in_f[fg_key].shape
        channel_order.append(slice(0, numinst_shape[0]))
        pred_keys += [fg_key]
        assert aff_shape[1:] == numinst_shape[1:], \
            'Please check: affinity and numinst shape do not match!'
    input_shape = aff_shape[1:]  # without first channel dimension

    # check if blocks should only be within bounding box
    if kwargs.get('only_bb'):
        mid = np.prod(kwargs['patchshape']) // 2
        mask = np.array(in_f[aff_key][mid])
        mask = mask > kwargs['patch_threshold']
        if np.sum(mask) == 0:
            logger.warning("bb empty")
            return
        if kwargs.get('ignore_small_comps', 0) > 0:
            labeled = ndimage.label(mask, np.ones([3] * len(input_shape)))[0]
            labels, counts = np.unique(labeled, return_counts=True)
            labels = labels[counts <= kwargs.get('ignore_small_comps')]
            labeled = replace(labeled, np.array(labels),
                              np.array([0] * len(labels)))
            print('num small comps: ', len(labels))
            # for label in labels:
            #     mask[labeled == label] = 0
            mask = labeled > 0
        min = np.min(np.transpose(np.nonzero(mask)), axis=0)
        max = np.max(np.transpose(np.nonzero(mask)), axis=0)
        # TODO for l1 data
        # min = np.array([31, 31, 0])
        # max = np.array([input_shape[0]-30, input_shape[1]-30, input_shape[2]])
        shape = max - min + 1
        bb_offset = min
    else:
        shape = input_shape
        bb_offset = [0] * len(shape)
    if len(shape) == 2:
        shape = (1, ) + tuple(shape)
        bb_offset = [0] * len(shape)
    logger.info("input shape: %s, bb cropped shape: %s, offset: %s",
                input_shape, shape, bb_offset)

    if pred_file.endswith('.hdf'):
        in_f.close()

    # create offset lists
    offsets = get_chessboard_offsets(shape, kwargs['chunksize'])
    #offsets = [offset + bb_offset for offset in offsets]
    logger.info('num blocks: %s', len(offsets))
    logger.info("%s", offsets)

    # create temporary zarr dataset for blocks (2^dim x shape)
    tmp_key = 'volumes/tmp_worker'
    skip_tmp_worker = False
    if not os.path.exists(result_file):
        compressor = Blosc(cname='zstd', clevel=3, shuffle=Blosc.BITSHUFFLE)
        f = zarr.open(result_file, mode='w')
        f.create_dataset(tmp_key,
                         shape=(2**len(shape), ) + tuple(shape),
                         compressor=compressor,
                         dtype=np.uint32,
                         chunks=(1, ) + tuple(kwargs['chunksize']))
    else:
        f = zarr.open(result_file, mode='r')
        if tmp_key in f:
            skip_tmp_worker = True

    def init(l):
        global mutex
        mutex = l

    if not skip_tmp_worker:
        mutex = Lock()
        if kwargs['num_parallel_blocks'] > 1:
            pool = Pool(processes=kwargs['num_parallel_blocks'],
                        initializer=init,
                        initargs=(mutex, ))
            pool.map(
                functools.partial(blockwise_vote_instances, pred_file,
                                  pred_keys, result_file, tmp_key, shape,
                                  channel_order, bb_offset, kwargs), offsets)
            pool.close()
            pool.join()
        else:
            kwargs['mutex'] = mutex
            for idx, offset in enumerate(offsets):
                # if idx < 7 or idx > 8:
                #     continue
                logger.info("start block idx: %s/%s (file %s)", idx,
                            len(offsets), sample)
                blockwise_vote_instances(pred_file, pred_keys, result_file,
                                         tmp_key, shape, channel_order,
                                         bb_offset, kwargs, offset)
    else:
        logger.info("skipping tmp_worker (blocks already exist?)")

    # stitch blocks
    res_key = kwargs.get('res_key', 'vote_instances')
    logger.info("%s", kwargs)
    instances = stitch_vote_instances(result_file, tmp_key, res_key,
                                      input_shape, bb_offset, **kwargs)

    # save mip
    save_mip = kwargs.get('save_mip', False)
    if save_mip:
        colored = color(np.max(instances, axis=0))
        io.imsave(os.path.join(result_folder, sample + '.png'),
                  colored.astype(np.uint8))

    # remove small components
    remove_small_comps = kwargs.get('remove_small_comps', 0)
    if remove_small_comps > 0:

        instances = remove_small_components(instances, remove_small_comps)
        instances = relabel(instances)
        io.imsave(os.path.join(result_folder, sample + '.tif'),
                  instances.astype(np.uint16),
                  plugin='tifffile')
        if save_mip:
            colored = color(np.max(instances, axis=0))
            io.imsave(os.path.join(result_folder, sample + '_cleaned.png'),
                      colored.astype(np.uint8))

    if kwargs['output_format'] == 'hdf':
        hf = h5py.File(os.path.join(result_folder, sample + '.hdf'), 'w')
        hf.create_dataset(res_key,
                          data=instances.astype(np.uint16),
                          dtype=np.uint16,
                          compression='gzip')
        if kwargs.get("dilate_instances", False):
            logger.info("dilating")
            instdil = np.copy(instances)
            for lbl in np.unique(instances):
                if lbl == 0:
                    continue
                label_mask = instdil == lbl
                dilated_label_mask = ndimage.binary_dilation(label_mask,
                                                             iterations=1)
                instdil[dilated_label_mask] = lbl
            hf.create_dataset(res_key + "_dil_1",
                              data=instdil.astype(np.uint16),
                              dtype=np.uint16,
                              compression='gzip')