|
| 1 | +import glob |
| 2 | +import os |
| 3 | +import re |
| 4 | + |
| 5 | +import numpy as np |
| 6 | +from numcodecs import blosc |
| 7 | + |
| 8 | +from aind_exaspim_image_compression.inference import ( |
| 9 | + build_volume_transform, |
| 10 | + load_model, |
| 11 | + predict, |
| 12 | +) |
| 13 | +from aind_exaspim_image_compression.utils import img_util, util |
| 14 | + |
| 15 | + |
| 16 | +def find_best_checkpoint(session_dir): |
| 17 | + """ |
| 18 | + Returns the best checkpoint written by a training session. |
| 19 | +
|
| 20 | + Checkpoints are named ``BM4DNet-<date>-<epoch>-<score>.pth`` by |
| 21 | + Trainer.save_model, where <score> is the checkpoint-selection score at save |
| 22 | + time and lower is better. The score can be negative (the cratio term is |
| 23 | + subtracted), so it is parsed with a regex that allows a leading minus rather |
| 24 | + than by splitting on "-". |
| 25 | +
|
| 26 | + Parameters |
| 27 | + ---------- |
| 28 | + session_dir : str |
| 29 | + Training session directory holding the BM4DNet-*.pth checkpoints. |
| 30 | +
|
| 31 | + Returns |
| 32 | + ------- |
| 33 | + str |
| 34 | + Path to the lowest-scoring (best) checkpoint. |
| 35 | + """ |
| 36 | + paths = glob.glob(os.path.join(session_dir, "BM4DNet-*.pth")) |
| 37 | + if not paths: |
| 38 | + raise FileNotFoundError( |
| 39 | + f"No BM4DNet-*.pth checkpoints found in {session_dir}" |
| 40 | + ) |
| 41 | + |
| 42 | + def score(path): |
| 43 | + m = re.search(r"-(-?\d+\.\d+)\.pth$", os.path.basename(path)) |
| 44 | + if m is None: |
| 45 | + raise ValueError(f"Cannot parse score from checkpoint: {path}") |
| 46 | + return float(m.group(1)) |
| 47 | + |
| 48 | + return min(paths, key=score) |
| 49 | + |
| 50 | + |
| 51 | +def evaluate(): |
| 52 | + # Resolve the checkpoint to evaluate (explicit path wins over auto-select). |
| 53 | + ckpt_path = checkpoint_path or find_best_checkpoint(session_dir) |
| 54 | + print("Checkpoint:", ckpt_path) |
| 55 | + |
| 56 | + # Load the model together with the intensity transform it was trained with |
| 57 | + # (load_model rebuilds the transform from the checkpoint's "transform" cfg). |
| 58 | + model, transform = load_model(ckpt_path, device=device) |
| 59 | + print("Transform:", transform.cfg) |
| 60 | + |
| 61 | + # Read the image. img_util.read handles s3://, gs://, and local zarr; point |
| 62 | + # img_path at a single 5D (t, c, z, y, x) multiscale level array (e.g. |
| 63 | + # ".../image.zarr/0"). Slicing the lazy zarr in get_patch fetches only the |
| 64 | + # requested region, so a crop avoids pulling the whole (huge) volume from S3. |
| 65 | + img = img_util.read(img_path) |
| 66 | + source_transform = img_util.get_ome_zarr_level_transform(img_path) |
| 67 | + source_scale = np.asarray(source_transform["scale"]) |
| 68 | + source_translation = np.asarray(source_transform["translation"]) |
| 69 | + crop_start = (0, 0, 0) |
| 70 | + if crop_center is not None: |
| 71 | + # Neuroglancer reports transformed spatial coordinates in (x, y, z) |
| 72 | + # order. Convert them to this level's integer (z, y, x) voxel indices. |
| 73 | + crop_center_voxel = img_util.ome_zarr_coordinate_to_voxel( |
| 74 | + crop_center, source_transform |
| 75 | + ) |
| 76 | + source_offset_zyx = source_translation[2:] / source_scale[2:] |
| 77 | + snapped_center_zyx = source_offset_zyx + np.asarray( |
| 78 | + crop_center_voxel |
| 79 | + ) |
| 80 | + crop_start, _ = img_util.get_start_end( |
| 81 | + crop_center_voxel, crop_shape, is_center=True |
| 82 | + ) |
| 83 | + crop_origin_zyx = source_offset_zyx + np.asarray(crop_start) |
| 84 | + crop_end = np.asarray(crop_start) + np.asarray(crop_shape) |
| 85 | + if np.any(np.asarray(crop_start) < 0) or np.any( |
| 86 | + crop_end > np.asarray(img.shape[2:]) |
| 87 | + ): |
| 88 | + raise ValueError( |
| 89 | + "Crop is outside the source level: " |
| 90 | + f"start={tuple(crop_start)}, end={tuple(crop_end)}, " |
| 91 | + f"source_shape={tuple(img.shape[2:])}" |
| 92 | + ) |
| 93 | + print( |
| 94 | + "Requested Neuroglancer crop center (x, y, z):", |
| 95 | + tuple(crop_center), |
| 96 | + ) |
| 97 | + print( |
| 98 | + "Neuroglancer spatial scale (x, y, z):", |
| 99 | + tuple(source_scale[2:][::-1].tolist()), |
| 100 | + source_transform["spatial_unit"], |
| 101 | + ) |
| 102 | + print("Crop center voxel (z, y, x):", crop_center_voxel) |
| 103 | + print( |
| 104 | + "Snapped crop center (x, y, z):", |
| 105 | + tuple(snapped_center_zyx[::-1].tolist()), |
| 106 | + ) |
| 107 | + print("Crop origin (z, y, x):", tuple(crop_start)) |
| 108 | + print( |
| 109 | + "Neuroglancer crop origin (x, y, z):", |
| 110 | + tuple(crop_origin_zyx[::-1].tolist()), |
| 111 | + ) |
| 112 | + raw = np.asarray( |
| 113 | + img_util.get_patch( |
| 114 | + img, crop_center_voxel, crop_shape, is_center=True |
| 115 | + ) |
| 116 | + ) |
| 117 | + else: |
| 118 | + raw = np.asarray(img[0, 0]) |
| 119 | + print("Volume shape:", raw.shape) |
| 120 | + |
| 121 | + # For a raw (non-background-subtracted) volume, use the supplied full-tile |
| 122 | + # offset. With background_offset=None, fall back to estimating from this |
| 123 | + # test subvolume for debugging only. |
| 124 | + if raw_input: |
| 125 | + volume_transform = build_volume_transform( |
| 126 | + transform, |
| 127 | + raw, |
| 128 | + percentile=0.1, |
| 129 | + offset=background_offset, |
| 130 | + ) |
| 131 | + print("Per-volume transform:", volume_transform.cfg) |
| 132 | + else: |
| 133 | + volume_transform = transform |
| 134 | + |
| 135 | + # Denoise the whole volume via overlapping tiled prediction. |
| 136 | + denoised = predict(raw, model, volume_transform, batch_size=batch_size) |
| 137 | + |
| 138 | + # Compression ratio, raw vs denoised, with the codec Zarr uses to store |
| 139 | + # chunks. clevel=5 matches the training-time codec (train.py). |
| 140 | + codec = blosc.Blosc(cname="zstd", clevel=clevel, shuffle=blosc.SHUFFLE) |
| 141 | + raw_cratio = img_util.compute_cratio(raw, codec) |
| 142 | + denoised_cratio = img_util.compute_cratio(denoised, codec) |
| 143 | + print(f"cratio (raw): {raw_cratio}") |
| 144 | + print(f"cratio (denoised): {denoised_cratio}") |
| 145 | + print(f"cratio gain: {denoised_cratio / raw_cratio:.2f}x") |
| 146 | + |
| 147 | + # Save side-by-side MIPs (XY/XZ/YZ) of the raw and denoised volumes. |
| 148 | + util.mkdir(output_dir) |
| 149 | + img_util.plot_mips( |
| 150 | + raw, output_path=os.path.join(output_dir, "raw_mips.png") |
| 151 | + ) |
| 152 | + img_util.plot_mips( |
| 153 | + denoised, output_path=os.path.join(output_dir, "denoised_mips.png") |
| 154 | + ) |
| 155 | + print("MIPs written to:", output_dir) |
| 156 | + |
| 157 | + # Optionally persist the denoised volume as a Zarr array. output_zarr may be |
| 158 | + # a local path or a cloud URL (s3://.../denoised.zarr); it is written with |
| 159 | + # the same zstd/clevel codec used to measure cratio, and reads back via |
| 160 | + # img_util.read at "<output_zarr>" (a plain array, no "/0" suffix). Writing |
| 161 | + # to S3 needs credentials (the default AWS chain), unlike the anonymous |
| 162 | + # public read of the input. |
| 163 | + if output_zarr is not None: |
| 164 | + crop_offset = np.asarray([0, 0, *crop_start]) |
| 165 | + output_translation = source_translation + source_scale * crop_offset |
| 166 | + print( |
| 167 | + "Output OME transform (t, c, z, y, x):", |
| 168 | + { |
| 169 | + "scale": tuple(source_scale.tolist()), |
| 170 | + "translation": tuple(output_translation.tolist()), |
| 171 | + "unit": source_transform["spatial_unit"], |
| 172 | + }, |
| 173 | + ) |
| 174 | + img_util.write_ome_zarr( |
| 175 | + denoised, |
| 176 | + output_zarr, |
| 177 | + scale=source_scale, |
| 178 | + translation=output_translation, |
| 179 | + spatial_unit=source_transform["spatial_unit"], |
| 180 | + ) |
| 181 | + print("Denoised Zarr written to:", output_zarr) |
| 182 | + |
| 183 | + |
| 184 | +if __name__ == "__main__": |
| 185 | + # Checkpoint. Point session_dir at a training session (the folder holding |
| 186 | + # the BM4DNet-*.pth files) to auto-select the best checkpoint. Set |
| 187 | + # checkpoint_path to a .pth to evaluate that file explicitly instead. |
| 188 | + session_dir = "/root/capsule/results/training-sessions/session-20260710_1719" |
| 189 | + checkpoint_path = "/root/capsule/results/training-sessions/session-20260710_1719/BM4DNet-20260710-499--19.965923.pth" |
| 190 | + |
| 191 | + # Test image. Any zarr readable by img_util.read, including an s3:// path; |
| 192 | + # give the full path to a single 5D multiscale level array. |
| 193 | + img_path = "s3://aind-open-data/exaSPIM_826511_2026-06-02_15-10-47/SPIM.ome.zarr/tile_000010_ch_488.zarr/0" |
| 194 | + |
| 195 | + # Region to evaluate. crop_center is the numeric (x, y, z) position shown by |
| 196 | + # Neuroglancer; the physical scale displayed beside each coordinate is read |
| 197 | + # from the source OME-Zarr. The position is converted to the nearest source |
| 198 | + # voxel before cropping. Each crop_shape dimension must be >= the model patch |
| 199 | + # size (64). Set crop_center=None only for a small, pre-cropped input volume. |
| 200 | + crop_center = (22464, -15914, 18711) |
| 201 | + crop_shape = (1024, 1024, 1024) |
| 202 | + |
| 203 | + # Use raw_input=True for volumes that were not background-subtracted. |
| 204 | + raw_input = True |
| 205 | + # Prefer the background offset precomputed from the full image tile's |
| 206 | + # lower-resolution data. None estimates from this test subvolume instead. |
| 207 | + background_offset = 37 |
| 208 | + |
| 209 | + # Output + misc |
| 210 | + output_dir = "/results/evaluation" |
| 211 | + # Where to persist the denoised volume as an OME-Zarr. Local path or a |
| 212 | + # cloud path (e.g. "s3://BUCKET/PATH/denoised.zarr"). Set to None to skip. |
| 213 | + output_zarr = "s3://aind-scratch-data/cameron.arshadi/denoising-experiments/outputs/BM4DNet-20260710-499--19.965923/826511_raw_crop.zarr" |
| 214 | + device = "cuda" |
| 215 | + batch_size = 32 |
| 216 | + clevel = 5 |
| 217 | + |
| 218 | + evaluate() |
0 commit comments