fix: improve inferer ckpt loading

Author: okunator

cellseg_models_pytorch/inference/_base_inferer.py

@@ -4,6 +4,7 @@
 from itertools import chain
 from pathlib import Path
 from typing import Callable, Dict, List, Tuple, Union
+from urllib.parse import urlparse
 
 import numpy as np
 import torch
@@ -150,17 +151,31 @@ def __init__(
 
         # load weights and set devices
         if checkpoint_path is not None:
-            ckpt = torch.load(
-                checkpoint_path, map_location=lambda storage, loc: storage
-            )
+            checkpoint_path = Path(checkpoint_path)
+            # check if path is url or local and load weigths to memory
+            if urlparse(checkpoint_path.as_posix()).scheme:
+                state_dict = torch.hub.load_state_dict_from_url(checkpoint_path)
+            else:
+                state_dict = torch.load(
+                    checkpoint_path, map_location=lambda storage, loc: storage
+                )
+
+            # if the checkpoint is from lightning, the ckpt file contains a lot of other
+            # stuff than just the state dict.
+            if "state_dict" in state_dict.keys():
+                state_dict = state_dict["state_dict"]
 
+            # try loading the weights to the model
             try:
-                self.model.load_state_dict(ckpt["state_dict"], strict=True)
+                msg = self.model.load_state_dict(state_dict, strict=True)
             except RuntimeError:
-                new_ckpt = self._strip_state_dict(ckpt)
-                self.model.load_state_dict(new_ckpt["state_dict"], strict=True)
+                new_ckpt = self._strip_state_dict(state_dict)
+                msg = self.model.load_state_dict(new_ckpt, strict=True)
             except BaseException as e:
-                print(e)
+                raise RuntimeError(f"Error when loading checkpoint: {e}")
+
+            print(f"Loading weights: {checkpoint_path} for inference.")
+            print(msg)
 
         assert device in ("cuda", "cpu", "mps")
         if device == "cpu":
@@ -213,6 +228,12 @@ def infer(self, mixed_precision: bool = False) -> None:
             `classes_type`, `classes_sem`, `offsets`. See more in the
             `FileHandler.save_masks` docs.
 
+
+        Parameters
+        ----------
+            mixed_precision : bool, default=False
+                If True, inference is performed with mixed precision.
+
         Attributes
         ----------
             - out_masks : Dict[str, Dict[str, np.ndarray]]
@@ -224,11 +245,6 @@ def infer(self, mixed_precision: bool = False) -> None:
                 The soft masks for each image. I.e. the soft predictions of the trained
                 model The keys are the image names and the values are dictionaries of
                 the soft masks. E.g. {"sample1": {"type": [H, W], "aux": [C, H, W]}}
-
-        Parameters
-        ----------
-            mixed_precision : bool, default=False
-                If True, inference is performed with mixed precision.
         """
         self.soft_masks = {}
         self.out_masks = {}
@@ -291,15 +307,13 @@ def infer(self, mixed_precision: bool = False) -> None:
     def _strip_state_dict(self, ckpt: Dict) -> OrderedDict:
         """Strip te first 'model.' (generated by lightning) from the state dict keys."""
         state_dict = OrderedDict()
-        for k, w in ckpt["state_dict"].items():
+        for k, w in ckpt.items():
             if "num_batches_track" not in k:
-                # new_key = k.strip("model")[1:]
                 spl = ["".join(kk) for kk in k.split(".")]
                 new_key = ".".join(spl[1:])
                 state_dict[new_key] = w
-        ckpt["state_dict"] = state_dict
 
-        return ckpt
+        return state_dict
 
     def _check_and_set_head_args(self) -> None:
         """Check the model has matching head names with the head args and set them."""