def create_skeletons( ctx, seg_layer_spec, dst_folder, timestamp, mip, teasar_scale, teasar_const, ids, ids_filepath, tick_threshold, chunk_xy, chunk_z, single_merge_mode, start_coord, end_coord, coord_mip, ): scheduler = ctx.obj["scheduler"] corgie_logger.debug("Setting up layers...") seg_stack = create_stack_from_spec(seg_layer_spec, name="src", readonly=True) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) object_ids = ids if ids_filepath is not None: object_ids = [] with open(ids_filepath, "r") as f: line = f.readline() while line: object_ids.append(int(line)) line = f.readline() if object_ids is None or len(object_ids) == 0: raise ValueError("Must specify ids to skeletonize") object_ids = list(object_ids) teasar_params = {"scale": teasar_scale, "const": teasar_const} seg_layer = seg_stack.get_layers_of_type("segmentation")[0] skeleton_job = SkeletonJob( seg_layer=seg_layer, dst_path=dst_folder, timestamp=timestamp, bcube=bcube, chunk_xy=chunk_xy, chunk_z=chunk_z, mip=mip, teasar_params=teasar_params, object_ids=object_ids, tick_threshold=tick_threshold, single_merge_mode=single_merge_mode, ) scheduler.register_job(skeleton_job, job_name="Skeletonize {}".format(bcube)) scheduler.execute_until_completion() result_report = f"Skeletonized {str(seg_layer)}. " corgie_logger.info(result_report)
def execute(self): corgie_logger.info(f"Merging skeletons at {self.dst_path}") fragment_filenames = self.cf.list(prefix=self.prefix, flat=True) skeleton_files = self.cf.get(fragment_filenames) skeletons = defaultdict(list) for skeleton_file in skeleton_files: try: colon_index = skeleton_file["path"].index(":") except ValueError: # File is full skeleton, not fragment continue seg_id = skeleton_file["path"][0:colon_index] skeleton_fragment = pickle.loads(skeleton_file["content"]) if not skeleton_fragment.empty(): skeletons[seg_id].append(skeleton_fragment) for seg_id, skeleton_fragments in skeletons.items(): skeleton = PrecomputedSkeleton.simple_merge( skeleton_fragments).consolidate() skeleton = kimimaro.postprocess(skeleton, self.dust_threshold, self.tick_threshold) skeleton.id = int(seg_id) self.cf.put(path=seg_id, content=skeleton.to_precomputed(), compress="gzip") corgie_logger.info(f"Finished skeleton {seg_id}")
def task_generator(self): chunks = self.dst_layer.break_bcube_into_chunks( bcube=self.bcube, chunk_xy=self.chunk_xy, chunk_z=1, mip=self.mip) if 'src_img' in self.src_specs[0]: tasks = [MergeRenderImageTask( src_layers=self.src_layers, src_specs=self.src_specs, dst_layer=self.dst_layer, mip=self.mip, pad=self.pad, bcube=input_chunk) for input_chunk in chunks] else: tasks = [MergeRenderMaskTask( src_layers=self.src_layers, src_specs=self.src_specs, dst_layer=self.dst_layer, mip=self.mip, pad=self.pad, bcube=input_chunk) for input_chunk in chunks] corgie_logger.info( f"Yielding render tasks for bcube: {self.bcube}, MIP: {self.mip}") yield tasks
def transform_skeletons( ctx, vector_field_spec, src_folder, dst_folder, field_mip, ids, ids_filepath, task_vertex_size, calculate_skeleton_lengths, mip0_field, ): scheduler = ctx.obj["scheduler"] corgie_logger.debug("Setting up layers...") vf_stack = create_stack_from_spec(vector_field_spec, name="src", readonly=True) skeleton_ids = ids if ids_filepath is not None: skeleton_ids = [] with open(ids_filepath, "r") as f: line = f.readline() while line: skeleton_ids.append(int(line)) line = f.readline() if len(skeleton_ids) == 0: skeleton_ids = None else: skeleton_ids = list(skeleton_ids) skeleton_length_file = None if calculate_skeleton_lengths: import time if not os.path.exists("skeleton_lengths"): os.makedirs("skeleton_lengths") skeleton_length_file = f"skeleton_lengths/skeleton_lengths_{int(time.time())}" vf_layer = vf_stack.get_layers_of_type("field")[0] transform_skeletons_job = TransformSkeletonsJob( vector_field_layer=vf_layer, src_path=src_folder, dst_path=dst_folder, field_mip=field_mip, skeleton_ids=skeleton_ids, task_vertex_size=task_vertex_size, skeleton_length_file=skeleton_length_file, mip0_field=mip0_field, ) scheduler.register_job( transform_skeletons_job, job_name="Transforming skeletons in {}".format(src_folder), ) scheduler.execute_until_completion() result_report = f"Transformed skeletons stored at {dst_folder}. " corgie_logger.info(result_report)
def execute(self): corgie_logger.info( f"Skeletonizing {self.seg_layer} at MIP{self.mip}, region: {self.bcube}" ) seg_data = self.seg_layer.read(bcube=self.bcube, mip=self.mip, timestamp=self.timestamp) resolution = self.seg_layer.cv[self.mip].resolution skeletons = kimimaro.skeletonize( seg_data, self.teasar_params, object_ids=self.object_ids, anisotropy=resolution, dust_threshold=self.dust_threshold, progress=False, fix_branching=self.fix_branching, fix_borders=self.fix_borders, fix_avocados=self.fix_avocados, ).values() minpt = self.bcube.minpt(self.mip) for skel in skeletons: skel.vertices[:] += minpt * resolution cf = CloudFiles(self.dst_path) for skel in skeletons: path = "{}:{}".format(skel.id, self.bcube.to_filename(self.mip)) cf.put( path=path, content=pickle.dumps(skel), compress="gzip", content_type="application/python-pickle", cache_control=False, )
def task_generator(self): for mip in self.mips: chunks = self.dst_stack.get_layers()[0].break_bcube_into_chunks( bcube=self.bcube, chunk_xy=self.chunk_xy, chunk_z=self.chunk_z, mip=mip, return_generator=True, ) tasks = (RenderTask( self.src_stack, self.dst_stack, blackout_masks=self.blackout_masks, render_masks=self.render_masks, mip=mip, pad=self.pad, bcube=input_chunk, additional_fields=self.additional_fields, preserve_zeros=self.preserve_zeros, seethrough_mask_layer=self.seethrough_mask_layer, seethrough_offset=self.seethrough_offset, ) for input_chunk in chunks) corgie_logger.info( f"Yielding render tasks for bcube: {self.bcube}, MIP: {mip}") yield tasks
def downsample(ctx, src_layer_spec, dst_layer_spec, mip_start, mip_end, chunk_xy, chunk_z, mips_per_task, start_coord, end_coord, coord_mip): scheduler = ctx.obj['scheduler'] corgie_logger.debug("Setting up Source and Destination layers...") src_layer = create_layer_from_spec(src_layer_spec, caller_name='src layer', readonly=True) if dst_layer_spec is None: corgie_logger.info("Destination layer not specified. Using Source layer " "as Destination.") dst_layer = src_layer dst_layer.readonly = False else: dst_layer = create_layer_from_spec(dst_layer_spec, caller_name='dst_layer layer', readonly=False, reference=src_layer, chunk_z=chunk_z, overwrite=True) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) downsample_job = DownsampleJob(src_layer=src_layer, dst_layer=dst_layer, mip_start=mip_start, mip_end=mip_end, bcube=bcube, chunk_xy=chunk_xy, chunk_z=chunk_z, mips_per_task=mips_per_task) # create scheduler and execute the job scheduler.register_job(downsample_job, job_name="downsample") scheduler.execute_until_completion() result_report = f"Downsampled {src_layer} from {mip_start} to {mip_end}. Result in {dst_layer}" corgie_logger.info(result_report)
def fill_nearest( ctx, src_layer_spec, dst_folder, chunk_xy, start_coord, end_coord, coord_mip, suffix, mip, radius, force_chunk_z=1, ): scheduler = ctx.obj["scheduler"] if suffix is None: suffix = "_seethrough" else: suffix = f"_{suffix}" crop, pad = 0, 0 corgie_logger.debug("Setting up layers...") src_stack = create_stack_from_spec(src_layer_spec, name="src", readonly=True) src_stack.folder = dst_folder dst_stack = stack.create_stack_from_reference( reference_stack=src_stack, folder=dst_folder, name="dst", types=["img", "mask"], readonly=False, suffix=suffix, overwrite=True, force_chunk_z=force_chunk_z, ) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) fill_nearest_job = FillNearestJob( src_stack=src_stack, dst_stack=dst_stack, bcube=bcube, radius=radius, mip=mip, chunk_xy=chunk_xy, ) # create scheduler and execute the job scheduler.register_job(fill_nearest_job, job_name="Fill Nearest Block {}".format(bcube)) scheduler.execute_until_completion() result_report = ( f"Rendered layers {[str(l) for l in src_stack.get_layers_of_type('img')]}. " f"Results in {[str(l) for l in dst_stack.get_layers_of_type('img')]}") corgie_logger.info(result_report)
def combine_masks( ctx, src_layer_spec, dst_layer_spec, exp, chunk_xy, chunk_z, force_chunk_xy, force_chunk_z, start_coord, end_coord, coord_mip, mip, pad, ): scheduler = ctx.obj["scheduler"] if not force_chunk_xy: force_chunk_xy = chunk_xy corgie_logger.debug("Setting up layers...") src_stack = create_stack_from_spec(src_layer_spec, name="src", readonly=True) reference_layer = src_stack.reference_layer dst_layer = create_layer_from_spec( dst_layer_spec, allowed_types=["mask"], default_type="mask", readonly=False, caller_name="dst_layer", reference=reference_layer, force_chunk_xy=force_chunk_xy, force_chunk_z=force_chunk_z, overwrite=True, ) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) combine_masks_job = CombineMasksJob( src_stack=src_stack, exp=json.loads(exp), dst_layer=dst_layer, mip=mip, bcube=bcube, pad=pad, chunk_xy=chunk_xy, chunk_z=chunk_z, ) # create scheduler and execute the job scheduler.register_job(combine_masks_job, job_name="Combine Masks {}".format(bcube)) scheduler.execute_until_completion() result_report = f"Results in {str(dst_layer)}" corgie_logger.info(result_report)
def spec_to_layer_dict_readonly(layer_specs): """Create dict of layers from a corgie spec indexed by unique id These layers will be readonly Args: layer_specs (dict): layer specs indexed by unique id spec_key (str): src/tgt """ layers = {} for k, s in layer_specs.items(): corgie_logger.info(f'Creating layer no. {k}') layers[k] = create_layer_from_dict(s, readonly=True) return layers
def execute(self): corgie_logger.info( f"Generate new skeleton vertices task for id {self.skeleton_id_str}" ) skeleton = get_skeleton(self.src_path, self.skeleton_id_str) if self.vertex_sort: vertex_sort = skeleton.vertices[:, 2].argsort() else: vertex_sort = np.arange(0, len(skeleton.vertices)) number_vertices = len(skeleton.vertices) index_points = list(range(0, number_vertices, self.task_vertex_size)) cf = CloudFiles(f"{self.dst_path}") array_filenames = [] for i in range(len(index_points)): start_index = index_points[i] if i + 1 == len(index_points): end_index = number_vertices else: end_index = index_points[i + 1] array_filenames.append( f"intermediary_arrays/{self.skeleton_id_str}:{start_index}-{end_index}" ) array_files = cf.get(array_filenames) # Dict to make sure arrays are concatenated in correct order array_dict = {} for array_file in array_files: array_dict[array_file["path"]] = pickle.loads( array_file["content"]) array_arrays = [] for array_filename in array_filenames: array_arrays.append(array_dict[array_filename]) array_arrays = np.concatenate(array_arrays) # Restore the correct order of the vertices restore_sort = vertex_sort.argsort() new_vertices = array_arrays[restore_sort] new_skeleton = Skeleton( vertices=new_vertices, edges=skeleton.edges, radii=skeleton.radius, vertex_types=skeleton.vertex_types, space=skeleton.space, transform=skeleton.transform, ) cf.put( path=self.skeleton_id_str, content=new_skeleton.to_precomputed(), compress="gzip", )
def task_generator(self): chunks = self.dst_stack.get_layers()[0].break_bcube_into_chunks( bcube=self.bcube, chunk_xy=self.chunk_xy, chunk_z=1, mip=self.mip) tasks = [MergeCopyTask(src_stack=self.src_stack, dst_stack=self.dst_stack, mip=self.mip, bcube=input_chunk, z_list=self.z_list) for input_chunk in chunks] corgie_logger.info(f"Yielding copy tasks for bcube: {self.bcube}, MIP: {self.mip}") yield tasks
def execute(self): corgie_logger.info( f"Normalizing {self.src_layer} at MIP{self.mip}, region: {self.bcube}" ) mean_data = self.mean_layer.read(self.bcube, mip=self.stats_mip) var_data = self.var_layer.read(self.bcube, mip=self.stats_mip) src_data = self.src_layer.read(self.bcube, mip=self.mip) mask_data = helpers.read_mask_list(mask_list=self.mask_layers, bcube=self.bcube, mip=self.mip) dst_data = (src_data - mean_data) / var_data.sqrt() if mask_data is not None: dst_data[mask_data] = self.mask_value self.dst_layer.write(dst_data, self.bcube, mip=self.mip)
def filter_skeletons( ctx, src_folder, dst_folder, ids, bad_sections, ids_filepath, z_start, z_end ): scheduler = ctx.obj["scheduler"] corgie_logger.debug("Setting up layers...") skeleton_ids = ids if ids_filepath is not None: skeleton_ids = [] with open(ids_filepath, "r") as f: line = f.readline() while line: skeleton_ids.append(int(line)) line = f.readline() if len(skeleton_ids) == 0: skeleton_ids = None else: skeleton_ids = list(skeleton_ids) transform_skeletons_job = FilterSkeletonsJob( src_path=src_folder, dst_path=dst_folder, skeleton_ids=skeleton_ids, bad_sections=bad_sections, z_start=z_start, z_end=z_end ) scheduler.register_job( transform_skeletons_job, job_name="Filtering skeletons in {}".format(src_folder), ) scheduler.execute_until_completion() result_report = f"Filtered skeletons stored at {dst_folder}. " corgie_logger.info(result_report)
def downsample_by_spec(ctx, src_layer_spec, spec_path, dst_layer_spec, mip_start, mip_end, chunk_xy, chunk_z, mips_per_task, start_coord, end_coord, coord_mip): scheduler = ctx.obj['scheduler'] corgie_logger.debug("Setting up Source and Destination layers...") src_layer = create_layer_from_spec(src_layer_spec, caller_name='src layer', readonly=True) with open(spec_path, 'r') as f: spec = set(json.load(f)) if dst_layer_spec is None: corgie_logger.info( "Destination layer not specified. Using Source layer " "as Destination.") dst_layer = src_layer dst_layer.readonly = False else: dst_layer = create_layer_from_spec(dst_layer_spec, caller_name='dst_layer layer', readonly=False, reference=src_layer, chunk_z=chunk_z, overwrite=True) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) for z in range(*bcube.z_range()): if z in spec: job_bcube = bcube.reset_coords(zs=z, ze=z + 1, in_place=False) downsample_job = DownsampleJob(src_layer=src_layer, dst_layer=dst_layer, mip_start=mip_start, mip_end=mip_end, bcube=job_bcube, chunk_xy=chunk_xy, chunk_z=chunk_z, mips_per_task=mips_per_task) # create scheduler and execute the job scheduler.register_job(downsample_job, job_name=f"Downsample {job_bcube}") scheduler.execute_until_completion() result_report = f"Downsampled {src_layer} from {mip_start} to {mip_end}. Result in {dst_layer}" corgie_logger.info(result_report)
def task_generator(self): chunks = self.seg_layer.break_bcube_into_chunks( bcube=self.bcube, chunk_xy=self.chunk_xy, chunk_z=self.chunk_z, mip=self.mip, readonly=True, ) tasks = [ SkeletonTask( self.seg_layer, self.dst_path, timestamp=self.timestamp, mip=self.mip, teasar_params=self.teasar_params, object_ids=self.object_ids, dust_threshold=self.dust_threshold, fix_branching=self.fix_branching, fix_borders=self.fix_borders, fix_avocados=self.fix_avocados, bcube=input_chunk, ) for input_chunk in chunks ] corgie_logger.info( f"Yielding skeletonization tasks for bcube: {self.bcube}, MIP: {self.mip}" ) yield tasks yield scheduling.wait_until_done if self.single_merge_mode: merge_tasks = [ MergeSkeletonTask( self.dst_path, self.mip, self.dust_threshold, self.tick_threshold, str(object_id), ) for object_id in self.object_ids ] yield merge_tasks else: yield [ MergeSkeletonTask(self.dst_path, self.mip, self.dust_threshold, self.tick_threshold) ]
def task_generator(self): chunks = self.dst_stack.get_layers()[0].break_bcube_into_chunks( bcube=self.bcube, chunk_xy=self.chunk_xy, chunk_z=self.chunk_z, mip=self.mip) tasks = [ CopyTask(self.src_stack, self.dst_stack, blackout_masks=self.blackout_masks, copy_masks=self.copy_masks, mip=self.mip, bcube=input_chunk) for input_chunk in chunks ] corgie_logger.info( f"Yielding copy tasks for bcube: {self.bcube}, MIP: {self.mip}") yield tasks
def task_generator(self): chunks = self.dst_layer.break_bcube_into_chunks( bcube=self.bcube, chunk_xy=self.chunk_xy, chunk_z=self.chunk_z, mip=self.mip, return_generator=True, ) tasks = (CopyLayerTask( self.src_layer, self.dst_layer, mip=self.mip, bcube=input_chunk, ) for input_chunk in chunks) corgie_logger.info( f"Yielding copy layer tasks for bcube: {self.bcube}, MIP: {self.mip}" ) yield tasks
def task_generator(self): chunks = self.dst_stack.get_layers()[0].break_bcube_into_chunks( bcube=self.bcube, chunk_xy=self.chunk_xy, chunk_z=1, mip=self.mip, return_generator=True, ) tasks = (FillNearestTask( self.src_stack, self.dst_stack, mip=self.mip, bcube=chunk, radius=self.radius, ) for chunk in chunks) corgie_logger.info( f"Yielding fill nearest tasks for bcube: {self.bcube}, MIP: {self.mip}" ) yield tasks
def task_generator(self): chunks = self.dst_stack.get_layers()[0].break_bcube_into_chunks( bcube=self.bcube, chunk_xy=self.chunk_xy, chunk_z=self.chunk_z, mip=self.mip) tasks = [ RenderTask(self.src_stack, self.dst_stack, blackout_masks=self.blackout_masks, render_masks=self.render_masks, mip=self.mip, pad=self.pad, bcube=input_chunk, additional_fields=self.additional_fields) for input_chunk in chunks ] corgie_logger.info( f"Yielding render tasks for bcube: {self.bcube}, MIP: {self.mip}") yield tasks
def execute(self): device = "cuda" if torch.cuda.is_available() else "cpu" padded_bcube = self.bcube.uncrop(self.pad, self.mip) for k, specs in enumerate(self.src_specs[::-1]): src_z = specs["src_z"] dst_z = self.bcube.z_range()[0] corgie_logger.info(f"Load fields for {padded_bcube}") # backwards compatible if not isinstance(specs["src_field"], list): specs["src_field"] = [specs["src_field"]] mask_layer = self.src_layers[str(specs["src_mask"])] field_ids = list(map(str, specs["src_field"])) corgie_logger.info(f"field ids={field_ids}") z_list = specs.get("src_field_z", [src_z] * len(field_ids)) fields = FieldSet([self.src_layers[n] for n in field_ids]) field = fields.read( bcube=padded_bcube, z_list=z_list, mip=self.mip, device=device ) bcube = padded_bcube.reset_coords(zs=src_z, ze=src_z + 1, in_place=False) mask_trans = helpers.percentile_trans_adjuster(field) mask_trans = mask_trans.round_to_mip(self.mip, mask_layer.data_mip) corgie_logger.debug(f"mask_trans: {mask_trans}") mask_bcube = bcube.translate( x_offset=mask_trans.y, y_offset=mask_trans.x, mip=self.mip ) corgie_logger.info(f"Load masks for {mask_bcube}") mask_id = specs["mask_id"] mask_layer.binarizer = helpers.Binarizer(["eq", mask_id]) mask = mask_layer.read(bcube=mask_bcube, mip=self.mip, device=device) mask = residuals.res_warp_img( mask.float(), field - mask_trans.to_tensor(device=field.device) ).tensor() mask = (mask > 0.4).bool() cropped_mask = helpers.crop(mask, self.pad) relabel_id = torch.as_tensor(specs.get("relabel_id", k + 1), dtype=torch.uint8) if k == 0: dst_img = cropped_mask * relabel_id dst_img[~cropped_mask] = 0 else: dst_img[cropped_mask] = cropped_mask[cropped_mask] * relabel_id self.dst_layer.write(dst_img.cpu(), bcube=self.bcube, mip=self.mip)
def copy( ctx, src_layer_spec, dst_folder, copy_masks, blackout_masks, chunk_xy, chunk_z, start_coord, end_coord, coord_mip, mip, suffix, force_chunk_xy, force_chunk_z, ): scheduler = ctx.obj["scheduler"] if suffix is None: suffix = "" else: suffix = f"_{suffix}" corgie_logger.debug("Setting up layers...") src_stack = create_stack_from_spec(src_layer_spec, name="src", readonly=True) if not force_chunk_xy: force_chunk_xy = chunk_xy if not force_chunk_z: force_chunk_z = chunk_z dst_stack = stack.create_stack_from_reference( reference_stack=src_stack, folder=dst_folder, name="dst", types=["img", "mask"], readonly=False, suffix=suffix, force_chunk_xy=force_chunk_xy, force_chunk_z=force_chunk_z, overwrite=True, ) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) copy_job = CopyJob( src_stack=src_stack, dst_stack=dst_stack, mip=mip, bcube=bcube, chunk_xy=chunk_xy, chunk_z=chunk_z, copy_masks=copy_masks, blackout_masks=blackout_masks, ) # create scheduler and execute the job scheduler.register_job(copy_job, job_name="Copy {}".format(bcube)) scheduler.execute_until_completion() result_report = ( f"Copied layers {[str(l) for l in src_stack.get_layers_of_type('img')]}. " f"to {[str(l) for l in dst_stack.get_layers_of_type('img')]}") corgie_logger.info(result_report)
def execute(self): corgie_logger.info( f"Starting transform skeleton vertices task for id {self.skeleton_id_str}" ) skeleton = get_skeleton(self.src_path, self.skeleton_id_str) if self.vertex_sort: vertex_sort = skeleton.vertices[:, 2].argsort() else: vertex_sort = np.arange(0, len(skeleton.vertices)) # How many vertices we will use at once to get a bcube to download from the vector field vertex_process_size = 50 vertices_to_transform = skeleton.vertices[ vertex_sort[self.start_vertex_index:self.end_vertex_index]] index_vertices = list( range(0, self.number_vertices, vertex_process_size)) new_vertices = [] for i in range(len(index_vertices)): if i + 1 == len(index_vertices): current_batch_vertices = vertices_to_transform[ index_vertices[i]:] else: current_batch_vertices = vertices_to_transform[ index_vertices[i]:index_vertices[i + 1]] field_resolution = np.array( self.vector_field_layer.resolution(self.field_mip)) bcube = get_bcube_from_vertices( vertices=current_batch_vertices, resolution=field_resolution, mip=self.field_mip, ) field_data = self.vector_field_layer.read( bcube=bcube, mip=self.field_mip).permute(2, 3, 0, 1) current_batch_vertices_to_mip = current_batch_vertices / field_resolution bcube_minpt = bcube.minpt(self.field_mip) field_indices = current_batch_vertices_to_mip.astype( np.int) - bcube_minpt vector_resolution = ( self.vector_field_layer.resolution(0) * np.array([ 2**(self.field_mip - self.vector_field_layer.data_mip), 2**(self.field_mip - self.vector_field_layer.data_mip), 1, ]) if self.mip0_field else self.vector_field_layer.resolution( self.field_mip)) vectors_to_add = [] corgie_logger.info(f"{field_data.shape}, {field_indices.max(0)}") for i in range(len(field_data.shape) - 1): if field_indices.max(0)[i] >= field_data.shape[i]: import pdb pdb.set_trace() for cur_field_index in field_indices: vector_at_point = field_data[cur_field_index[0], cur_field_index[1], cur_field_index[2]] # Each vector is stored in [Y,X] format vectors_to_add.append([ int(vector_resolution[0] * vector_at_point[1].item()), int(vector_resolution[1] * vector_at_point[0].item()), 0, ]) vectors_to_add = np.array(vectors_to_add) current_batch_warped_vertices = current_batch_vertices + vectors_to_add new_vertices.append(current_batch_warped_vertices) new_vertices = np.concatenate(new_vertices) cf = CloudFiles(f"{self.dst_path}/intermediary_arrays/") cf.put( path= f"{self.skeleton_id_str}:{self.start_vertex_index}-{self.end_vertex_index}", content=pickle.dumps(new_vertices), )
def normalize_by_spec(ctx, src_layer_spec, spec_path, dst_folder, stats_mip, mip_start, mip_end, chunk_xy, chunk_z, start_coord, end_coord, coord_mip, suffix, recompute_stats, mask_value): if chunk_z != 1: raise NotImplemented("Compute Statistics command currently only \ supports per-section statistics.") result_report = "" scheduler = ctx.obj['scheduler'] if suffix is None: suffix = '_norm' else: suffix = f"_{suffix}" if stats_mip is None: stats_mip = mip_end with open(spec_path, 'r') as f: spec = set(json.load(f)) corgie_logger.debug("Setting up layers...") src_stack = create_stack_from_spec(src_layer_spec, name='src', readonly=True) # dst_stack = stack.create_stack_from_reference(reference_stack=src_stack, # folder=dst_folder, name="dst", types=["img"], readonly=False, # suffix=suffix, overwrite=True) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) img_layers = src_stack.get_layers_of_type("img") mask_layers = src_stack.get_layers_of_type("mask") field_layers = src_stack.get_layers_of_type("field") assert len(field_layers) == 0 for l in img_layers: mean_layer = l.get_sublayer( name=f"mean{suffix}", path=os.path.join(dst_folder, f"mean{suffix}"), layer_type="section_value", ) var_layer = l.get_sublayer(name=f"var{suffix}", path=os.path.join(dst_folder, f"var{suffix}"), layer_type="section_value") if recompute_stats: for z in range(*bcube.z_range()): if z in spec: job_bcube = bcube.reset_coords(zs=z, ze=z + 1, in_place=False) compute_stats_job = ComputeStatsJob( src_layer=l, mask_layers=mask_layers, mean_layer=mean_layer, var_layer=var_layer, bcube=job_bcube, mip=stats_mip, chunk_xy=chunk_xy, chunk_z=chunk_z) # create scheduler and execute the job scheduler.register_job( compute_stats_job, job_name= f"Compute Stats. Layer: {l}, Bcube: {job_bcube}") scheduler.execute_until_completion() dst_layer = l.get_sublayer(name=f"{l.name}{suffix}", path=os.path.join(dst_folder, "img", f"{l.name}{suffix}"), layer_type=l.get_layer_type(), dtype='float32', overwrite=True) for z in range(*bcube.z_range()): if z in spec: job_bcube = bcube.reset_coords(zs=z, ze=z + 1, in_place=False) result_report += f"Normalized {l} -> {dst_layer}\n" for mip in range(mip_start, mip_end + 1): normalize_job = NormalizeJob(src_layer=l, mask_layers=mask_layers, dst_layer=deepcopy(dst_layer), mean_layer=mean_layer, var_layer=var_layer, stats_mip=stats_mip, mip=mip, bcube=job_bcube, chunk_xy=chunk_xy, chunk_z=chunk_z, mask_value=mask_value) # create scheduler and execute the job scheduler.register_job( normalize_job, job_name=f"Normalize {job_bcube}, MIP {mip}") scheduler.execute_until_completion() corgie_logger.info(result_report)
def task_generator(self): skeletons = self.get_skeletons(self.src_path) transform_vertex_tasks = [] generate_new_skeleton_tasks = [] for skeleton_id_str, skeleton in skeletons.items(): number_vertices = len(skeleton.vertices) # Vector field chunks are typically chunked by 1 in z, so we process # the skeleton's vertices in z-order for maximum download efficiency. index_points = list( range(0, number_vertices, self.task_vertex_size)) for i in range(len(index_points)): start_vertex_index = index_points[i] if i + 1 == len(index_points): end_vertex_index = number_vertices else: end_vertex_index = index_points[i + 1] transform_vertex_tasks.append( TransformSkeletonVerticesTask( vector_field_layer=self.vector_field_layer, skeleton_id_str=skeleton_id_str, src_path=self.src_path, dst_path=self.dst_path, field_mip=self.field_mip, start_vertex_index=start_vertex_index, end_vertex_index=end_vertex_index, vertex_sort=True, mip0_field=self.mip0_field, )) generate_new_skeleton_tasks.append( GenerateNewSkeletonTask( skeleton_id_str=skeleton_id_str, src_path=self.src_path, dst_path=self.dst_path, task_vertex_size=self.task_vertex_size, vertex_sort=True, )) corgie_logger.info( f"Yielding transform skeleton vertex tasks for skeletons in {self.src_path}" ) yield transform_vertex_tasks yield scheduling.wait_until_done corgie_logger.info(f"Generating skeletons to {self.dst_path}") yield generate_new_skeleton_tasks yield scheduling.wait_until_done # TODO: Delete intermediary vertex files if self.skeleton_length_file is not None: new_skeletons = self.get_skeletons(self.dst_path) corgie_logger.info( f"Calculating skeleton lengths to {self.skeleton_length_file}") with open(self.skeleton_length_file, "w") as f: f.write( "Skeleton id, Original Skeleton Length (nm), New Skeleton Length (nm)\n" ) for skeleton_id_str in skeletons: original_skeleton = skeletons[skeleton_id_str] new_skeleton = new_skeletons[skeleton_id_str] verts = new_skeleton.vertices vl = [verts[i] for i in range(verts.shape[0])] f.write( f"{skeleton_id_str},{int(original_skeleton.cable_length())},{int(new_skeleton.cable_length())}\n" )
def align( ctx, src_layer_spec, dst_folder, render_pad, render_chunk_xy, processor_spec, pad, crop, processor_mip, chunk_xy, start_coord, end_coord, coord_mip, bad_starter_path, block_size, stitch_size, vote_dist, consensus_threshold, blur_sigma, kernel_size, blend_xy, force_chunk_xy, suffix, seethrough_spec, seethrough_limit, seethrough_spec_mip, decay_dist, blur_rate, restart_stage, restart_suffix, ): scheduler = ctx.obj["scheduler"] if suffix is None: suffix = "_aligned" else: suffix = f"_{suffix}" if (restart_suffix is None) or (restart_stage == 0): restart_suffix = suffix if crop is None: crop = pad corgie_logger.debug("Setting up layers...") # TODO: store stitching images in layer other than even & odd if vote_dist + stitch_size - 2 >= block_size: raise exceptions.CorgieException( "block_size too small for stitching + voting requirements (stitch_size + vote_dist)" ) corgie_logger.debug("Setting up layers...") src_stack = create_stack_from_spec(src_layer_spec, name="src", readonly=True) src_stack.folder = dst_folder if force_chunk_xy is None: force_chunk_xy = chunk_xy dst_stack = stack.create_stack_from_reference( reference_stack=src_stack, folder=dst_folder, name="dst", types=["img", "mask"], readonly=False, suffix=restart_suffix, force_chunk_xy=force_chunk_xy, overwrite=True, ) even_stack = stack.create_stack_from_reference( reference_stack=src_stack, folder=os.path.join(dst_folder, "even"), name="even", types=["img", "mask"], readonly=False, suffix=suffix, force_chunk_xy=force_chunk_xy, overwrite=True, ) odd_stack = stack.create_stack_from_reference( reference_stack=src_stack, folder=os.path.join(dst_folder, "odd"), name="odd", types=["img", "mask"], readonly=False, suffix=suffix, force_chunk_xy=force_chunk_xy, overwrite=True, ) corgie_logger.debug("Done!") bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) corgie_logger.debug("Calculating blocks...") skip_list = [] if bad_starter_path is not None: with open(bad_starter_path) as f: line = f.readline() while line: skip_list.append(int(line)) line = f.readline() blocks = get_blocks( start=bcube.z_range()[0], stop=bcube.z_range()[1], block_size=block_size, block_overlap=1, skip_list=skip_list, src_stack=src_stack, even_stack=even_stack, odd_stack=odd_stack, ) stitch_blocks = [b.overlap(stitch_size) for b in blocks[1:]] corgie_logger.debug("All Blocks") for block, stitch_block in zip(blocks, [None] + stitch_blocks): corgie_logger.debug(block) corgie_logger.debug(f"Stitch {stitch_block}") corgie_logger.debug("\n") max_blur_mip = (math.ceil(math.log(decay_dist * blur_rate + 1, 2)) + processor_mip[-1]) corgie_logger.debug(f"Max blur mip for stitching field: {max_blur_mip}") # Set all field names, adjusting for restart suffix block_field_name = f"field{suffix}" stitch_estimated_suffix = f"_stitch_estimated{suffix}" stitch_estimated_name = f"field{stitch_estimated_suffix}" stitch_corrected_name = f"stitch_corrected{suffix}" stitch_corrected_field = None composed_name = f"composed{suffix}" if restart_stage <= 2: stitch_estimated_suffix = f"_stitch_estimated{restart_suffix}" stitch_estimated_name = f"field{stitch_estimated_suffix}" stitch_corrected_name = f"stitch_corrected{restart_suffix}" if restart_stage <= 3: composed_name = f"composed{restart_suffix}" render_method = helpers.PartialSpecification( f=RenderJob, pad=render_pad, chunk_xy=render_chunk_xy, chunk_z=1, render_masks=False, ) cf_method = helpers.PartialSpecification( f=ComputeFieldJob, pad=pad, crop=crop, processor_mip=processor_mip, processor_spec=processor_spec, chunk_xy=chunk_xy, blend_xy=blend_xy, chunk_z=1, ) if seethrough_spec != tuple(): assert seethrough_spec_mip is not None seethrough_method = helpers.PartialSpecification( f=SeethroughCompareJob, mip=seethrough_spec_mip, processor_spec=seethrough_spec, chunk_xy=chunk_xy, pad=pad, crop=pad, seethrough_limit=seethrough_limit, ) else: seethrough_method = None #restart_stage = 4 #import pdb; pdb.set_trace() if restart_stage == 0: corgie_logger.debug("Aligning blocks...") for block in blocks: block_bcube = block.get_bcube(bcube) # Use copies of src & dst so that aligning the stitching blocks # is not affected by these block fields. # Copying also allows local compute to not modify objects for other tasks align_block_job_forv = AlignBlockJob( src_stack=deepcopy(block.src_stack), dst_stack=deepcopy(block.dst_stack), bcube=block_bcube, render_method=render_method, cf_method=cf_method, vote_dist=vote_dist, seethrough_method=seethrough_method, suffix=suffix, copy_start=True, use_starters=True, backward=False, consensus_threshold=consensus_threshold, blur_sigma=blur_sigma, kernel_size=kernel_size, ) scheduler.register_job( align_block_job_forv, job_name=f"Forward Align {block} {block_bcube}", ) scheduler.execute_until_completion() corgie_logger.debug("Done!") if restart_stage <= 1: corgie_logger.debug("Aligning stitching blocks...") for stitch_block in stitch_blocks: block_bcube = stitch_block.get_bcube(bcube) # These blocks will have block-aligned images, but not # the block_fields that warped them. align_block_job_forv = AlignBlockJob( src_stack=deepcopy(stitch_block.src_stack), dst_stack=deepcopy(stitch_block.dst_stack), bcube=block_bcube, render_method=render_method, cf_method=cf_method, vote_dist=vote_dist, seethrough_method=seethrough_method, suffix=stitch_estimated_suffix, copy_start=False, use_starters=False, backward=False, consensus_threshold=consensus_threshold, blur_sigma=blur_sigma, kernel_size=kernel_size, ) scheduler.register_job( align_block_job_forv, job_name=f"Stitch Align {stitch_block} {block_bcube}", ) scheduler.execute_until_completion() corgie_logger.debug("Done!") # Add in the stitch_estimated fields that were just created above even_stack.create_sublayer( stitch_estimated_name, layer_type="field", overwrite=False, ) odd_stack.create_sublayer( stitch_estimated_name, layer_type="field", overwrite=False, ) if restart_stage <= 2: if stitch_size > 1: corgie_logger.debug("Voting over stitching blocks") stitch_corrected_field = dst_stack.create_sublayer( stitch_corrected_name, layer_type="field", overwrite=True) for stitch_block in stitch_blocks: stitch_estimated_field = stitch_block.dst_stack[ stitch_estimated_name] block_bcube = bcube.reset_coords( zs=stitch_block.start, ze=stitch_block.start + 1, in_place=False, ) z_offsets = [ z - block_bcube.z_range()[0] for z in range(stitch_block.start, stitch_block.stop) ] vote_stitch_job = VoteJob( input_fields=[stitch_estimated_field], output_field=stitch_corrected_field, chunk_xy=chunk_xy, bcube=block_bcube, z_offsets=z_offsets, mip=processor_mip[-1], consensus_threshold=consensus_threshold, blur_sigma=blur_sigma, kernel_size=kernel_size, ) scheduler.register_job( vote_stitch_job, job_name=f"Stitching Vote {stitch_block} {block_bcube}", ) scheduler.execute_until_completion() corgie_logger.debug("Done!") for stitch_block in stitch_blocks: block_bcube = bcube.reset_coords(zs=stitch_block.start, ze=stitch_block.start + 1, in_place=False) field_to_downsample = stitch_block.dst_stack[stitch_estimated_name] if stitch_corrected_field is not None: field_to_downsample = stitch_corrected_field # Hack for fafb field_info = field_to_downsample.get_info() for scale in field_info['scales']: scale['chunk_sizes'][-1][-1] = 1 scale['encoding'] = 'raw' field_to_downsample.cv.store_info(field_info) field_to_downsample.cv.fetch_info() downsample_field_job = DownsampleJob( src_layer=field_to_downsample, mip_start=processor_mip[-1], mip_end=max_blur_mip, bcube=block_bcube, chunk_xy= chunk_xy, # TODO: This probably needs to be modified at highest mips chunk_z=1, mips_per_task=2, ) scheduler.register_job( downsample_field_job, job_name=f"Downsample stitching field {block_bcube}", ) scheduler.execute_until_completion() corgie_logger.debug("Done!") # Add in the block-align fields even_stack.create_sublayer( block_field_name, layer_type="field", overwrite=False, ) odd_stack.create_sublayer( block_field_name, layer_type="field", overwrite=False, ) composed_field = dst_stack.create_sublayer(composed_name, layer_type="field", overwrite=True) if (restart_stage > 2) and (stitch_size > 1): stitch_corrected_field = dst_stack.create_sublayer( stitch_corrected_name, layer_type="field", overwrite=False) if restart_stage <= 3: corgie_logger.debug("Stitching blocks...") for block, stitch_block in zip(blocks[1:], stitch_blocks): block_bcube = block.broadcastable().get_bcube(bcube) block_list = block.get_neighbors(dist=decay_dist) corgie_logger.debug(f"src_block: {block}") corgie_logger.debug(f"influencing blocks: {block_list}") z_list = [b.stop for b in block_list] # stitch_corrected_field used if there is multi-section block overlap, # which requires voting to produce a corrected field. # If there is only single-section block overlap, then use # stitch_estimated_fields from each stitch_block if stitch_corrected_field is not None: stitching_fields = [stitch_corrected_field] else: # Order with furthest block first (convention of FieldSet). stitching_fields = [ stitch_block.dst_stack[stitch_estimated_name], stitch_block.src_stack[stitch_estimated_name], ] broadcast_job = BroadcastJob( block_field=block.dst_stack[block_field_name], stitching_fields=stitching_fields, output_field=composed_field, chunk_xy=chunk_xy, bcube=block_bcube, pad=pad, z_list=z_list, mip=processor_mip[-1], decay_dist=decay_dist, blur_rate=blur_rate, ) scheduler.register_job(broadcast_job, job_name=f"Broadcast {block} {block_bcube}") scheduler.execute_until_completion() corgie_logger.debug("Done!") if len(blocks) > 1: block_bcube = blocks[0].get_bcube(bcube) copy_job = CopyLayerJob( src_layer=even_stack[block_field_name], dst_layer=composed_field, mip=processor_mip[-1], bcube=block_bcube, chunk_xy=chunk_xy, chunk_z=1, ) scheduler.register_job( copy_job, job_name=f"Copy first block_field to composed_field location") scheduler.execute_until_completion() corgie_logger.debug("Done!") if restart_stage <= 4: if len(blocks) == 1: block_bcube = blocks[0].get_bcube(bcube) render_job = RenderJob( src_stack=src_stack, dst_stack=dst_stack, mips=processor_mip[-1], pad=pad, bcube=block_bcube, chunk_xy=chunk_xy, chunk_z=1, render_masks=True, blackout_masks=False, additional_fields=[even_stack[block_field_name]], ) scheduler.register_job( render_job, job_name=f"Render first block {block_bcube}") else: block_bcube = bcube.reset_coords(zs=blocks[0].start, ze=blocks[-1].stop, in_place=False) render_job = RenderJob( src_stack=src_stack, dst_stack=dst_stack, mips=processor_mip[-1], pad=pad, bcube=block_bcube, chunk_xy=chunk_xy, chunk_z=1, render_masks=True, blackout_masks=True, additional_fields=[composed_field], ) scheduler.register_job(render_job, job_name=f"Render all blocks {block_bcube}") scheduler.execute_until_completion() corgie_logger.debug("Done!") result_report = ( f"Aligned layers {[str(l) for l in src_stack.get_layers_of_type('img')]}. " f"Results in {[str(l) for l in dst_stack.get_layers_of_type('img')]}") corgie_logger.info(result_report)
def seethrough_block( ctx, src_layer_spec, dst_folder, chunk_xy, start_coord, end_coord, coord_mip, suffix, seethrough_spec, seethrough_limit, seethrough_spec_mip, force_chunk_z=1, ): scheduler = ctx.obj["scheduler"] if suffix is None: suffix = "_seethrough" else: suffix = f"_{suffix}" crop, pad = 0, 0 corgie_logger.debug("Setting up layers...") src_stack = create_stack_from_spec(src_layer_spec, name="src", readonly=True) src_stack.folder = dst_folder dst_stack = stack.create_stack_from_reference( reference_stack=src_stack, folder=dst_folder, name="dst", types=["img", "mask"], readonly=False, suffix=suffix, overwrite=True, force_chunk_z=force_chunk_z, ) render_method = helpers.PartialSpecification( f=RenderJob, pad=pad, chunk_xy=chunk_xy, chunk_z=1, render_masks=False, ) seethrough_method = helpers.PartialSpecification( f=SeethroughCompareJob, mip=seethrough_spec_mip, processor_spec=seethrough_spec, chunk_xy=chunk_xy, pad=pad, crop=pad, seethrough_limit=seethrough_limit, ) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) seethrough_block_job = SeethroughBlockJob( src_stack=src_stack, dst_stack=dst_stack, bcube=bcube, render_method=render_method, seethrough_method=seethrough_method, suffix=suffix, ) # create scheduler and execute the job scheduler.register_job(seethrough_block_job, job_name="Seethrough Block {}".format(bcube)) scheduler.execute_until_completion() result_report = ( f"Rendered layers {[str(l) for l in src_stack.get_layers_of_type('img')]}. " f"Results in {[str(l) for l in dst_stack.get_layers_of_type('img')]}") corgie_logger.info(result_report)
def execute(self): device = "cuda" if torch.cuda.is_available() else "cpu" # Field padding padded_bcube = self.bcube.uncrop(self.pad, self.mip) for k, specs in enumerate(self.src_specs[::-1]): src_z = specs["src_z"] dst_z = self.bcube.z_range()[0] corgie_logger.info(f"Load fields for {padded_bcube}") # backwards compatible if not isinstance(specs["src_field"], list): specs["src_field"] = [specs["src_field"]] mask_layer = self.src_layers[str(specs["src_mask"])] field_ids = list(map(str, specs["src_field"])) corgie_logger.info(f"field ids={field_ids}") z_list = specs.get("src_field_z", [src_z] * len(field_ids)) fields = FieldSet([self.src_layers[n] for n in field_ids]) field = fields.read( bcube=padded_bcube, z_list=z_list, mip=self.mip, device=device ) bcube = padded_bcube.reset_coords(zs=src_z, ze=src_z + 1, in_place=False) # Extend image/mask cutout to account for field spread render_pad = int((field.max_vector() - field.min_vector()).max().ceil().tensor().item()) snap_factor = 2 ** (max(self.mip, mask_layer.data_mip) - self.mip) render_pad = math.ceil(render_pad / snap_factor) * snap_factor render_pad = min(render_pad, 4096) # Safety render_bcube = bcube.uncrop(render_pad, self.mip) corgie_logger.debug(f"render_pad: {render_pad}") # Move image/mask cutout to account for field drift img_trans = helpers.percentile_trans_adjuster(field) mask_trans = img_trans.round_to_mip(self.mip, mask_layer.data_mip) corgie_logger.debug(f"img_trans: {img_trans} | mask_trans: {mask_trans}") img_bcube = render_bcube.translate( x_offset=img_trans.y, y_offset=img_trans.x, mip=self.mip ) mask_bcube = render_bcube.translate( x_offset=mask_trans.y, y_offset=mask_trans.x, mip=self.mip ) if render_pad > 0: field = torch.nn.functional.pad(field, [render_pad, render_pad, render_pad, render_pad], mode='replicate') corgie_logger.info(f"Load masks for {mask_bcube}") mask_id = specs["mask_id"] mask_layer.binarizer = helpers.Binarizer(["eq", mask_id]) mask = mask_layer.read(bcube=mask_bcube, mip=self.mip, device=device) mask = residuals.res_warp_img( mask.float(), field - mask_trans.to_tensor(device=field.device) ).tensor() mask = (mask > 0.4).bool() cropped_mask = helpers.crop(mask, self.pad + render_pad) corgie_logger.info(f"Load image for {img_bcube}") if cropped_mask.sum() == 0: cropped_img = torch.zeros_like(cropped_mask, dtype=torch.float) else: img_layer = self.src_layers[str(specs["src_img"])] img = img_layer.read(bcube=img_bcube, mip=self.mip, device=device) img = residuals.res_warp_img( img.float(), field - img_trans.to_tensor(device=field.device) ) cropped_img = helpers.crop(img, self.pad + render_pad) # write to composite image if k == 0: dst_img = cropped_img dst_img[~cropped_mask] = 0 else: dst_img[cropped_mask] = cropped_img[cropped_mask] self.dst_layer.write(dst_img.cpu(), bcube=self.bcube, mip=self.mip)
def align_block( ctx, src_layer_spec, dst_folder, vote_dist, render_pad, render_chunk_xy, processor_spec, pad, crop, processor_mip, chunk_xy, start_coord, end_coord, coord_mip, blend_xy, force_chunk_xy, suffix, copy_start, use_starters, seethrough_spec, seethrough_limit, seethrough_spec_mip, mode, chunk_z=1, ): scheduler = ctx.obj["scheduler"] if suffix is None: suffix = "_aligned" else: suffix = f"_{suffix}" if crop is None: crop = pad corgie_logger.debug("Setting up layers...") src_stack = create_stack_from_spec(src_layer_spec, name="src", readonly=True) src_stack.folder = dst_folder force_chunk_xy = chunk_xy if force_chunk_xy else None dst_stack = stack.create_stack_from_reference( reference_stack=src_stack, folder=dst_folder, name="dst", types=["img", "mask"], readonly=False, suffix=suffix, force_chunk_xy=force_chunk_xy, overwrite=True, ) render_method = helpers.PartialSpecification( f=RenderJob, pad=render_pad, chunk_xy=render_chunk_xy, chunk_z=1, render_masks=False, ) if seethrough_spec != tuple(): assert seethrough_spec_mip is not None seethrough_method = helpers.PartialSpecification( f=SeethroughCompareJob, mip=seethrough_spec_mip, processor_spec=seethrough_spec, chunk_xy=chunk_xy, pad=pad, crop=pad, seethrough_limit=seethrough_limit, ) else: seethrough_method = None cf_method = helpers.PartialSpecification( f=ComputeFieldJob, pad=pad, crop=crop, processor_mip=processor_mip, processor_spec=processor_spec, chunk_xy=chunk_xy, blend_xy=blend_xy, chunk_z=1, ) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) if mode == "bidirectional": z_mid = (bcube.z_range()[1] + bcube.z_range()[0]) // 2 bcube_back = bcube.reset_coords(ze=z_mid, in_place=False) bcube_forv = bcube.reset_coords(zs=z_mid, in_place=False) align_block_job_back = AlignBlockJob( src_stack=src_stack, dst_stack=dst_stack, bcube=bcube_back, render_method=render_method, cf_method=cf_method, seethrough_method=seethrough_method, suffix=suffix, copy_start=copy_start, backward=True, vote_dist=vote_dist, use_starters=use_starters, ) scheduler.register_job( align_block_job_back, job_name="Backward Align Block {}".format(bcube), ) align_block_job_forv = AlignBlockJob( src_stack=src_stack, dst_stack=deepcopy(dst_stack), bcube=bcube_forv, render_method=render_method, cf_method=cf_method, seethrough_method=seethrough_method, suffix=suffix, copy_start=True, backward=False, vote_dist=vote_dist, use_starters=use_starters, ) scheduler.register_job( align_block_job_forv, job_name="Forward Align Block {}".format(bcube), ) else: align_block_job = AlignBlockJob( src_stack=src_stack, dst_stack=dst_stack, bcube=bcube, render_method=render_method, cf_method=cf_method, seethrough_method=seethrough_method, suffix=suffix, copy_start=copy_start, backward=mode == "backward", vote_dist=vote_dist, use_starters=use_starters, ) # create scheduler and execute the job scheduler.register_job(align_block_job, job_name="Align Block {}".format(bcube)) scheduler.execute_until_completion() result_report = ( f"Aligned layers {[str(l) for l in src_stack.get_layers_of_type('img')]}. " f"Results in {[str(l) for l in dst_stack.get_layers_of_type('img')]}" ) corgie_logger.info(result_report)
def align_block(ctx, src_layer_spec, tgt_layer_spec, dst_folder, render_pad, render_chunk_xy, processor_spec, pad, crop, processor_mip, chunk_xy, start_coord, end_coord, coord_mip, suffix, copy_start, mode, chunk_z=1): scheduler = ctx.obj['scheduler'] if suffix is None: suffix = '_aligned' else: suffix = f"_{suffix}" if crop is None: crop = pad corgie_logger.debug("Setting up layers...") src_stack = create_stack_from_spec(src_layer_spec, name='src', readonly=True) tgt_stack = create_stack_from_spec(tgt_layer_spec, name='tgt', readonly=True, reference=src_stack) dst_stack = stack.create_stack_from_reference(reference_stack=src_stack, folder=dst_folder, name="dst", types=["img", "mask"], readonly=False, suffix=suffix) render_method = helpers.PartialSpecification( f=RenderJob, pad=render_pad, chunk_xy=render_chunk_xy, chunk_z=1, blackout_masks=False, render_masks=True, mip=min(processor_mip) ) cf_method = helpers.PartialSpecification( f=ComputeFieldJob, pad=pad, crop=crop, processor_mip=processor_mip, processor_spec=processor_spec, chunk_xy=chunk_xy, chunk_z=1 ) bcube = get_bcube_from_coords(start_coord, end_coord, coord_mip) if mode == 'bidirectional': z_mid = (bcube.z_range()[1] + bcube.z_range()[0]) // 2 bcube_back = bcube.reset_coords(ze=z_mid, in_place=False) bcube_forv = bcube.reset_coords(zs=z_mid, in_place=False) align_block_job_back = AlignBlockJob(src_stack=src_stack, tgt_stack=tgt_stack, dst_stack=dst_stack, bcube=bcube_back, render_method=render_method, cf_method=cf_method, suffix=suffix, copy_start=copy_start, backward=True) scheduler.register_job(align_block_job_back, job_name="Backward Align Block {}".format(bcube)) align_block_job_forv = AlignBlockJob(src_stack=src_stack, tgt_stack=tgt_stack, dst_stack=deepcopy(dst_stack), bcube=bcube_forv, render_method=render_method, cf_method=cf_method, suffix=suffix, copy_start=True, backward=False) scheduler.register_job(align_block_job_forv, job_name="Forward Align Block {}".format(bcube)) else: align_block_job = AlignBlockJob(src_stack=src_stack, tgt_stack=tgt_stack, dst_stack=dst_stack, bcube=bcube, render_method=render_method, cf_method=cf_method, suffix=suffix, copy_start=copy_start, backward=mode=='backward') # create scheduler and execute the job scheduler.register_job(align_block_job, job_name="Align Block {}".format(bcube)) scheduler.execute_until_completion() result_report = f"Aligned layers {[str(l) for l in src_stack.get_layers_of_type('img')]}. " \ f"Results in {[str(l) for l in dst_stack.get_layers_of_type('img')]}" corgie_logger.info(result_report)