Add more tests for backbone ckpt

Author: gitttt-1234

tests/inference/test_predictors.py

@@ -216,11 +216,32 @@ def test_topdown_predictor(
     )
 
     assert np.all(np.abs(head_layer_ckpt - model_weights) < 1e-6)
-    print(
-        f"centered instance model: ",
-        predictor.inference_model.instance_peaks.torch_model.model,
+
+    # load only backbone and head ckpt as None - centered instance
+    predictor = Predictor.from_model_paths(
+        [minimal_instance_ckpt],
+        backbone_ckpt_path=Path(minimal_instance_ckpt) / "best.ckpt",
+        head_ckpt_path=None,
+        peak_threshold=0.03,
+        max_instances=6,
+        preprocess_config=OmegaConf.create(preprocess_config),
+    )
+
+    ckpt = torch.load(Path(minimal_instance_ckpt) / "best.ckpt")
+    backbone_ckpt = ckpt["state_dict"][
+        "model.backbone.enc.encoder_stack.0.blocks.0.weight"
+    ][0, 0, :].numpy()
+
+    model_weights = (
+        next(predictor.inference_model.instance_peaks.torch_model.model.parameters())[
+            0, 0, :
+        ]
+        .detach()
+        .numpy()
     )
 
+    assert np.all(np.abs(backbone_ckpt - model_weights) < 1e-6)
+
     # check loading diff head ckpt for centroid
     preprocess_config = {
         "is_rgb": False,
@@ -238,11 +259,32 @@ def test_topdown_predictor(
         preprocess_config=OmegaConf.create(preprocess_config),
     )
 
-    print(
-        f"centroid model: ", predictor.inference_model.centroid_crop.torch_model.model
+    ckpt = torch.load(Path(minimal_instance_ckpt) / "best.ckpt")
+    backbone_ckpt = ckpt["state_dict"][
+        "model.backbone.enc.encoder_stack.0.blocks.0.weight"
+    ][0, 0, :].numpy()
+
+    model_weights = (
+        next(predictor.inference_model.centroid_crop.torch_model.model.parameters())[
+            0, 0, :
+        ]
+        .detach()
+        .numpy()
+    )
+
+    assert np.all(np.abs(backbone_ckpt - model_weights) < 1e-6)
+
+    # load only backbone and head ckpt as None - centroid
+    predictor = Predictor.from_model_paths(
+        [minimal_instance_centroid_ckpt],
+        backbone_ckpt_path=Path(minimal_instance_centroid_ckpt) / "best.ckpt",
+        head_ckpt_path=None,
+        peak_threshold=0.03,
+        max_instances=6,
+        preprocess_config=OmegaConf.create(preprocess_config),
     )
 
-    ckpt = torch.load(Path(minimal_instance_ckpt) / "best.ckpt")
+    ckpt = torch.load(Path(minimal_instance_centroid_ckpt) / "best.ckpt")
     backbone_ckpt = ckpt["state_dict"][
         "model.backbone.enc.encoder_stack.0.blocks.0.weight"
     ][0, 0, :].numpy()
@@ -261,7 +303,6 @@ def test_topdown_predictor(
 def test_single_instance_predictor(
     minimal_instance,
     minimal_instance_ckpt,
-    minimal_instance_centroid_ckpt,
     minimal_instance_bottomup_ckpt,
 ):
     """Test SingleInstancePredictor module."""
@@ -453,6 +494,55 @@ def test_single_instance_predictor(
         # save the original config back
         OmegaConf.save(_config, f"{minimal_instance_ckpt}/training_config.yaml")
 
+    _config = OmegaConf.load(f"{minimal_instance_ckpt}/training_config.yaml")
+
+    config = _config.copy()
+
+    try:
+        head_config = config.model_config.head_configs.centered_instance
+        del config.model_config.head_configs.centered_instance
+        OmegaConf.update(
+            config, "model_config.head_configs.single_instance", head_config
+        )
+        del config.model_config.head_configs.single_instance.confmaps.anchor_part
+        OmegaConf.update(config, "data_config.preprocessing.scale", 0.9)
+
+        OmegaConf.save(config, f"{minimal_instance_ckpt}/training_config.yaml")
+
+        # check loading diff head ckpt
+        preprocess_config = {
+            "is_rgb": False,
+            "crop_hw": None,
+            "max_width": None,
+            "max_height": None,
+        }
+
+        predictor = Predictor.from_model_paths(
+            [minimal_instance_bottomup_ckpt],
+            backbone_ckpt_path=Path(minimal_instance_ckpt) / "best.ckpt",
+            head_ckpt_path=None,
+            peak_threshold=0.03,
+            max_instances=6,
+            preprocess_config=OmegaConf.create(preprocess_config),
+        )
+
+        ckpt = torch.load(Path(minimal_instance_ckpt) / "best.ckpt")
+        backbone_ckpt = ckpt["state_dict"][
+            "model.backbone.enc.encoder_stack.0.blocks.0.weight"
+        ][0, 0, :].numpy()
+
+        model_weights = (
+            next(predictor.inference_model.torch_model.model.parameters())[0, 0, :]
+            .detach()
+            .numpy()
+        )
+
+        assert np.all(np.abs(backbone_ckpt - model_weights) < 1e-6)
+
+    finally:
+        # save the original config back
+        OmegaConf.save(_config, f"{minimal_instance_ckpt}/training_config.yaml")
+
 
 def test_bottomup_predictor(
     minimal_instance, minimal_instance_bottomup_ckpt, minimal_instance_ckpt
@@ -611,3 +701,26 @@ def test_bottomup_predictor(
     print(model_weights)
 
     assert np.all(np.abs(head_layer_ckpt - model_weights) < 1e-6)
+
+    # load only backbone and head ckpt as None
+    predictor = Predictor.from_model_paths(
+        [minimal_instance_bottomup_ckpt],
+        backbone_ckpt_path=Path(minimal_instance_ckpt) / "best.ckpt",
+        head_ckpt_path=None,
+        peak_threshold=0.03,
+        max_instances=6,
+        preprocess_config=OmegaConf.create(preprocess_config),
+    )
+
+    ckpt = torch.load(Path(minimal_instance_ckpt) / "best.ckpt")
+    backbone_ckpt = ckpt["state_dict"][
+        "model.backbone.enc.encoder_stack.0.blocks.0.weight"
+    ][0, 0, :].numpy()
+
+    model_weights = (
+        next(predictor.inference_model.torch_model.model.parameters())[0, 0, :]
+        .detach()
+        .numpy()
+    )
+
+    assert np.all(np.abs(backbone_ckpt - model_weights) < 1e-6)