Plot the true particles sizes for tomography (#121)

src/cryoemservices/services/cryolo.py

@@ -286,14 +286,12 @@ def cryolo(self, rw, header: dict, message: dict):
                 "image_command": "picked_particles_3d_apng",
                 "file": cryolo_params.input_path,
                 "coordinates_file": cryolo_params.output_path,
-                "diameter_pixels": cryolo_params.cryolo_box_size,
                 "box_size": cryolo_params.cryolo_box_size,
             }
             central_slice_parameters = {
                 "image_command": "picked_particles_3d_central_slice",
                 "file": cryolo_params.input_path,
                 "coordinates_file": cryolo_params.output_path,
-                "diameter_pixels": cryolo_params.cryolo_box_size,
                 "box_size": cryolo_params.cryolo_box_size,
             }
             rw.send_to("images", movie_parameters)

src/cryoemservices/services/images.py

@@ -1,27 +1,19 @@
 from __future__ import annotations
 
 from importlib.metadata import entry_points
-from typing import Any, Callable, NamedTuple
+from typing import Callable
 
 import workflows.recipe
 from workflows.services.common_service import CommonService
 
 
-class PluginInterface(NamedTuple):
-    rw: workflows.recipe.wrapper.RecipeWrapper
-    parameters: Callable
-    message: dict[str, Any]
-
-
 class Images(CommonService):
     """
     A service that generates images and thumbnails.
     Plugin functions can be registered under the entry point
     'cryoemservices.services.images.plugins'. The contract is that a plugin function
-    takes a single argument of type PluginInterface, and returns a truthy value
+    takes a parameters callable, and returns a truthy value
     to acknowledge success, and a falsy value to reject the related message.
-    Functions may choose to return a list of files that were generated, but
-    this is optional at this time.
     """
 
     # Human readable service name
@@ -69,9 +61,7 @@ def parameters(key: str):
             return
 
         try:
-            result = self.image_functions[command](
-                PluginInterface(rw, parameters, message)
-            )
+            result = self.image_functions[command](parameters)
         except (PermissionError, FileNotFoundError) as e:
             self.log.error(f"Command {command!r} raised {e}", exc_info=True)
             rw.transport.nack(header)

src/cryoemservices/services/images_plugins.py

@@ -3,6 +3,7 @@
 import logging
 import time
 from pathlib import Path
+from typing import Callable
 
 import mrcfile
 import numpy as np
@@ -15,13 +16,20 @@
 logger.setLevel(logging.INFO)
 
 
-def mrc_to_jpeg(plugin_params):
-    filename = plugin_params.parameters("file")
-    allframes = plugin_params.parameters("all_frames")
-    if not filename or filename == "None":
-        logger.error("Skipping mrc to jpeg conversion: filename not specified")
+def required_parameters(parameters: Callable, required_keys: list[str]):
+    """Make sure all parameters which aren't nullable are there"""
+    for param_key in required_keys:
+        if parameters(param_key) in [False, "False", None, "None"]:
+            logger.error(f"Required key {param_key} not valid: {parameters(param_key)}")
+            return False
+    return True
+
+
+def mrc_to_jpeg(plugin_params: Callable):
+    if not required_parameters(plugin_params, ["file"]):
         return False
-    filepath = Path(filename)
+    filepath = Path(plugin_params("file"))
+    allframes = plugin_params("all_frames")
     if not filepath.is_file():
         logger.error(f"File {filepath} not found")
         return False
@@ -86,31 +94,30 @@ def mrc_to_jpeg(plugin_params):
     timing = time.perf_counter() - start
 
     logger.info(
-        f"Converted mrc to jpeg {filename} -> {outfile} in {timing:.1f} seconds",
+        f"Converted mrc to jpeg {filepath} -> {outfile} in {timing:.1f} seconds",
         extra={"image-processing-time": timing},
     )
     if outfiles:
         return outfiles
     return outfile
 
 
-def picked_particles(plugin_params):
-    basefilename = Path(plugin_params.parameters("file"))
+def picked_particles(plugin_params: Callable):
+    if not required_parameters(plugin_params, ["file", "diameter", "outfile"]):
+        return False
+    basefilename = Path(plugin_params("file"))
     if basefilename.suffix == ".jpeg":
         logger.info(f"Replacing jpeg extension with mrc extension for {basefilename}")
         basefilename = basefilename.with_suffix(".mrc")
-    coords = plugin_params.parameters("coordinates")
-    selected_coords = plugin_params.parameters("selected_coordinates")
-    pixel_size = plugin_params.parameters("pixel_size")
+    coords = plugin_params("coordinates")
+    selected_coords = plugin_params("selected_coordinates")
+    pixel_size = plugin_params("pixel_size")
     if not pixel_size:
         # Legacy case of zocalo-relion
-        pixel_size = plugin_params.parameters("angpix")
-    diam = plugin_params.parameters("diameter")
-    contrast_factor = plugin_params.parameters("contrast_factor") or 6
-    outfile = plugin_params.parameters("outfile")
-    if not outfile:
-        logger.error(f"Outfile incorrectly specified: {outfile}")
-        return False
+        pixel_size = plugin_params("angpix") or 1
+    diam = plugin_params("diameter")
+    contrast_factor = plugin_params("contrast_factor") or 6
+    outfile = plugin_params("outfile")
     if not basefilename.is_file():
         logger.error(f"File {basefilename} not found")
         return False
@@ -124,9 +131,6 @@ def picked_particles(plugin_params):
             f"File {basefilename} could not be opened. It may be corrupted or not in mrc format"
         )
         return False
-    except FileNotFoundError:
-        logger.error(f"File {basefilename} could not be opened")
-        return False
     mean = np.mean(data)
     sdev = np.std(data)
     sigma_min = mean - 3 * sdev
@@ -181,14 +185,11 @@ def picked_particles(plugin_params):
     return outfile
 
 
-def mrc_central_slice(plugin_params):
-    filename = plugin_params.parameters("file")
-    skip_rescaling = plugin_params.parameters("skip_rescaling")
-
-    if not filename or filename == "None":
-        logger.error("Skipping mrc to jpeg conversion: filename not specified")
+def mrc_central_slice(plugin_params: Callable):
+    if not required_parameters(plugin_params, ["file"]):
         return False
-    filepath = Path(filename)
+    filepath = Path(plugin_params("file"))
+    skip_rescaling = plugin_params("skip_rescaling")
     if not filepath.is_file():
         logger.error(f"File {filepath} not found")
         return False
@@ -235,21 +236,17 @@ def mrc_central_slice(plugin_params):
     timing = time.perf_counter() - start
 
     logger.info(
-        f"Converted mrc to jpeg {filename} -> {outfile} in {timing:.1f} seconds",
+        f"Converted mrc to jpeg {filepath} -> {outfile} in {timing:.1f} seconds",
         extra={"image-processing-time": timing},
     )
     return outfile
 
 
-def mrc_to_apng(plugin_params):
-    filename = plugin_params.parameters("file")
-    skip_rescaling = plugin_params.parameters("skip_rescaling")
-
-    if not filename or filename == "None":
-        logger.error("Skipping mrc to jpeg conversion: filename not specified")
+def mrc_to_apng(plugin_params: Callable):
+    if not required_parameters(plugin_params, ["file"]):
         return False
-    filepath = Path(filename)
-
+    filepath = Path(plugin_params("file"))
+    skip_rescaling = plugin_params("skip_rescaling")
     if not filepath.is_file():
         logger.error(f"File {filepath} not found")
         return False
@@ -294,19 +291,16 @@ def mrc_to_apng(plugin_params):
         return False
     timing = time.perf_counter() - start
     logger.info(
-        f"Converted mrc to apng {filename} -> {outfile} in {timing:.1f} seconds"
+        f"Converted mrc to apng {filepath} -> {outfile} in {timing:.1f} seconds"
     )
     return outfile
 
 
 def particles_3d_in_frame(
     framedata: np.ndarray,
     framenum: int,
-    radius: int,
     radius_box: int,
-    pick_x_coords: np.ndarray,
-    pick_y_coords: np.ndarray,
-    pick_z_coords: np.ndarray,
+    coords_file: Path,
 ):
     """Function to plot the pick locations of particles in a frame of a volume"""
     # Rescale the frame image
@@ -325,27 +319,45 @@ def particles_3d_in_frame(
     colour_im = im.convert("RGB")
     dim = ImageDraw.Draw(colour_im)
 
+    # Find the coordinates of the picks
+    try:
+        all_coords: pd.DataFrame = starfile.read(coords_file)["cryolo"]
+        pick_x_coords = all_coords["CoordinateX"].to_numpy(dtype=float)
+        pick_y_coords = all_coords["CoordinateY"].to_numpy(dtype=float)
+        pick_z_coords = all_coords["CoordinateZ"].to_numpy(dtype=float)
+        pick_width = all_coords["EstWidth"].to_numpy(dtype=float)
+        pick_height = all_coords["EstHeight"].to_numpy(dtype=float)
+        frame_count = all_coords["NumBoxes"].to_numpy(dtype=int)
+    except KeyError:
+        logger.warning(f"Cannot find picks in {coords_file}")
+        colour_im.thumbnail((512, 512))
+        return colour_im
+
     # Find the locations of picks which cross into this frame
-    picks_appearing_in_frame = np.abs(pick_z_coords - framenum) < radius
+    picks_appearing_in_frame = np.abs(pick_z_coords - framenum) - frame_count // 2 <= 0
     frame_x_coords = pick_x_coords[picks_appearing_in_frame]
     frame_y_coords = pick_y_coords[picks_appearing_in_frame]
     if len(frame_x_coords) and len(frame_y_coords):
         # Red circles for all coordinates
         for cid in range(len(frame_x_coords)):
             # Scale the pick radius based on the distance of frame from centre
-            cid_radius = np.sqrt(
-                radius**2
+            scaled_width = np.sqrt(
+                pick_width[cid] ** 2
+                - (pick_z_coords[picks_appearing_in_frame][cid] - framenum) ** 2
+            )
+            scaled_height = np.sqrt(
+                pick_height[cid] ** 2
                 - (pick_z_coords[picks_appearing_in_frame][cid] - framenum) ** 2
             )
             dim.ellipse(
                 [
                     (
-                        frame_x_coords[cid] + radius_box - cid_radius,
-                        frame_y_coords[cid] + radius_box - cid_radius,
+                        frame_x_coords[cid] + radius_box - scaled_width / 2,
+                        frame_y_coords[cid] + radius_box - scaled_height / 2,
                     ),
                     (
-                        frame_x_coords[cid] + radius_box + cid_radius,
-                        frame_y_coords[cid] + radius_box + cid_radius,
+                        frame_x_coords[cid] + radius_box + scaled_width / 2,
+                        frame_y_coords[cid] + radius_box + scaled_height / 2,
                     ),
                 ],
                 width=4,
@@ -356,19 +368,16 @@ def particles_3d_in_frame(
     return colour_im
 
 
-def picked_particles_3d_central_slice(plugin_params):
-    filename = plugin_params.parameters("file")
-    coords_file = plugin_params.parameters("coordinates_file")
-    diameter_pixels = plugin_params.parameters("diameter_pixels")
-    box_size = plugin_params.parameters("box_size")
+def picked_particles_3d_central_slice(plugin_params: Callable):
+    if not required_parameters(plugin_params, ["file", "coordinates_file", "box_size"]):
+        return False
+    filename = plugin_params("file")
+    coords_file = plugin_params("coordinates_file")
+    box_size = plugin_params("box_size")
+    radius_box = box_size // 2
     if not Path(filename).is_file() or not Path(coords_file).is_file():
         logger.error(f"File {filename} or {coords_file} not found")
         return False
-    radius = diameter_pixels // 2
-    radius_box = box_size // 2
-
-    outfile = str(Path(coords_file).with_suffix("")) + "_thumbnail.jpeg"
-
     try:
         with mrcfile.open(filename) as mrc:
             data = mrc.data
@@ -378,16 +387,6 @@ def picked_particles_3d_central_slice(plugin_params):
     if not len(data.shape) == 3:
         logger.error(f"File {filename} is not a 3D volume")
         return False
-    try:
-        all_coords: pd.DataFrame = starfile.read(coords_file)["cryolo"]
-        pick_x_coords = all_coords["CoordinateX"].to_numpy(dtype=float)
-        pick_y_coords = all_coords["CoordinateY"].to_numpy(dtype=float)
-        pick_z_coords = all_coords["CoordinateZ"].to_numpy(dtype=float)
-    except KeyError:
-        logger.warning(f"Cannot find picks in {coords_file}")
-        pick_x_coords = np.array([])
-        pick_y_coords = np.array([])
-        pick_z_coords = np.array([])
 
     # Extract central slice
     total_slices = data.shape[0]
@@ -397,12 +396,11 @@ def picked_particles_3d_central_slice(plugin_params):
     colour_im = particles_3d_in_frame(
         framedata=central_slice_data,
         framenum=central_slice_index,
-        radius=radius,
         radius_box=radius_box,
-        pick_x_coords=pick_x_coords,
-        pick_y_coords=pick_y_coords,
-        pick_z_coords=pick_z_coords,
+        coords_file=Path(coords_file),
     )
+
+    outfile = str(Path(coords_file).with_suffix("")) + "_thumbnail.jpeg"
     try:
         colour_im.save(outfile)
     except FileNotFoundError:
@@ -413,19 +411,16 @@ def picked_particles_3d_central_slice(plugin_params):
     return outfile
 
 
-def picked_particles_3d_apng(plugin_params):
-    filename = plugin_params.parameters("file")
-    coords_file = plugin_params.parameters("coordinates_file")
-    diameter_pixels = plugin_params.parameters("diameter_pixels")
-    box_size = plugin_params.parameters("box_size")
+def picked_particles_3d_apng(plugin_params: Callable):
+    if not required_parameters(plugin_params, ["file", "coordinates_file", "box_size"]):
+        return False
+    filename = plugin_params("file")
+    coords_file = plugin_params("coordinates_file")
+    box_size = plugin_params("box_size")
+    radius_box = box_size // 2
     if not Path(filename).is_file() or not Path(coords_file).is_file():
         logger.error(f"File {filename} or {coords_file} not found")
         return False
-    radius = diameter_pixels // 2
-    radius_box = box_size // 2
-
-    outfile = str(Path(coords_file).with_suffix("")) + "_movie.png"
-
     try:
         with mrcfile.open(filename) as mrc:
             data = mrc.data
@@ -435,29 +430,18 @@ def picked_particles_3d_apng(plugin_params):
     if not len(data.shape) == 3:
         logger.error(f"File {filename} is not a 3D volume")
         return False
-    try:
-        all_coords: pd.DataFrame = starfile.read(coords_file)["cryolo"]
-        pick_x_coords = all_coords["CoordinateX"].to_numpy(dtype=float)
-        pick_y_coords = all_coords["CoordinateY"].to_numpy(dtype=float)
-        pick_z_coords = all_coords["CoordinateZ"].to_numpy(dtype=float)
-    except KeyError:
-        logger.warning(f"Cannot find picks in {coords_file}")
-        pick_x_coords = np.array([])
-        pick_y_coords = np.array([])
-        pick_z_coords = np.array([])
 
     images_to_append = []
     for framenum, frame in enumerate(data):
         colour_im = particles_3d_in_frame(
             framedata=frame,
             framenum=framenum,
-            radius=radius,
             radius_box=radius_box,
-            pick_x_coords=pick_x_coords,
-            pick_y_coords=pick_y_coords,
-            pick_z_coords=pick_z_coords,
+            coords_file=Path(coords_file),
         )
         images_to_append.append(colour_im)
+
+    outfile = str(Path(coords_file).with_suffix("")) + "_movie.png"
     try:
         im_frame0 = images_to_append[0]
         im_frame0.save(outfile, save_all=True, append_images=images_to_append[1:])