create_consensus_from_json.py

import zarr
import argparse
import numpy as np
from empanada.inference.tracker import InstanceTracker
from empanada.inference import filters
from empanada.inference.patterns import fill_volume
from empanada.consensus import merge_objects_from_trackers

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('zarr_store', type=str, metavar='zarr_store', help='Path to zarr directory')
    parser.add_argument('zarr_key', type=str, metavar='zarr_key', help='Name to save the consensus segmentation in zarr.')
    parser.add_argument('tracker_json_files', type=str, nargs='+', metavar='tracker_json_files', 
                        help='At least 2 json tracker files to create consensus from')
    parser.add_argument('-v', type=int, dest='pixel_vote_thr', metavar='pixel_vote_thr', 
                        help='Number between 1 and number of trackers given. Min number of votes.')
    parser.add_argument('-c', type=float, dest='cluster_iou_thr', metavar='cluster_iou_thr', 
                        help='Minimum iou required between objects for them to form a cluster.')
    parser.add_argument('-s', type=int, dest='min_size', metavar='min_size', default=500,
                        help='Minimum object size, in voxels, in the final 3d segmentation')
    parser.add_argument('-e', type=int, dest='min_span', metavar='min_span', default=4,
                        help='Minimum number of consecutive slices that object must appear on in final 3d segmentation')
    parser.add_argument('--one-view', action='store_true',
                        help='Same as Permit detection found in 1 stack from napari plugin')
    args = parser.parse_args()

    # load the json files into trackers
    trackers = []
    for fp in args.tracker_json_files:
        tracker = InstanceTracker()
        tracker.load_from_json(fp)
        trackers.append(tracker)

    # some helpful parameters
    class_id = trackers[0].class_id
    label_divisor = trackers[0].label_divisor
    shape = trackers[0].shape3d

    # create the consensus from run length encodings
    consensus = merge_objects_from_trackers(
        trackers, 
        args.pixel_vote_thr, 
        args.cluster_iou_thr, 
        args.one_view
    )
    consensus_tracker = InstanceTracker(class_id, label_divisor, shape, 'xy')
    consensus_tracker.instances = consensus

    # filter out small objects and pancakes
    filters.remove_small_objects(consensus_tracker, min_size=args.min_size)
    filters.remove_pancakes(consensus_tracker, min_span=args.min_span)
    

    data = zarr.open(args.zarr_store, mode='r+')
    consensus_vol = data.create_dataset(
        args.zarr_key, shape=shape, dtype=np.uint32,
        overwrite=True, chunks=(256, 256, 256)
    )
    fill_volume(consensus_vol, consensus_tracker.instances, processes=4)

    print('Finished')