From 566ceed05335c9c6743cfacc424494b87bdb4afc Mon Sep 17 00:00:00 2001 From: rhoadesScholar Date: Tue, 13 Aug 2024 16:47:34 -0400 Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20Fix=20COSEM=20model=20loa?= =?UTF-8?q?ding.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/cellmap_models/pytorch/__init__.py | 3 +-- src/cellmap_models/pytorch/cosem/README.md | 2 +- src/cellmap_models/pytorch/cosem/load_model.py | 1 + 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/cellmap_models/pytorch/__init__.py b/src/cellmap_models/pytorch/__init__.py index 5cdcb2d..963639f 100755 --- a/src/cellmap_models/pytorch/__init__.py +++ b/src/cellmap_models/pytorch/__init__.py @@ -1,2 +1 @@ -from . import cosem -from . import cellpose +from . import cosem, cellpose, untrained_models diff --git a/src/cellmap_models/pytorch/cosem/README.md b/src/cellmap_models/pytorch/cosem/README.md index ef9c8e3..1629ccb 100755 --- a/src/cellmap_models/pytorch/cosem/README.md +++ b/src/cellmap_models/pytorch/cosem/README.md @@ -25,7 +25,7 @@ Each model has a separate backbone and single layer prediction head. The `backbo ```python import cellmap_models.cosem as cosem_models model = cosem_models.load_model('setup04/1820500') -backnone = model.backbone +backbone = model.backbone head = model.prediction_head ``` diff --git a/src/cellmap_models/pytorch/cosem/load_model.py b/src/cellmap_models/pytorch/cosem/load_model.py index 5d0c8b6..1857087 100755 --- a/src/cellmap_models/pytorch/cosem/load_model.py +++ b/src/cellmap_models/pytorch/cosem/load_model.py @@ -81,6 +81,7 @@ def load_model(checkpoint_name: str) -> torch.nn.Module: new_checkpoint["model"].pop(key) continue new_key = key.replace("architecture.", "") + new_key = new_key.replace("unet.", "backbone.") new_checkpoint["model"][new_key] = new_checkpoint["model"].pop(key) model.load_state_dict(new_checkpoint["model"]) model.eval()