Instantiate AstroPredictor with in-memory model weights

src/deepdisc/astrodet/astrodet.py

@@ -255,7 +255,7 @@
         inputs = cv2.imread("input.jpg")
         outputs = pred(inputs)
     """
-    def __init__(self, cfg, lazy=False, cfglazy=None):
+    def __init__(self, cfg, lazy=False, cfglazy=None, checkpoint=None):
         self.cfg = copy.deepcopy(cfg) # cfg can be modified by model
 
         if "model" in self.cfg: # This is when were using a LazyConfig-style model in the solo config
@@ -271,8 +271,14 @@
             self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
 
         checkpointer = DetectionCheckpointer(self.model)
-        checkpointer.load(cfg.MODEL.WEIGHTS)
 
+        # If we provide AstroPredictor with a checkpoint already loaded in memory
+        # just simply load the weights into the model.
+        if checkpoint:
+            checkpointer._load_model(checkpoint)
+        else:
+            checkpointer.load(cfg.train.init_checkpoint)
+
         self.aug = T.ResizeShortestEdge(
             [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
         )

src/deepdisc/inference/predictors.py

@@ -1,7 +1,7 @@
 import deepdisc.astrodet.astrodet as toolkit
 
 
-def return_predictor_transformer(cfg):
+def return_predictor_transformer(cfg, checkpoint=None):
     """This function returns a trained model and its config file.
 
     Used for models with lazy config files. Also assumes a cascade roi head structure.
@@ -15,7 +15,7 @@
         torch model
 
     """
-    predictor = toolkit.AstroPredictor(cfg)
+    predictor = toolkit.AstroPredictor(cfg, checkpoint=checkpoint)
     return predictor