From 5457adde39a516457abee60864533b5080fe98e4 Mon Sep 17 00:00:00 2001 From: yxd92326 Date: Thu, 26 Feb 2026 15:43:47 +0000 Subject: [PATCH 1/9] Numpy 2 allowed --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d7123d5e..57f91cb1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ dependencies = [ "icebreaker-em", "ispyb>=11.1.2", "mrcfile", - "numpy<2", + "numpy", "pillow", "plotly", "pydantic>=2", From 46dbedfc64a9219a679e640215a3eba057923561 Mon Sep 17 00:00:00 2001 From: yxd92326 Date: Mon, 2 Mar 2026 11:50:14 +0000 Subject: [PATCH 2/9] Low mem version of easymode --- .../easymode_segmentation.py | 107 ++++++++++++++++++ src/cryoemservices/services/easymode.py | 4 +- 2 files changed, 110 insertions(+), 1 deletion(-) create mode 100644 src/cryoemservices/pipeliner_plugins/easymode_segmentation.py diff --git a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py new file mode 100644 index 00000000..cd14c371 --- /dev/null +++ b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py @@ -0,0 +1,107 @@ +import os + +import mrcfile +import numpy as np +import psutil +from scipy.ndimage import zoom + + +def _segment_tomogram_instance(volume, model, tile_size, overlap, logger): + (pz, py, px) = tile_size + (oz, oy, ox) = overlap + d, h, w = volume.shape + sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox + z_boxes = max(1, (d + sz - 1) // sz) + y_boxes = max(1, (h + sy - 1) // sy) + x_boxes = max(1, (w + sx - 1) // sx) + out = np.zeros((d, h, w), dtype=np.float32) + wgt = np.zeros((d, h, w), dtype=np.float32) + for zi in range(z_boxes): + for yi in range(y_boxes): + for xi in range(x_boxes): + logger.info(f"Starting {xi} {yi} {zi}") + z_start = zi * sz - oz + y_start = yi * sy - oy + x_start = xi * sx - ox + vz0 = max(0, z_start) + vy0 = max(0, y_start) + vx0 = max(0, x_start) + vz1 = min(d, z_start + pz) + vy1 = min(h, y_start + py) + vx1 = min(w, x_start + px) + extracted = volume[vz0:vz1, vy0:vy1, vx0:vx1] + tile = np.zeros((pz, py, px), dtype=volume.dtype) + tz0 = vz0 - z_start + ty0 = vy0 - y_start + tx0 = vx0 - x_start + tile[ + tz0 : tz0 + extracted.shape[0], + ty0 : ty0 + extracted.shape[1], + tx0 : tx0 + extracted.shape[2], + ] = extracted + logger.info(f"{extracted.shape} {tile.shape}") + z_pos, y_pos, x_pos = zi * sz, yi * sy, xi * sx + tile = np.expand_dims(np.array([tile]), axis=-1) + segmented_tile = model.predict(tile, verbose=0, batch_size=1) + logger.info( + f"Segmented at {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" + ) + segmented_tile = segmented_tile.squeeze(-1) + center = segmented_tile[0, oz : oz + sz, oy : oy + sy, ox : ox + sx] + z_end = min(z_pos + sz, d) + y_end = min(y_pos + sy, h) + x_end = min(x_pos + sx, w) + az, ay, ax = z_end - z_pos, y_end - y_pos, x_end - x_pos + out[z_pos:z_end, y_pos:y_end, x_pos:x_end] += center[:az, :ay, :ax] + wgt[z_pos:z_end, y_pos:y_end, x_pos:x_end] += 1 + wgt[wgt == 0] = 1 + return out / wgt + + +def _pad_volume(volume, min_pad=16, div=32): + j, k, l = volume.shape + pads = [] + for n in (j, k, l): + total_pad = max(2 * min_pad, ((n + 2 * min_pad + div - 1) // div) * div - n) + before = total_pad // 2 + after = total_pad - before + pads.append((before, after)) + padded = np.pad(volume, pads, mode="reflect") + return padded, tuple(pads) + + +def segment_tomogram( + model, tomogram_path, tta=1, batch_size=1, model_apix=10, input_apix=10, logger=None +): + logger.info("Reading tomogram") + with mrcfile.open(tomogram_path) as m: + volume = m.data.astype(np.float32) + oj, ok, ol = volume.shape + logger.info("Read tomogram") + scale = float(input_apix) / float(model_apix) + volume = zoom(volume, scale, order=1) + _k_margin = min(int(0.2 * ok), 64) + _l_margin = min(int(0.2 * ol), 64) + volume -= np.mean(volume[:, _k_margin:-_k_margin, _l_margin:-_l_margin]) + volume /= np.std(volume[:, _k_margin:-_k_margin, _l_margin:-_l_margin]) + 1e-7 + volume, padding = _pad_volume(volume) + tile_size = (min(256, oj), min(256, ok), min(256, ol)) + overlap = ( + 0 if tile_size[0] == oj else 48, + 0 if tile_size[1] == ok else 48, + 0 if tile_size[2] == ol else 48, + ) + logger.info("Doing instance") + segmented_volume = _segment_tomogram_instance( + volume, model, tile_size, overlap, logger + ) + logger.info("Done instance") + (j0, j1), (k0, k1), (l0, l1) = padding + segmented_volume = segmented_volume[ + j0 : segmented_volume.shape[0] - j1, + k0 : segmented_volume.shape[1] - k1, + l0 : segmented_volume.shape[2] - l1, + ] + sj, sk, sl = segmented_volume.shape + segmented_volume = zoom(segmented_volume, (oj / sj, ok / sk, ol / sl), order=1) + return segmented_volume, input_apix diff --git a/src/cryoemservices/services/easymode.py b/src/cryoemservices/services/easymode.py index eb4a6267..07cd9e08 100644 --- a/src/cryoemservices/services/easymode.py +++ b/src/cryoemservices/services/easymode.py @@ -7,10 +7,10 @@ import tensorflow as tf from easymode.core import config as easymode_config from easymode.core.distribution import get_model, load_model -from easymode.segmentation.inference import segment_tomogram from pydantic import BaseModel, Field, ValidationError from workflows.recipe import wrap_subscribe +from cryoemservices.pipeliner_plugins.easymode_segmentation import segment_tomogram from cryoemservices.services.common_service import CommonService from cryoemservices.util.models import MockRW from cryoemservices.util.relion_service_options import RelionServiceOptions @@ -120,9 +120,11 @@ def easymode(self, rw, header: dict, message: dict): batch_size=easymode_params.batch_size, model_apix=model_metadata["apix"], input_apix=easymode_params.pixel_size, + logger=self.log, ) # Convert to int8 and save mrc + self.log.info("Saving output") segmented_volume = (segmented_volume * 127).astype(np.int8) with mrcfile.new(output_tomograms[feature], overwrite=True) as mrc: mrc.set_data(segmented_volume) From b7b2301ef5e5656a1f295d4ba4111e94e95c2824 Mon Sep 17 00:00:00 2001 From: yxd92326 Date: Wed, 4 Mar 2026 14:47:41 +0000 Subject: [PATCH 3/9] Shuffle things around --- pyproject.toml | 6 +- .../easymode_segmentation.py | 96 +++++++++++-------- 2 files changed, 59 insertions(+), 43 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 57f91cb1..3ab5d899 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,8 +57,12 @@ dev = [ "pytest-datafiles", "pytest-mock", ] -torch = [ +easymode = [ # "easymode@git+https://github.com/mgflast/easymode@126163fcb22e469f8fc8791d10ab5c6f0b8b9070", + "huggingface_hub", + "tensorflow[and-cuda]", +] +torch = [ "membrain-seg", "topaz-em", ] diff --git a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py index cd14c371..c6bd75e8 100644 --- a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py +++ b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py @@ -6,56 +6,64 @@ from scipy.ndimage import zoom +def segment_loop(xi, yi, zi, volume, model, tile_size, overlap, logger): + (pz, py, px) = tile_size + (oz, oy, ox) = overlap + d, h, w = volume.shape # 672, 1440, 2016 + sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox # non-overlapped patch size, 160 + + z_start = zi * sz - oz + y_start = yi * sy - oy + x_start = xi * sx - ox + vz0 = max(0, z_start) + vy0 = max(0, y_start) + vx0 = max(0, x_start) + vz1 = min(d, z_start + pz) + vy1 = min(h, y_start + py) + vx1 = min(w, x_start + px) + tile = np.zeros((pz, py, px), dtype=volume.dtype) + tile[ + vz0 - z_start : vz1 - z_start, + vy0 - y_start : vy1 - y_start, + vx0 - x_start : vx1 - x_start, + ] = volume[vz0:vz1, vy0:vy1, vx0:vx1] + tile = np.expand_dims(np.array([tile]), axis=-1) + segmented_tile = model.predict(tile, verbose=0, batch_size=1).squeeze(-1) + logger.info( + f"Segmented {xi} {yi} {zi} at {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" + ) + z_end = min(zi * sz + sz, d) + y_end = min(yi * sy + sy, h) + x_end = min(xi * sx + sx, w) + az = z_end - zi * sz + ay = y_end - yi * sy + ax = x_end - xi * sx + return ( + segmented_tile[0, oz : oz + sz, oy : oy + sy, ox : ox + sx][:az, :ay, :ax] * 127 + ).astype(np.int8) + + def _segment_tomogram_instance(volume, model, tile_size, overlap, logger): (pz, py, px) = tile_size (oz, oy, ox) = overlap - d, h, w = volume.shape - sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox - z_boxes = max(1, (d + sz - 1) // sz) - y_boxes = max(1, (h + sy - 1) // sy) - x_boxes = max(1, (w + sx - 1) // sx) - out = np.zeros((d, h, w), dtype=np.float32) - wgt = np.zeros((d, h, w), dtype=np.float32) + d, h, w = volume.shape # 672, 1440, 2016 + sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox # non-overlapped patch size, 160 + z_boxes = max(1, (d + sz - 1) // sz) # 5 + y_boxes = max(1, (h + sy - 1) // sy) # 9 + x_boxes = max(1, (w + sx - 1) // sx) # 13 + logger.info(f"{volume.shape}, {z_boxes}, {y_boxes}, {x_boxes}") + out = np.ones((d, h, w), dtype=np.int8) * 0 for zi in range(z_boxes): for yi in range(y_boxes): for xi in range(x_boxes): - logger.info(f"Starting {xi} {yi} {zi}") - z_start = zi * sz - oz - y_start = yi * sy - oy - x_start = xi * sx - ox - vz0 = max(0, z_start) - vy0 = max(0, y_start) - vx0 = max(0, x_start) - vz1 = min(d, z_start + pz) - vy1 = min(h, y_start + py) - vx1 = min(w, x_start + px) - extracted = volume[vz0:vz1, vy0:vy1, vx0:vx1] - tile = np.zeros((pz, py, px), dtype=volume.dtype) - tz0 = vz0 - z_start - ty0 = vy0 - y_start - tx0 = vx0 - x_start - tile[ - tz0 : tz0 + extracted.shape[0], - ty0 : ty0 + extracted.shape[1], - tx0 : tx0 + extracted.shape[2], - ] = extracted - logger.info(f"{extracted.shape} {tile.shape}") z_pos, y_pos, x_pos = zi * sz, yi * sy, xi * sx - tile = np.expand_dims(np.array([tile]), axis=-1) - segmented_tile = model.predict(tile, verbose=0, batch_size=1) - logger.info( - f"Segmented at {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" - ) - segmented_tile = segmented_tile.squeeze(-1) - center = segmented_tile[0, oz : oz + sz, oy : oy + sy, ox : ox + sx] z_end = min(z_pos + sz, d) y_end = min(y_pos + sy, h) x_end = min(x_pos + sx, w) - az, ay, ax = z_end - z_pos, y_end - y_pos, x_end - x_pos - out[z_pos:z_end, y_pos:y_end, x_pos:x_end] += center[:az, :ay, :ax] - wgt[z_pos:z_end, y_pos:y_end, x_pos:x_end] += 1 - wgt[wgt == 0] = 1 - return out / wgt + out[z_pos:z_end, y_pos:y_end, x_pos:x_end] = segment_loop( + xi, yi, zi, volume, model, tile_size, overlap, logger + ) + return out def _pad_volume(volume, min_pad=16, div=32): @@ -91,11 +99,15 @@ def segment_tomogram( 0 if tile_size[1] == ok else 48, 0 if tile_size[2] == ol else 48, ) - logger.info("Doing instance") + logger.info( + f"Doing instance {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" + ) segmented_volume = _segment_tomogram_instance( volume, model, tile_size, overlap, logger ) - logger.info("Done instance") + logger.info( + f"Done instance {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" + ) (j0, j1), (k0, k1), (l0, l1) = padding segmented_volume = segmented_volume[ j0 : segmented_volume.shape[0] - j1, From b1ab76eb2b061cdad84d3ef1c2b953de9e67201e Mon Sep 17 00:00:00 2001 From: yxd92326 Date: Fri, 6 Mar 2026 12:12:42 +0000 Subject: [PATCH 4/9] Restore back to older state --- .../easymode_segmentation.py | 94 +++++++++++++++---- 1 file changed, 78 insertions(+), 16 deletions(-) diff --git a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py index c6bd75e8..c285a46b 100644 --- a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py +++ b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py @@ -21,26 +21,23 @@ def segment_loop(xi, yi, zi, volume, model, tile_size, overlap, logger): vz1 = min(d, z_start + pz) vy1 = min(h, y_start + py) vx1 = min(w, x_start + px) + + extracted = volume[vz0:vz1, vy0:vy1, vx0:vx1] tile = np.zeros((pz, py, px), dtype=volume.dtype) + tz0 = vz0 - z_start + ty0 = vy0 - y_start + tx0 = vx0 - x_start tile[ - vz0 - z_start : vz1 - z_start, - vy0 - y_start : vy1 - y_start, - vx0 - x_start : vx1 - x_start, - ] = volume[vz0:vz1, vy0:vy1, vx0:vx1] + tz0 : tz0 + extracted.shape[0], + ty0 : ty0 + extracted.shape[1], + tx0 : tx0 + extracted.shape[2], + ] = extracted tile = np.expand_dims(np.array([tile]), axis=-1) segmented_tile = model.predict(tile, verbose=0, batch_size=1).squeeze(-1) logger.info( f"Segmented {xi} {yi} {zi} at {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" ) - z_end = min(zi * sz + sz, d) - y_end = min(yi * sy + sy, h) - x_end = min(xi * sx + sx, w) - az = z_end - zi * sz - ay = y_end - yi * sy - ax = x_end - xi * sx - return ( - segmented_tile[0, oz : oz + sz, oy : oy + sy, ox : ox + sx][:az, :ay, :ax] * 127 - ).astype(np.int8) + return segmented_tile[0, oz : oz + sz, oy : oy + sy, ox : ox + sx] def _segment_tomogram_instance(volume, model, tile_size, overlap, logger): @@ -52,18 +49,62 @@ def _segment_tomogram_instance(volume, model, tile_size, overlap, logger): y_boxes = max(1, (h + sy - 1) // sy) # 9 x_boxes = max(1, (w + sx - 1) // sx) # 13 logger.info(f"{volume.shape}, {z_boxes}, {y_boxes}, {x_boxes}") - out = np.ones((d, h, w), dtype=np.int8) * 0 + out = np.ones((d, h, w), dtype=np.float32) * 0 + wgt = np.ones((d, h, w), dtype=np.int8) * 0 for zi in range(z_boxes): for yi in range(y_boxes): for xi in range(x_boxes): + logger.info(f"Starting {xi} {yi} {zi}") + z_start = zi * sz - oz + y_start = yi * sy - oy + x_start = xi * sx - ox + vz0 = max(0, z_start) + vy0 = max(0, y_start) + vx0 = max(0, x_start) + vz1 = min(d, z_start + pz) + vy1 = min(h, y_start + py) + vx1 = min(w, x_start + px) + extracted = volume[vz0:vz1, vy0:vy1, vx0:vx1] + tile = np.zeros((pz, py, px), dtype=volume.dtype) + tz0 = vz0 - z_start + ty0 = vy0 - y_start + tx0 = vx0 - x_start + tile[ + tz0 : tz0 + extracted.shape[0], + ty0 : ty0 + extracted.shape[1], + tx0 : tx0 + extracted.shape[2], + ] = extracted + logger.info(f"{extracted.shape} {tile.shape}") z_pos, y_pos, x_pos = zi * sz, yi * sy, xi * sx + tile = np.expand_dims(np.array([tile]), axis=-1) + segmented_tile = model.predict(tile, verbose=0, batch_size=1) + logger.info( + f"Segmented at {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" + ) + segmented_tile = segmented_tile.squeeze(-1) + center = segmented_tile[0, oz : oz + sz, oy : oy + sy, ox : ox + sx] + z_end = min(z_pos + sz, d) + y_end = min(y_pos + sy, h) + x_end = min(x_pos + sx, w) + az, ay, ax = z_end - z_pos, y_end - y_pos, x_end - x_pos + out[z_pos:z_end, y_pos:y_end, x_pos:x_end] += center[:az, :ay, :ax] + + """z_pos, y_pos, x_pos = zi * sz, yi * sy, xi * sx z_end = min(z_pos + sz, d) y_end = min(y_pos + sy, h) x_end = min(x_pos + sx, w) - out[z_pos:z_end, y_pos:y_end, x_pos:x_end] = segment_loop( + az, ay, ax = z_end - z_pos, y_end - y_pos, x_end - x_pos + seg_vol_float = segment_loop( xi, yi, zi, volume, model, tile_size, overlap, logger ) - return out + print(np.min(seg_vol_float), np.max(seg_vol_float)) + out[z_pos:z_end, y_pos:y_end, x_pos:x_end] += seg_vol_float[ + :az, :ay, :ax + ] # (seg_vol_float * 127).astype(np.int8)""" + wgt[z_pos:z_end, y_pos:y_end, x_pos:x_end] += 1 + print(np.min(out), np.max(out), np.max(wgt)) + wgt[wgt == 0] = 1 + return out # / wgt def _pad_volume(volume, min_pad=16, div=32): @@ -99,6 +140,7 @@ def segment_tomogram( 0 if tile_size[1] == ok else 48, 0 if tile_size[2] == ol else 48, ) + print(tile_size, overlap) logger.info( f"Doing instance {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" ) @@ -115,5 +157,25 @@ def segment_tomogram( l0 : segmented_volume.shape[2] - l1, ] sj, sk, sl = segmented_volume.shape + print(scale, oj / sj, ok / sk, ol / sl, volume.shape, segmented_volume.shape) segmented_volume = zoom(segmented_volume, (oj / sj, ok / sk, ol / sl), order=1) return segmented_volume, input_apix + + +""" +from cryoemservices.pipeliner_plugins import easymode_segmentation +from easymode.core.distribution import get_model, load_model +model_path, model_metadata = get_model("ribosome") +model = load_model(model_path) +import logging +logger = logging.getLogger() +seg_vol = easymode_segmentation.segment_tomogram(model, "/dls/m06/data/2026/bi41204-8/processed/raw3/relion_murfey/Denoise/job007/tomograms/Position_1_stack_Vol.denoised.mrc", input_apix=6.2, logger=logger) + +import mrcfile +import numpy as np +with mrcfile.new("/dls/m12/data/2026/cm44187-1/tmp/test_vol4.mrc", overwrite=True) as mrc: + #mrc.set_data((seg_vol[0] * 127).astype(np.int8)) + mrc.set_data(seg_vol_2[0]) + mrc.voxel_size = 6.2 + +""" From a2e6368bd13e4e39d282d96cb72c09544909ff2f Mon Sep 17 00:00:00 2001 From: yxd92326 Date: Fri, 6 Mar 2026 14:39:33 +0000 Subject: [PATCH 5/9] Known working state --- .../easymode_segmentation.py | 103 ++++++------------ src/cryoemservices/services/easymode.py | 1 - 2 files changed, 31 insertions(+), 73 deletions(-) diff --git a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py index c285a46b..e60b883a 100644 --- a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py +++ b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py @@ -1,12 +1,9 @@ -import os - import mrcfile import numpy as np -import psutil from scipy.ndimage import zoom -def segment_loop(xi, yi, zi, volume, model, tile_size, overlap, logger): +def segment_loop(xi, yi, zi, volume, model, tile_size, overlap): (pz, py, px) = tile_size (oz, oy, ox) = overlap d, h, w = volume.shape # 672, 1440, 2016 @@ -34,13 +31,13 @@ def segment_loop(xi, yi, zi, volume, model, tile_size, overlap, logger): ] = extracted tile = np.expand_dims(np.array([tile]), axis=-1) segmented_tile = model.predict(tile, verbose=0, batch_size=1).squeeze(-1) - logger.info( - f"Segmented {xi} {yi} {zi} at {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" - ) + # logger.info( + # f"Segmented {xi} {yi} {zi} at {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" + # ) return segmented_tile[0, oz : oz + sz, oy : oy + sy, ox : ox + sx] -def _segment_tomogram_instance(volume, model, tile_size, overlap, logger): +def _segment_tomogram_instance(volume, model, tile_size, overlap): (pz, py, px) = tile_size (oz, oy, ox) = overlap d, h, w = volume.shape # 672, 1440, 2016 @@ -48,63 +45,24 @@ def _segment_tomogram_instance(volume, model, tile_size, overlap, logger): z_boxes = max(1, (d + sz - 1) // sz) # 5 y_boxes = max(1, (h + sy - 1) // sy) # 9 x_boxes = max(1, (w + sx - 1) // sx) # 13 - logger.info(f"{volume.shape}, {z_boxes}, {y_boxes}, {x_boxes}") - out = np.ones((d, h, w), dtype=np.float32) * 0 - wgt = np.ones((d, h, w), dtype=np.int8) * 0 + # logger.info(f"{volume.shape}, {z_boxes}, {y_boxes}, {x_boxes}") + out = np.ones((d, h, w), dtype=np.int8) * 0 for zi in range(z_boxes): for yi in range(y_boxes): for xi in range(x_boxes): - logger.info(f"Starting {xi} {yi} {zi}") - z_start = zi * sz - oz - y_start = yi * sy - oy - x_start = xi * sx - ox - vz0 = max(0, z_start) - vy0 = max(0, y_start) - vx0 = max(0, x_start) - vz1 = min(d, z_start + pz) - vy1 = min(h, y_start + py) - vx1 = min(w, x_start + px) - extracted = volume[vz0:vz1, vy0:vy1, vx0:vx1] - tile = np.zeros((pz, py, px), dtype=volume.dtype) - tz0 = vz0 - z_start - ty0 = vy0 - y_start - tx0 = vx0 - x_start - tile[ - tz0 : tz0 + extracted.shape[0], - ty0 : ty0 + extracted.shape[1], - tx0 : tx0 + extracted.shape[2], - ] = extracted - logger.info(f"{extracted.shape} {tile.shape}") + # logger.info(f"Starting {xi} {yi} {zi}") z_pos, y_pos, x_pos = zi * sz, yi * sy, xi * sx - tile = np.expand_dims(np.array([tile]), axis=-1) - segmented_tile = model.predict(tile, verbose=0, batch_size=1) - logger.info( - f"Segmented at {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" - ) - segmented_tile = segmented_tile.squeeze(-1) - center = segmented_tile[0, oz : oz + sz, oy : oy + sy, ox : ox + sx] - z_end = min(z_pos + sz, d) - y_end = min(y_pos + sy, h) - x_end = min(x_pos + sx, w) - az, ay, ax = z_end - z_pos, y_end - y_pos, x_end - x_pos - out[z_pos:z_end, y_pos:y_end, x_pos:x_end] += center[:az, :ay, :ax] - - """z_pos, y_pos, x_pos = zi * sz, yi * sy, xi * sx z_end = min(z_pos + sz, d) y_end = min(y_pos + sy, h) x_end = min(x_pos + sx, w) az, ay, ax = z_end - z_pos, y_end - y_pos, x_end - x_pos seg_vol_float = segment_loop( - xi, yi, zi, volume, model, tile_size, overlap, logger - ) - print(np.min(seg_vol_float), np.max(seg_vol_float)) - out[z_pos:z_end, y_pos:y_end, x_pos:x_end] += seg_vol_float[ - :az, :ay, :ax - ] # (seg_vol_float * 127).astype(np.int8)""" - wgt[z_pos:z_end, y_pos:y_end, x_pos:x_end] += 1 - print(np.min(out), np.max(out), np.max(wgt)) - wgt[wgt == 0] = 1 - return out # / wgt + xi, yi, zi, volume, model, tile_size, overlap + )[:az, :ay, :ax] + out[z_pos:z_end, y_pos:y_end, x_pos:x_end] += ( + seg_vol_float * 127 + ).astype(np.int8) + return out def _pad_volume(volume, min_pad=16, div=32): @@ -120,13 +78,18 @@ def _pad_volume(volume, min_pad=16, div=32): def segment_tomogram( - model, tomogram_path, tta=1, batch_size=1, model_apix=10, input_apix=10, logger=None + model, + tomogram_path, + tta=1, + batch_size=1, + model_apix=10, + input_apix=10, # , logger=None ): - logger.info("Reading tomogram") + # logger.info("Reading tomogram") with mrcfile.open(tomogram_path) as m: volume = m.data.astype(np.float32) oj, ok, ol = volume.shape - logger.info("Read tomogram") + # logger.info("Read tomogram") scale = float(input_apix) / float(model_apix) volume = zoom(volume, scale, order=1) _k_margin = min(int(0.2 * ok), 64) @@ -140,16 +103,13 @@ def segment_tomogram( 0 if tile_size[1] == ok else 48, 0 if tile_size[2] == ol else 48, ) - print(tile_size, overlap) - logger.info( - f"Doing instance {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" - ) - segmented_volume = _segment_tomogram_instance( - volume, model, tile_size, overlap, logger - ) - logger.info( - f"Done instance {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" - ) + # logger.info( + # f"Doing instance {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" + # ) + segmented_volume = _segment_tomogram_instance(volume, model, tile_size, overlap) + # logger.info( + # f"Done instance {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" + # ) (j0, j1), (k0, k1), (l0, l1) = padding segmented_volume = segmented_volume[ j0 : segmented_volume.shape[0] - j1, @@ -157,7 +117,6 @@ def segment_tomogram( l0 : segmented_volume.shape[2] - l1, ] sj, sk, sl = segmented_volume.shape - print(scale, oj / sj, ok / sk, ol / sl, volume.shape, segmented_volume.shape) segmented_volume = zoom(segmented_volume, (oj / sj, ok / sk, ol / sl), order=1) return segmented_volume, input_apix @@ -173,9 +132,9 @@ def segment_tomogram( import mrcfile import numpy as np -with mrcfile.new("/dls/m12/data/2026/cm44187-1/tmp/test_vol4.mrc", overwrite=True) as mrc: +with mrcfile.new("/dls/m12/data/2026/cm44187-1/tmp/test_vol6.mrc", overwrite=True) as mrc: #mrc.set_data((seg_vol[0] * 127).astype(np.int8)) - mrc.set_data(seg_vol_2[0]) + mrc.set_data(seg_vol[0]) mrc.voxel_size = 6.2 """ diff --git a/src/cryoemservices/services/easymode.py b/src/cryoemservices/services/easymode.py index 07cd9e08..2f088b9a 100644 --- a/src/cryoemservices/services/easymode.py +++ b/src/cryoemservices/services/easymode.py @@ -120,7 +120,6 @@ def easymode(self, rw, header: dict, message: dict): batch_size=easymode_params.batch_size, model_apix=model_metadata["apix"], input_apix=easymode_params.pixel_size, - logger=self.log, ) # Convert to int8 and save mrc From f5962fd363506da6eb513aa08af797c21b97c515 Mon Sep 17 00:00:00 2001 From: yxd92326 Date: Fri, 6 Mar 2026 14:53:35 +0000 Subject: [PATCH 6/9] Bit of refactoring --- .../easymode_segmentation.py | 52 ++++++------------- 1 file changed, 17 insertions(+), 35 deletions(-) diff --git a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py index e60b883a..29b456c4 100644 --- a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py +++ b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py @@ -3,11 +3,10 @@ from scipy.ndimage import zoom -def segment_loop(xi, yi, zi, volume, model, tile_size, overlap): +def segment_tile(xi, yi, zi, volume, model, tile_size, overlap): (pz, py, px) = tile_size (oz, oy, ox) = overlap - d, h, w = volume.shape # 672, 1440, 2016 - sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox # non-overlapped patch size, 160 + sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox # non-overlapped patch size z_start = zi * sz - oz y_start = yi * sy - oy @@ -15,48 +14,39 @@ def segment_loop(xi, yi, zi, volume, model, tile_size, overlap): vz0 = max(0, z_start) vy0 = max(0, y_start) vx0 = max(0, x_start) - vz1 = min(d, z_start + pz) - vy1 = min(h, y_start + py) - vx1 = min(w, x_start + px) + vz1 = min(volume.shape[0], z_start + pz) + vy1 = min(volume.shape[1], y_start + py) + vx1 = min(volume.shape[2], x_start + px) - extracted = volume[vz0:vz1, vy0:vy1, vx0:vx1] tile = np.zeros((pz, py, px), dtype=volume.dtype) - tz0 = vz0 - z_start - ty0 = vy0 - y_start - tx0 = vx0 - x_start tile[ - tz0 : tz0 + extracted.shape[0], - ty0 : ty0 + extracted.shape[1], - tx0 : tx0 + extracted.shape[2], - ] = extracted + vz0 - z_start : vz1 - z_start, + vy0 - y_start : vy1 - y_start, + vx0 - x_start : vx1 - x_start, + ] = volume[vz0:vz1, vy0:vy1, vx0:vx1] tile = np.expand_dims(np.array([tile]), axis=-1) segmented_tile = model.predict(tile, verbose=0, batch_size=1).squeeze(-1) - # logger.info( - # f"Segmented {xi} {yi} {zi} at {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" - # ) return segmented_tile[0, oz : oz + sz, oy : oy + sy, ox : ox + sx] def _segment_tomogram_instance(volume, model, tile_size, overlap): (pz, py, px) = tile_size (oz, oy, ox) = overlap - d, h, w = volume.shape # 672, 1440, 2016 - sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox # non-overlapped patch size, 160 - z_boxes = max(1, (d + sz - 1) // sz) # 5 - y_boxes = max(1, (h + sy - 1) // sy) # 9 - x_boxes = max(1, (w + sx - 1) // sx) # 13 - # logger.info(f"{volume.shape}, {z_boxes}, {y_boxes}, {x_boxes}") - out = np.ones((d, h, w), dtype=np.int8) * 0 + d, h, w = volume.shape + sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox # non-overlapped patch size + z_boxes = max(1, (d + sz - 1) // sz) + y_boxes = max(1, (h + sy - 1) // sy) + x_boxes = max(1, (w + sx - 1) // sx) + out = np.ones(volume.shape, dtype=np.int8) * 0 for zi in range(z_boxes): for yi in range(y_boxes): for xi in range(x_boxes): - # logger.info(f"Starting {xi} {yi} {zi}") z_pos, y_pos, x_pos = zi * sz, yi * sy, xi * sx z_end = min(z_pos + sz, d) y_end = min(y_pos + sy, h) x_end = min(x_pos + sx, w) az, ay, ax = z_end - z_pos, y_end - y_pos, x_end - x_pos - seg_vol_float = segment_loop( + seg_vol_float = segment_tile( xi, yi, zi, volume, model, tile_size, overlap )[:az, :ay, :ax] out[z_pos:z_end, y_pos:y_end, x_pos:x_end] += ( @@ -83,13 +73,11 @@ def segment_tomogram( tta=1, batch_size=1, model_apix=10, - input_apix=10, # , logger=None + input_apix=10, ): - # logger.info("Reading tomogram") with mrcfile.open(tomogram_path) as m: volume = m.data.astype(np.float32) oj, ok, ol = volume.shape - # logger.info("Read tomogram") scale = float(input_apix) / float(model_apix) volume = zoom(volume, scale, order=1) _k_margin = min(int(0.2 * ok), 64) @@ -103,13 +91,7 @@ def segment_tomogram( 0 if tile_size[1] == ok else 48, 0 if tile_size[2] == ol else 48, ) - # logger.info( - # f"Doing instance {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" - # ) segmented_volume = _segment_tomogram_instance(volume, model, tile_size, overlap) - # logger.info( - # f"Done instance {psutil.Process(os.getpid()).memory_info().rss / 1024**2}" - # ) (j0, j1), (k0, k1), (l0, l1) = padding segmented_volume = segmented_volume[ j0 : segmented_volume.shape[0] - j1, From c95076955b431e20bd73f84c5936d9a5391ace54 Mon Sep 17 00:00:00 2001 From: yxd92326 Date: Fri, 6 Mar 2026 16:44:44 +0000 Subject: [PATCH 7/9] Remove example --- .../easymode_segmentation.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py index 29b456c4..9232f7cc 100644 --- a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py +++ b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py @@ -101,22 +101,3 @@ def segment_tomogram( sj, sk, sl = segmented_volume.shape segmented_volume = zoom(segmented_volume, (oj / sj, ok / sk, ol / sl), order=1) return segmented_volume, input_apix - - -""" -from cryoemservices.pipeliner_plugins import easymode_segmentation -from easymode.core.distribution import get_model, load_model -model_path, model_metadata = get_model("ribosome") -model = load_model(model_path) -import logging -logger = logging.getLogger() -seg_vol = easymode_segmentation.segment_tomogram(model, "/dls/m06/data/2026/bi41204-8/processed/raw3/relion_murfey/Denoise/job007/tomograms/Position_1_stack_Vol.denoised.mrc", input_apix=6.2, logger=logger) - -import mrcfile -import numpy as np -with mrcfile.new("/dls/m12/data/2026/cm44187-1/tmp/test_vol6.mrc", overwrite=True) as mrc: - #mrc.set_data((seg_vol[0] * 127).astype(np.int8)) - mrc.set_data(seg_vol[0]) - mrc.voxel_size = 6.2 - -""" From afe7adf392f92f2ab08616dd28a897272de597b6 Mon Sep 17 00:00:00 2001 From: yxd92326 Date: Wed, 11 Mar 2026 10:16:42 +0000 Subject: [PATCH 8/9] Treat tile size as minimum vol --- .../easymode_segmentation.py | 40 +++++++++++-------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py index 9232f7cc..67485dd0 100644 --- a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py +++ b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py @@ -4,9 +4,11 @@ def segment_tile(xi, yi, zi, volume, model, tile_size, overlap): - (pz, py, px) = tile_size (oz, oy, ox) = overlap - sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox # non-overlapped patch size + # non-overlapped patch sizes + sz = tile_size - 2 * oz + sy = tile_size - 2 * oy + sx = tile_size - 2 * ox z_start = zi * sz - oz y_start = yi * sy - oy @@ -14,11 +16,11 @@ def segment_tile(xi, yi, zi, volume, model, tile_size, overlap): vz0 = max(0, z_start) vy0 = max(0, y_start) vx0 = max(0, x_start) - vz1 = min(volume.shape[0], z_start + pz) - vy1 = min(volume.shape[1], y_start + py) - vx1 = min(volume.shape[2], x_start + px) + vz1 = min(volume.shape[0], z_start + tile_size) + vy1 = min(volume.shape[1], y_start + tile_size) + vx1 = min(volume.shape[2], x_start + tile_size) - tile = np.zeros((pz, py, px), dtype=volume.dtype) + tile = np.zeros((tile_size, tile_size, tile_size), dtype=volume.dtype) tile[ vz0 - z_start : vz1 - z_start, vy0 - y_start : vy1 - y_start, @@ -30,10 +32,13 @@ def segment_tile(xi, yi, zi, volume, model, tile_size, overlap): def _segment_tomogram_instance(volume, model, tile_size, overlap): - (pz, py, px) = tile_size (oz, oy, ox) = overlap d, h, w = volume.shape - sz, sy, sx = pz - 2 * oz, py - 2 * oy, px - 2 * ox # non-overlapped patch size + # non-overlapped patch sizes + sz = tile_size - 2 * oz + sy = tile_size - 2 * oy + sx = tile_size - 2 * ox + z_boxes = max(1, (d + sz - 1) // sz) y_boxes = max(1, (h + sy - 1) // sy) x_boxes = max(1, (w + sx - 1) // sx) @@ -49,17 +54,19 @@ def _segment_tomogram_instance(volume, model, tile_size, overlap): seg_vol_float = segment_tile( xi, yi, zi, volume, model, tile_size, overlap )[:az, :ay, :ax] - out[z_pos:z_end, y_pos:y_end, x_pos:x_end] += ( + out[z_pos:z_end, y_pos:y_end, x_pos:x_end] = ( seg_vol_float * 127 ).astype(np.int8) return out -def _pad_volume(volume, min_pad=16, div=32): +def _pad_volume(volume, min_pad=16, div=32, min_size=None): j, k, l = volume.shape pads = [] for n in (j, k, l): total_pad = max(2 * min_pad, ((n + 2 * min_pad + div - 1) // div) * div - n) + if min_size and total_pad + n < min_size: + total_pad = min_size - n before = total_pad // 2 after = total_pad - before pads.append((before, after)) @@ -74,6 +81,7 @@ def segment_tomogram( batch_size=1, model_apix=10, input_apix=10, + tile_size=256, ): with mrcfile.open(tomogram_path) as m: volume = m.data.astype(np.float32) @@ -84,12 +92,11 @@ def segment_tomogram( _l_margin = min(int(0.2 * ol), 64) volume -= np.mean(volume[:, _k_margin:-_k_margin, _l_margin:-_l_margin]) volume /= np.std(volume[:, _k_margin:-_k_margin, _l_margin:-_l_margin]) + 1e-7 - volume, padding = _pad_volume(volume) - tile_size = (min(256, oj), min(256, ok), min(256, ol)) + volume, padding = _pad_volume(volume, min_size=tile_size) overlap = ( - 0 if tile_size[0] == oj else 48, - 0 if tile_size[1] == ok else 48, - 0 if tile_size[2] == ol else 48, + 0 if tile_size >= oj else 48, + 0 if tile_size >= ok else 48, + 0 if tile_size >= ol else 48, ) segmented_volume = _segment_tomogram_instance(volume, model, tile_size, overlap) (j0, j1), (k0, k1), (l0, l1) = padding @@ -98,6 +105,5 @@ def segment_tomogram( k0 : segmented_volume.shape[1] - k1, l0 : segmented_volume.shape[2] - l1, ] - sj, sk, sl = segmented_volume.shape - segmented_volume = zoom(segmented_volume, (oj / sj, ok / sk, ol / sl), order=1) + segmented_volume = zoom(segmented_volume, 1 / scale, order=1) return segmented_volume, input_apix From 74b16a0e69a5482ae15200ee4a1e4a460ccf4bf8 Mon Sep 17 00:00:00 2001 From: yxd92326 Date: Wed, 15 Apr 2026 09:56:19 +0100 Subject: [PATCH 9/9] Rewrite very slightly --- src/cryoemservices/pipeliner_plugins/easymode_segmentation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py index 67485dd0..8bed7ad4 100644 --- a/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py +++ b/src/cryoemservices/pipeliner_plugins/easymode_segmentation.py @@ -105,5 +105,6 @@ def segment_tomogram( k0 : segmented_volume.shape[1] - k1, l0 : segmented_volume.shape[2] - l1, ] - segmented_volume = zoom(segmented_volume, 1 / scale, order=1) + sj, sk, sl = segmented_volume.shape + segmented_volume = zoom(segmented_volume, (oj / sj, ok / sk, ol / sl), order=1) return segmented_volume, input_apix