Allow user provided crYOLO model (#311)

Search a configured directory for a crYOLO model before falling back on a configured default

Author: d-j-hatton

src/murfey/server/api.py

@@ -49,6 +49,7 @@
 from murfey.server.config import from_file, settings
 from murfey.server.gain import Camera, prepare_eer_gain, prepare_gain
 from murfey.server.murfey_db import murfey_db
+from murfey.server.spa.api import _cryolo_model_path
 from murfey.util.db import (
     AutoProcProgram,
     ClientEnvironment,
@@ -997,8 +998,11 @@ async def request_spa_preprocessing(
             Path(secure_filename(str(mrc_out))).parent.mkdir(
                 parents=True, exist_ok=True
             )
+        recipe_name = machine_config.recipes.get(
+            "em-spa-preprocess", "em-spa-preprocess"
+        )
         zocalo_message = {
-            "recipes": ["em-spa-preprocess"],
+            "recipes": [recipe_name],
             "parameters": {
                 "feedback_queue": machine_config.feedback_queue,
                 "node_creator_queue": machine_config.node_creator_queue,
@@ -1019,6 +1023,7 @@ async def request_spa_preprocessing(
                 "particle_diameter": proc_params["particle_diameter"] or 0,
                 "fm_int_file": proc_file.eer_fractionation_file,
                 "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs,
+                "cryolo_model_weights": str(_cryolo_model_path(visit_name)),
             },
         }
         # log.info(f"Sending Zocalo message {zocalo_message}")

src/murfey/server/config.py

@@ -17,6 +17,7 @@ class MachineConfig(BaseModel):
     rsync_basepath: Path
     murfey_db_credentials: str
     crypto_key: str
+    default_model: Path
     display_name: str = ""
     image_path: Optional[Path] = None
     software_versions: Dict[str, str] = {}
@@ -41,11 +42,23 @@ class MachineConfig(BaseModel):
     processed_extra_directory: str = ""
     plugin_packages: Dict[str, Path] = {}
     software_settings_output_directories: Dict[str, List[str]] = {}
+    recipes: Dict[str, str] = {
+        "em-spa-bfactor": "em-spa-bfactor",
+        "em-spa-class2d": "em-spa-class2d",
+        "em-spa-class3d": "em-spa-class3d",
+        "em-spa-preprocess": "em-spa-preprocess",
+        "em-spa-refine": "em-spa-refine",
+        "em-tomo-preprocess": "em-tomo-preprocess",
+        "em-tomo-align": "em-tomo-align",
+    }
 
     # Find and download upstream directories
     upstream_data_directories: List[Path] = []  # Previous sessions
     upstream_data_download_directory: Optional[Path] = None  # Set by microscope config
     upstream_data_tiff_locations: List[str] = ["processed"]  # Location of CLEM TIFFs
+
+    model_search_directory: str = "processing"
+
     failure_queue: str = ""
     auth_key: str = ""
     auth_algorithm: str = ""
@@ -93,6 +106,7 @@ def get_machine_config() -> MachineConfig:
         rsync_basepath=Path("dls/tmp"),
         murfey_db_credentials="",
         crypto_key="",
+        default_model="/tmp/weights.h5",
     )
     if settings.murfey_machine_configuration:
         microscope = get_microscope()

src/murfey/server/spa/__init__.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from datetime import datetime
+from functools import lru_cache
+from pathlib import Path
+
+from fastapi import APIRouter
+from sqlmodel import select
+
+from murfey.server import get_machine_config
+from murfey.server.murfey_db import murfey_db
+from murfey.util.db import Session
+
+# Create APIRouter class object
+router = APIRouter()
+
+
+@lru_cache(maxsize=5)
+def _cryolo_model_path(visit: str) -> Path:
+    machine_config = get_machine_config()
+    if machine_config.model_search_directory:
+        visit_directory = (
+            machine_config.rsync_basepath
+            / (machine_config.rsync_module or "data")
+            / str(datetime.now().year)
+            / visit
+        )
+        possible_models = list(
+            (visit_directory / machine_config.model_search_directory).glob("*.h5")
+        )
+        if possible_models:
+            return sorted(possible_models, key=lambda x: x.stat().st_ctime)[-1]
+    return machine_config.default_model
+
+
+@router.get("/sessions/{session_id}/cryolo_model")
+def get_cryolo_model_path(session_id: int, db=murfey_db):
+    visit = db.exec(select(Session).where(Session.session_id == session_id)).one().visit
+    return {"model_path": _cryolo_model_path(visit)}