diff --git a/specs/tommy/lee_fly_cns_segmentation_test001/230608_adjust_affinities.cue b/specs/tommy/lee_fly_cns_segmentation_test001/230608_adjust_affinities.cue new file mode 100644 index 000000000..b62751bb7 --- /dev/null +++ b/specs/tommy/lee_fly_cns_segmentation_test001/230608_adjust_affinities.cue @@ -0,0 +1,83 @@ +#AFF_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity" +#BACKUP_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity/backup_for_adjustments" + +// #AFF_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity/cutouts/test001" +// #BACKUP_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity/cutouts/test001/backup" +#BLACKOUT_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/resin_mask_final_surgery" +#SNAP_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/img_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_surgery_consecutive3" +#THRESHOLD_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/aff_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_surgery" + +"@type": "mazepa.execute_on_gcp_with_sqs" +worker_image: "us.gcr.io/zetta-lee-fly-vnc-001/zetta_utils:tmacrina_mask_affinities_x3" +worker_cluster_name: "zutils-cns" +worker_cluster_region: "us-east1" +worker_cluster_project: "zetta-lee-fly-vnc-001" +worker_resources: { + memory: "18560Mi" +} +worker_replicas: 512 +local_test: false + +#AFF: { + "@type": "build_cv_layer" + path: #AFF_PATH + cv_kwargs: {"delete_black_uploads": false} // ws+agg require all affinities present, so cannot issue DELETES +} +#BACKUP: { + "@type": "build_cv_layer" + path: #BACKUP_PATH + info_reference_path: #AFF_PATH +} +#BLACKOUT_MASK: { + "@type": "build_cv_layer" + path: #BLACKOUT_MSK_PATH + data_resolution: [256, 256, 45] + interpolation_mode: "mask" +} +#SNAP_MASK: { + "@type": "build_cv_layer" + path: #SNAP_MSK_PATH + data_resolution: [64, 64, 45] + interpolation_mode: "mask" +} +#THRESHOLD_MASK: { + "@type": "build_cv_layer" + path: #THRESHOLD_MSK_PATH + data_resolution: [64, 64, 45] + interpolation_mode: "mask" +} + +#BBOX: { + "@type": "BBox3D.from_coords" + // start_coord: [1341 - 144, 288, 2990] + // end_coord: [1341, 288 + 144, 2990 + 2*13] + // resolution: [256, 256, 45] + start_coord: [0, 0, 0] + // end_coord: [4096, 4608, 7010] // actual + end_coord: [4176, 4752, 7020] // divisible by processing_chunk_size + resolution: [256, 256, 45] +} + +target: { + "@type": "build_subchunkable_apply_flow" + // "@mode": "partial" + bbox: #BBOX + dst: #AFF + dst_resolution: [16, 16, 45] + processing_chunk_sizes: [[144 * 16, 144 * 16, 13]] + // must be one pixel for lowest-res mask in xy and >=1 in z for snap mask breaking + processing_crop_pads: [[16, 16, 1]] + processing_blend_pads: [[0, 0, 0]] + op: { + "@type": "AdjustAffinitiesOp" + } + op_kwargs: { + aff_layer: #AFF + aff_backup_layer: #BACKUP + blackout_mask_layer: #BLACKOUT_MASK + snap_mask_layer: #SNAP_MASK + threshold_mask_layer: #THRESHOLD_MASK + threshold_value: 0.85 + fill_value: 0 + } +} diff --git a/specs/tommy/lee_fly_cns_segmentation_test001/230608_detect_three_image_masks.cue b/specs/tommy/lee_fly_cns_segmentation_test001/230608_detect_three_image_masks.cue new file mode 100644 index 000000000..9b5898f12 --- /dev/null +++ b/specs/tommy/lee_fly_cns_segmentation_test001/230608_detect_three_image_masks.cue @@ -0,0 +1,55 @@ +// Create a "snap" mask at locations of three consecutive image masks + +#SRC_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/img_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_surgery" +#DST_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/img_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_surgery_consecutive3" + +#BBOX: { + "@type": "BBox3D.from_coords" + // start_coord: [9216, 0, 4000] + // end_coord: [10240, 1024, 4064] + start_coord: [0, 0, 0] + end_coord: [16384, 18432, 7010] + resolution: [64, 64, 45] +} + +"@type": "mazepa.execute_on_gcp_with_sqs" +worker_image: "us.gcr.io/zetta-lee-fly-vnc-001/zetta_utils:tmacrina_mask_affinities_x2" +worker_cluster_name: "zutils-cns" +worker_cluster_region: "us-east1" +worker_cluster_project: "zetta-lee-fly-vnc-001" +worker_resources: { + memory: "18560Mi" +} +worker_replicas: 128 +local_test: false + +target: { + "@type": "build_subchunkable_apply_flow" + bbox: #BBOX + dst_resolution: [64, 64, 45] + processing_chunk_sizes: [[1024, 1024, 64]] + // need to manually set crop/pad for consecutive detection + processing_crop_pads: [[0, 0, 2]] + processing_blend_pads: [[0, 0, 0]] + + expand_bbox: true + + fn: { + "@type": "detect_consecutive_masks" + "@mode": "partial" + } + op_kwargs: { + src: { + "@type": "build_cv_layer" + path: #SRC_PATH + } + num_consecutive: 3 + } + + dst: { + "@type": "build_cv_layer" + path: #DST_PATH + info_reference_path: #SRC_PATH + on_info_exists: "overwrite" + } +} diff --git a/specs/tommy/lee_fly_cns_segmentation_test001/230610_detect_three_resin_masks.cue b/specs/tommy/lee_fly_cns_segmentation_test001/230610_detect_three_resin_masks.cue new file mode 100644 index 000000000..761c33479 --- /dev/null +++ b/specs/tommy/lee_fly_cns_segmentation_test001/230610_detect_three_resin_masks.cue @@ -0,0 +1,55 @@ +// Create a "snap" mask at locations of three consecutive image masks + +#SRC_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/resin_mask_final_surgery" +#DST_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/resin_mask_final_surgery_consecutive3" + +#BBOX: { + "@type": "BBox3D.from_coords" + // start_coord: [9216, 0, 4000] + // end_coord: [10240, 1024, 4064] + start_coord: [0, 0, 0] + end_coord: [4096, 4608, 7010] + resolution: [256, 256, 45] +} + +"@type": "mazepa.execute_on_gcp_with_sqs" +worker_image: "us.gcr.io/zetta-lee-fly-vnc-001/zetta_utils:tmacrina_mask_affinities_x3" +worker_cluster_name: "zutils-cns" +worker_cluster_region: "us-east1" +worker_cluster_project: "zetta-lee-fly-vnc-001" +worker_resources: { + memory: "18560Mi" +} +worker_replicas: 32 +local_test: false + +target: { + "@type": "build_subchunkable_apply_flow" + bbox: #BBOX + dst_resolution: [256, 256, 45] + processing_chunk_sizes: [[2048, 2048, 64]] + // need to manually set crop/pad for consecutive detection + processing_crop_pads: [[0, 0, 2]] + processing_blend_pads: [[0, 0, 0]] + + expand_bbox: true + + fn: { + "@type": "detect_consecutive_masks" + "@mode": "partial" + } + op_kwargs: { + src: { + "@type": "build_cv_layer" + path: #SRC_PATH + } + num_consecutive: 3 + } + + dst: { + "@type": "build_cv_layer" + path: #DST_PATH + info_reference_path: #SRC_PATH + on_info_exists: "overwrite" + } +} diff --git a/specs/tommy/lee_fly_cns_segmentation_test001/230612_copy_affinities.cue b/specs/tommy/lee_fly_cns_segmentation_test001/230612_copy_affinities.cue new file mode 100644 index 000000000..a7724aa8e --- /dev/null +++ b/specs/tommy/lee_fly_cns_segmentation_test001/230612_copy_affinities.cue @@ -0,0 +1,77 @@ +#AFF_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity" +#BACKUP_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity/backup_for_adjustments" +#COPY_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity/cutouts/test002" +#AFF_COPY_PATH: #COPY_PATH + "/affinity" +#BACKUP_COPY_PATH: #COPY_PATH + "/backup" + +#BBOX: { + "@type": "BBox3D.from_coords" + start_coord: [1602, 306, 2834] + end_coord: [1602 + 18, 306 + 18, 2834 + 13] + resolution: [256, 256, 45] +} + +"@type": "mazepa.execute_on_gcp_with_sqs" +worker_image: "us.gcr.io/zetta-lee-fly-vnc-001/zetta_utils:tmacrina_mask_affinities_x3" +worker_cluster_name: "zutils-cns" +worker_cluster_region: "us-east1" +worker_cluster_project: "zetta-lee-fly-vnc-001" +worker_resources: { + memory: "18560Mi" +} +worker_replicas: 1 +local_test: true + +#COPY_TEMPLATE: { + "@type": "build_subchunkable_apply_flow" + bbox: #BBOX + dst_resolution: [16, 16, 45] + processing_chunk_sizes: [[144 * 2, 144 * 2, 13]] + processing_crop_pads: [[0, 0, 0]] + processing_blend_pads: [[0, 0, 0]] + expand_bbox: true + + fn: { + "@type": "lambda" + lambda_str: "lambda src: src" + } + op_kwargs: _ + dst: _ +} + +target: { + "@type": "mazepa.seq_flow" + stages: [ + #COPY_TEMPLATE & { + op_kwargs: { + src: { + "@type": "build_cv_layer" + path: #AFF_PATH + } + } + dst: { + "@type": "build_cv_layer" + path: #AFF_COPY_PATH + info_reference_path: #AFF_PATH + cv_kwargs: {"delete_black_uploads": false} + } + }, + + #COPY_TEMPLATE & { + op_kwargs: { + src: { + "@type": "build_cv_layer" + path: #BACKUP_PATH + } + } + dst: { + "@type": "build_cv_layer" + path: #BACKUP_COPY_PATH + info_reference_path: #BACKUP_PATH + cv_kwargs: {"delete_black_uploads": false} + } + }, + + ] + +} diff --git a/specs/tommy/lee_fly_cns_segmentation_test001/230612_rework_affinities.cue b/specs/tommy/lee_fly_cns_segmentation_test001/230612_rework_affinities.cue new file mode 100644 index 000000000..570eb7814 --- /dev/null +++ b/specs/tommy/lee_fly_cns_segmentation_test001/230612_rework_affinities.cue @@ -0,0 +1,85 @@ +#AFF_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity" +#BACKUP_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity/backup_for_adjustments" +#ORIGINAL_BLACK_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/resin_mask_final_surgery" +#BLACK_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/resin_mask_final_surgery_consecutive3" +#SNAP_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/img_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_surgery_consecutive3" +#THRESHOLD_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/aff_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_surgery" + +"@type": "mazepa.execute_on_gcp_with_sqs" +worker_image: "us.gcr.io/zetta-lee-fly-vnc-001/zetta_utils:tmacrina_mask_affinities_x4" +worker_cluster_name: "zutils-cns" +worker_cluster_region: "us-east1" +worker_cluster_project: "zetta-lee-fly-vnc-001" +worker_resources: { + memory: "18560Mi" +} +worker_replicas: 512 +local_test: false + +#AFF: { + "@type": "build_cv_layer" + path: #AFF_PATH + cv_kwargs: {"delete_black_uploads": false} // ws+agg require all affinities present, so cannot issue DELETES +} +#BACKUP: { + "@type": "build_cv_layer" + path: #BACKUP_PATH + info_reference_path: #AFF_PATH +} +#ORIGINAL_BLACK_MASK: { + "@type": "build_cv_layer" + path: #ORIGINAL_BLACK_MSK_PATH + data_resolution: [256, 256, 45] + interpolation_mode: "mask" +} +#BLACK_MASK: { + "@type": "build_cv_layer" + path: #BLACK_MSK_PATH + data_resolution: [256, 256, 45] + interpolation_mode: "mask" +} +#SNAP_MASK: { + "@type": "build_cv_layer" + path: #SNAP_MSK_PATH + data_resolution: [64, 64, 45] + interpolation_mode: "mask" +} +#THRESHOLD_MASK: { + "@type": "build_cv_layer" + path: #THRESHOLD_MSK_PATH + data_resolution: [64, 64, 45] + interpolation_mode: "mask" +} + +#BBOX: { + "@type": "BBox3D.from_coords" + start_coord: [0, 0, 0] + // end_coord: [4096, 4608, 7010] // actual + end_coord: [4176, 4752, 7020] // divisible by processing_chunk_size + resolution: [256, 256, 45] +} + +target: { + "@type": "build_subchunkable_apply_flow" + // "@mode": "partial" + bbox: #BBOX + dst: #AFF + dst_resolution: [16, 16, 45] + processing_chunk_sizes: [[144 * 16, 144 * 16, 13]] + // must be one pixel for lowest-res mask in xy and >=1 in z for snap mask breaking + processing_crop_pads: [[16, 16, 1]] + processing_blend_pads: [[0, 0, 0]] + op: { + "@type": "ReworkAffinitiesOp" + } + op_kwargs: { + aff_layer: #AFF + aff_backup_layer: #BACKUP + original_black_mask_layer: #ORIGINAL_BLACK_MASK + black_mask_layer: #BLACK_MASK + snap_mask_layer: #SNAP_MASK + threshold_mask_layer: #THRESHOLD_MASK + threshold_value: 0.85 + fill_value: 0 + } +} diff --git a/specs/tommy/lee_fly_cns_segmentation_test001/230612_rework_affinities_test.cue b/specs/tommy/lee_fly_cns_segmentation_test001/230612_rework_affinities_test.cue new file mode 100644 index 000000000..6c837fbf7 --- /dev/null +++ b/specs/tommy/lee_fly_cns_segmentation_test001/230612_rework_affinities_test.cue @@ -0,0 +1,90 @@ +// #AFF_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity" +// #BACKUP_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity/backup_for_adjustments" + +#AFF_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity/cutouts/test002/affinity" +#BACKUP_PATH: "gs://zetta_lee_fly_cns_001_kisuk/final/v2/affinity/cutouts/test002/backup" +#ORIGINAL_BLACK_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/resin_mask_final_surgery" +#BLACK_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/resin_mask_final_surgery_consecutive3" +#SNAP_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/img_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_surgery_consecutive3" +#THRESHOLD_MSK_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/aff_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_surgery" + +"@type": "mazepa.execute_on_gcp_with_sqs" +worker_image: "us.gcr.io/zetta-lee-fly-vnc-001/zetta_utils:tmacrina_mask_affinities_x3" +worker_cluster_name: "zutils-cns" +worker_cluster_region: "us-east1" +worker_cluster_project: "zetta-lee-fly-vnc-001" +worker_resources: { + memory: "18560Mi" +} +worker_replicas: 1 +local_test: true + +#AFF: { + "@type": "build_cv_layer" + path: #AFF_PATH + cv_kwargs: {"delete_black_uploads": false} // ws+agg require all affinities present, so cannot issue DELETES +} +#BACKUP: { + "@type": "build_cv_layer" + path: #BACKUP_PATH + info_reference_path: #AFF_PATH +} +#ORIGINAL_BLACK_MASK: { + "@type": "build_cv_layer" + path: #ORIGINAL_BLACK_MSK_PATH + data_resolution: [256, 256, 45] + interpolation_mode: "mask" +} +#BLACK_MASK: { + "@type": "build_cv_layer" + path: #BLACK_MSK_PATH + data_resolution: [256, 256, 45] + interpolation_mode: "mask" +} +#SNAP_MASK: { + "@type": "build_cv_layer" + path: #SNAP_MSK_PATH + data_resolution: [64, 64, 45] + interpolation_mode: "mask" +} +#THRESHOLD_MASK: { + "@type": "build_cv_layer" + path: #THRESHOLD_MSK_PATH + data_resolution: [64, 64, 45] + interpolation_mode: "mask" +} + +#BBOX: { + "@type": "BBox3D.from_coords" + // start_coord: [1638, 486, 5837] + // end_coord: [1638 + 18, 486 + 18, 5837 + 13] + // resolution: [256, 256, 45] + start_coord: [1602, 306, 2834] + end_coord: [1602 + 18, 306 + 18, 2834 + 13] + resolution: [256, 256, 45] +} + +target: { + "@type": "build_subchunkable_apply_flow" + // "@mode": "partial" + bbox: #BBOX + dst: #AFF + dst_resolution: [16, 16, 45] + processing_chunk_sizes: [[144 * 2, 144 * 2, 13]] + // must be one pixel for lowest-res mask in xy and >=1 in z for snap mask breaking + processing_crop_pads: [[16, 16, 1]] + processing_blend_pads: [[0, 0, 0]] + op: { + "@type": "ReworkAffinitiesOp" + } + op_kwargs: { + aff_layer: #AFF + aff_backup_layer: #BACKUP + original_black_mask_layer: #ORIGINAL_BLACK_MASK + black_mask_layer: #BLACK_MASK + snap_mask_layer: #SNAP_MASK + threshold_mask_layer: #THRESHOLD_MASK + threshold_value: 0.85 + fill_value: 0 + } +} diff --git a/specs/tommy/lee_fly_cns_segmentation_test001/detect_three_image_masks.cue b/specs/tommy/lee_fly_cns_segmentation_test001/detect_three_image_masks.cue new file mode 100644 index 000000000..b6c7d1eff --- /dev/null +++ b/specs/tommy/lee_fly_cns_segmentation_test001/detect_three_image_masks.cue @@ -0,0 +1,62 @@ +// +// Handy variables +#SRC_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/img_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001" +#DST_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/img_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_3consecutive_v2" + +#BBOX: { + "@type": "BBox3D.from_coords" + start_coord: [5120, 1024, 6000] + end_coord: [7168, 3072, 6011] + resolution: [64, 64, 45] +} + +// Execution parameters +"@type": "mazepa.execute_on_gcp_with_sqs" +worker_image: "us.gcr.io/zetta-research/zetta_utils:sergiy_all_p39_x189" +worker_cluster_name: "zutils-x3" +worker_cluster_region: "us-east1" +worker_cluster_project: "zetta-research" +worker_resources: { + memory: "18560Mi" + //"nvidia.com/gpu": "1" +} +worker_replicas: 10 +local_test: true // set to `false` execute remotely + +target: { + "@type": "build_subchunkable_apply_flow" + bbox: #BBOX + + dst_resolution: [64, 64, 45] + + processing_chunk_sizes: [[1024, 1024, 12]] + // need to manually set crop/pad for consecutive detection + processing_crop_pads: [[0, 0, 2]] + processing_blend_pads: [[0, 0, 0]] + + expand_bbox: true + + // Specification for the operation we're performing + fn: { + "@type": "detect_consecutive_masks" + "@mode": "partial" + } + // Specification for the inputs to the operation + op_kwargs: { + src: { + "@type": "build_cv_layer" + path: #SRC_PATH + } + num_consecutive: 3 + } + + // Specification of the output layer. Subchunkable expects + // a single output layer. If multiple output layers are + // needed, refer to advanced examples. + dst: { + "@type": "build_cv_layer" + path: #DST_PATH + info_reference_path: #SRC_PATH + on_info_exists: "overwrite" + } +} diff --git a/specs/tommy/lee_fly_cns_segmentation_test001/upsample_add_one.cue b/specs/tommy/lee_fly_cns_segmentation_test001/upsample_add_one.cue new file mode 100644 index 000000000..ed5565bee --- /dev/null +++ b/specs/tommy/lee_fly_cns_segmentation_test001/upsample_add_one.cue @@ -0,0 +1,68 @@ +// +// Handy variables +#SRC_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/img_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_3consecutive" +#DST_PATH: "gs://zetta_lee_fly_cns_001_alignment_temp/aced/med_x1/try_x0/img_mask_stage1_v2_64nm_try_x1_iter300_rig200_lr0.001_3consecutive_semantic_v2" + +#BBOX: { + "@type": "BBox3D.from_coords" + start_coord: [1024 * 20, 8 * 1024, 2500] + end_coord: [38 * 1024, 12 * 1024, 3500] + resolution: [16, 16, 45] +} + +// Execution parameters +"@type": "mazepa.execute_on_gcp_with_sqs" +worker_image: "us.gcr.io/zetta-research/zetta_utils:sergiy_all_p39_x189" +worker_cluster_name: "zutils-x3" +worker_cluster_region: "us-east1" +worker_cluster_project: "zetta-research" +worker_resources: { + memory: "18560Mi" + //"nvidia.com/gpu": "1" +} +worker_replicas: 10 +local_test: false // set to `false` execute remotely + +target: { + // We're applying subchunkable processing flow, + "@type": "build_subchunkable_apply_flow" + bbox: #BBOX + + // What resolution is our destination? + dst_resolution: [16, 16, 45] + + // How do we chunk/crop/blend? + processing_chunk_sizes: [[4 * 1024, 4 * 1024, 1]] + processing_crop_pads: [[0, 0, 0]] + processing_blend_pads: [[0, 0, 0]] + + // We want to expand the input bbox to be evenly divisible + // by chunk size + expand_bbox: true + + // Specification for the operation we're performing + fn: { + "@type": "lambda" + lambda_str: "lambda src: src + 2" + } + // Specification for the inputs to the operation + op_kwargs: { + src: { + "@type": "build_cv_layer" + path: #SRC_PATH + data_resolution: [64, 64, 45] + interpolation_mode: "mask" + } + } + + // Specification of the output layer. Subchunkable expects + // a single output layer. If multiple output layers are + // needed, refer to advanced examples. + dst: { + "@type": "build_cv_layer" + path: #DST_PATH + info_reference_path: #SRC_PATH + info_field_overrides: {"type": "segmentation"} + on_info_exists: "overwrite" + } +} diff --git a/zetta_utils/mazepa_layer_processing/__init__.py b/zetta_utils/mazepa_layer_processing/__init__.py index 77a99f52a..441874242 100644 --- a/zetta_utils/mazepa_layer_processing/__init__.py +++ b/zetta_utils/mazepa_layer_processing/__init__.py @@ -7,6 +7,7 @@ VolumetricOpProtocol, ) from . import alignment +from . import segmentation from .common import ( ChunkedApplyFlowSchema, CallableOperation, diff --git a/zetta_utils/mazepa_layer_processing/segmentation/__init__.py b/zetta_utils/mazepa_layer_processing/segmentation/__init__.py new file mode 100644 index 000000000..a9ac823b7 --- /dev/null +++ b/zetta_utils/mazepa_layer_processing/segmentation/__init__.py @@ -0,0 +1 @@ +from . import masks diff --git a/zetta_utils/mazepa_layer_processing/segmentation/masks/__init__.py b/zetta_utils/mazepa_layer_processing/segmentation/masks/__init__.py new file mode 100644 index 000000000..cc8cc0f2e --- /dev/null +++ b/zetta_utils/mazepa_layer_processing/segmentation/masks/__init__.py @@ -0,0 +1,2 @@ +from .masks import detect_consecutive_masks +from .affinities import AdjustAffinitiesOp diff --git a/zetta_utils/mazepa_layer_processing/segmentation/masks/affinities.py b/zetta_utils/mazepa_layer_processing/segmentation/masks/affinities.py new file mode 100644 index 000000000..b60c6b72c --- /dev/null +++ b/zetta_utils/mazepa_layer_processing/segmentation/masks/affinities.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +from typing import Sequence + +import attrs +import torch + +from zetta_utils import builder, mazepa, tensor_ops +from zetta_utils.geometry.vec import Vec3D +from zetta_utils.layer.volumetric.index import VolumetricIndex +from zetta_utils.layer.volumetric.layer import VolumetricLayer + + +@builder.register("AdjustAffinitiesOp") +@mazepa.taskable_operation_cls +@attrs.frozen +class AdjustAffinitiesOp: + crop_pad: Sequence[int] = (0, 0, 0) + + def get_input_resolution( # pylint: disable=no-self-use + self, dst_resolution: Vec3D[float] + ) -> Vec3D[float]: + return dst_resolution + + def with_added_crop_pad(self, crop_pad: Vec3D[int]) -> AdjustAffinitiesOp: + return attrs.evolve(self, crop_pad=Vec3D[int](*self.crop_pad) + crop_pad) + + def __call__( + self, + idx: VolumetricIndex, + dst: VolumetricLayer, + aff_layer: VolumetricLayer, + aff_backup_layer: VolumetricLayer, + blackout_mask_layer: VolumetricLayer, + snap_mask_layer: VolumetricLayer, + threshold_mask_layer: VolumetricLayer, + threshold_value: float = 0.85, + fill_value: float = 0, + rework_mode: bool = False, + rework_mask_layer: VolumetricLayer | None = None, + ) -> None: + """ + Adjust affinities using masks, making sure to backup any adjusted affinities. + + :param idx: VolumetricIndex for the operation + :param src: layer with affinities + :param dst: layer where adjusted affinities will be written sparsely + :param aff_backup_layer: layer where original affinities will be stored + if chunk is adjusted + :param aff_backup_layer: layer where original affinities will be stored + if chunk is adjusted + :param blackout_mask_layer: one-hot mask to black out all affinities + :param snap_mask_layer: one-hot mask for which segments cannot span its + boundary ("snaps" objects at its boundary) + :param threshold_mask_layer: one-hot mask which will adjust affinities below + `threshold_value` + :param threshold_value: value for thresholding + :param fill_value: value to set the adjusted affinities + :param rework_mode: bool to use rework_mask_layer; requires `rework_mask_layer` + :param rework_mask_layer: one-hot mask of locations that need to be reworked + """ + idx_padded = idx.padded(self.crop_pad) + if rework_mode: + if rework_mask_layer is None: + raise ValueError("`rework_mode` requires `rework_mask_layer` to be given.") + rework_mask = rework_mask_layer[idx_padded] + if (rework_mask != 0).sum().item() == 0: + return + blackout_mask = blackout_mask_layer[idx_padded] + snap_mask = snap_mask_layer[idx_padded] + threshold_mask = threshold_mask_layer[idx_padded] + blackout_mask_has_data = (blackout_mask != 0).sum().item() > 0 + snap_mask_has_data = (snap_mask != 0).sum().item() > 0 + threshold_mask_has_data = (threshold_mask != 0).sum().item() > 0 + any_mask_has_data = any( + [blackout_mask_has_data, snap_mask_has_data, threshold_mask_has_data] + ) + if any_mask_has_data: + aff = aff_layer[idx_padded] + # Check the backup, in case we've run this operation before. + # Otherwise, we may overwrite any previous backup. + # This relies on the rest of the adjustment functions being idempotent. + # This assumption may not be true if a restarted operation changes + # any of the parameters (e.g. later runs lower the threshold_value). + aff_backup = aff_backup_layer[idx_padded] + # A 0-valued affinity indicates there may have been a backup + aff[aff == 0] = aff_backup[aff == 0] + if aff.sum() > 0: + aff_backup = aff.clone() + if snap_mask_has_data: + aff = adjust_affinities_across_mask_boundary( + src=aff, mask=snap_mask, fill_value=fill_value + ) + if threshold_mask_has_data: + aff = adjust_thresholded_affinities_in_mask( + src=aff, + mask=threshold_mask, + threshold_value=threshold_value, + fill_value=fill_value, + ) + if blackout_mask_has_data: + blackout_mask = torch.cat([blackout_mask, blackout_mask, blackout_mask], dim=0) + aff[blackout_mask] = fill_value + aff_adjusted_mask = aff != aff_backup + # store only affinities that are different + # assumes that a 0 affinity is never changed + aff_backup[~aff_adjusted_mask] = 0 + aff_backup_layer[idx] = tensor_ops.crop(aff_backup, self.crop_pad) + dst[idx] = tensor_ops.crop(aff, self.crop_pad) + + +@builder.register("adjust_affinities_across_mask_boundary") +def adjust_affinities_across_mask_boundary( + src: torch.Tensor, + mask: torch.Tensor, + fill_value: float = 0, +) -> torch.Tensor: + """ + Adjust affinities that span masked and unmasked voxel pairs. + + Note: At each voxel, we store the affinities to neighboring voxels in + the negative cardinal directions. Meaning that location (1, 1, 1) in + the affinity map contains the following affinities: + (0, 1, 1) - (1, 1, 1) + (1, 0, 1) - (1, 1, 1) + (1, 1, 0) - (1, 1, 1) + + :param src: input Tensor + :param mask: one-hot mask matching dimensions of src + :param fill_value: value to set the adjusted affinities + """ + result = src + mask_compare = mask[0, 1:, :, :] != mask[0, :-1, :, :] + aff_mask = torch.nn.functional.pad(mask_compare, (0, 0, 0, 0, 1, 0)) + result[0, aff_mask] = fill_value + mask_compare = mask[0, :, 1:, :] != mask[0, :, :-1, :] + aff_mask = torch.nn.functional.pad(mask_compare, (0, 0, 1, 0, 0, 0)) + result[1, aff_mask] = fill_value + mask_compare = mask[0, :, :, 1:] != mask[0, :, :, :-1] + aff_mask = torch.nn.functional.pad(mask_compare, (1, 0, 0, 0, 0, 0)) + result[2, aff_mask] = fill_value + result = result.to(src.dtype) + return result + + +@builder.register("adjust_thresholded_affinities_in_mask") +def adjust_thresholded_affinities_in_mask( + src: torch.Tensor, + mask: torch.Tensor, + threshold_value: float = 0.85, + fill_value: float = 0, +) -> torch.Tensor: + """ + Adjust affinities below `threshold_value` that are stored at masked voxels. + + :param src: input Tensor + :param mask: one-hot mask matching single-channel dimensions of src + :param threshold_value: affinities below this value will be set to `fill_value` + :param fill_value: value to set the adjusted affinities + """ + result = src + aff_mask = torch.cat([mask, mask, mask], dim=0) + threshold_mask = torch.logical_and(src < threshold_value, aff_mask) + result[threshold_mask] = fill_value + return result diff --git a/zetta_utils/mazepa_layer_processing/segmentation/masks/masks.py b/zetta_utils/mazepa_layer_processing/segmentation/masks/masks.py new file mode 100644 index 000000000..d1ae8cb01 --- /dev/null +++ b/zetta_utils/mazepa_layer_processing/segmentation/masks/masks.py @@ -0,0 +1,39 @@ +import torch + +from zetta_utils import builder + + +@builder.register("detect_consecutive_masks") +def detect_consecutive_masks( + src: torch.Tensor, + num_consecutive: int = 3, +) -> torch.Tensor: + """ + Identify locations in `src` that participate in `num_consecutive` sections + of a mask. Using `num_consecutive=3`, this is identical to the boolean function: + + `(z-2 & z-1 & z) | (z-1 & z & z+1) | (z & z+1 & z+2)` + + :param src: tensor + :param num_consecutive: number z sections that a location must + be masked to be considered part of the consecutive mask + + """ + if src.shape[-1] < num_consecutive: + raise ValueError( + f"Expected the number of masks to be >={num_consecutive}, but got {src.shape[-1]}" + ) + # sum over z with kernel of size num_consecutive + # pylint: disable=invalid-name + c, x, y, z = src.shape + reshaped_src = src.reshape(x * y, c, z) + kernel = torch.ones((1, 1, num_consecutive), dtype=src.dtype) + # use full padding to help with propagation later + src_sum = torch.nn.functional.conv1d(reshaped_src, kernel, padding=num_consecutive - 1) + is_consecutive = src_sum >= num_consecutive + # propagate is_consecutive mask to all participating sections + # only needs to participate in one consecutive range to be masked + participates = torch.logical_or( + *(is_consecutive[:, :, i : z + i] for i in range(num_consecutive)) + ) + return participates.reshape(c, x, y, z).to(src.dtype)