Merge branch

Author: gitttt-1234

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "av",
     "kornia",
     "hydra-core",
-    "sleap-io>=0.1.0",
+    "sleap-io==0.1.10",
 ]
 dynamic = ["version", "readme"]
 

sleap_nn/training/model_trainer.py

@@ -390,22 +390,8 @@ def train(
         self,
         backbone_trained_ckpts_path: Optional[str] = None,
         head_trained_ckpts_path: Optional[str] = None,
-        delete_bin_files_after_training: bool = True,
-        chunks_dir_path: Optional[str] = None,
     ):
-        """Initiate the training by calling the fit method of Trainer.
-
-        Args:
-            backbone_trained_ckpts_path: Path of the `ckpt` file with which the backbone
-                 is initialized. If `None`, random init is used.
-            head_trained_ckpts_path: Path of the `ckpt` file with which the head layers
-                 are initialized. If `None`, random init is used.
-            delete_bin_files_after_training: If `False`, the `bin` files are retained after
-                training. Else, the `bin` files are deleted.
-            chunks_dir_path: Path to chunks dir (this dir should contain `train_chunks`
-                and `val_chunks` folder.). If `None`, `bin` files are generated.
-
-        """
+        """Initiate the training by calling the fit method of Trainer."""
         logger = []
 
         if self.config.trainer_config.save_ckpt:

tests/training/test_model_trainer.py

@@ -376,58 +376,6 @@ def test_trainer_load_trained_ckpts(config, tmp_path, minimal_instance_ckpt):
     assert np.all(np.abs(head_layer_ckpt - model_ckpt) < 1e-6)
 
 
-@pytest.mark.skipif(
-    sys.platform.startswith("li"),
-    reason="Flaky test (The training test runs on Ubuntu for a long time: >6hrs and then fails.)",
-)
-# TODO: Revisit this test later (Failing on ubuntu)
-def test_reuse_bin_files(config, tmp_path: str):
-    """Test reusing `.bin` files."""
-    # Centroid model
-    centroid_config = config.copy()
-    head_config = config.model_config.head_configs.centered_instance
-    OmegaConf.update(centroid_config, "model_config.head_configs.centroid", head_config)
-    del centroid_config.model_config.head_configs.centered_instance
-    del centroid_config.model_config.head_configs.centroid["confmaps"].part_names
-
-    OmegaConf.update(
-        centroid_config,
-        "trainer_config.save_ckpt_path",
-        f"{tmp_path}/test_model_trainer/",
-    )
-
-    if (Path(centroid_config.trainer_config.save_ckpt_path) / "best.ckpt").exists():
-        os.remove(
-            (
-                Path(centroid_config.trainer_config.save_ckpt_path) / "best.ckpt"
-            ).as_posix()
-        )
-        os.remove(
-            (
-                Path(centroid_config.trainer_config.save_ckpt_path) / "last.ckpt"
-            ).as_posix()
-        )
-        shutil.rmtree(
-            (
-                Path(centroid_config.trainer_config.save_ckpt_path) / "lightning_logs"
-            ).as_posix()
-        )
-
-    OmegaConf.update(centroid_config, "trainer_config.save_ckpt", True)
-    OmegaConf.update(centroid_config, "trainer_config.use_wandb", False)
-    OmegaConf.update(centroid_config, "trainer_config.max_epochs", 1)
-    OmegaConf.update(centroid_config, "trainer_config.steps_per_epoch", 10)
-
-    # test reusing bin files
-    trainer1 = ModelTrainer(centroid_config)
-    trainer1.train(delete_bin_files_after_training=False)
-
-    trainer2 = ModelTrainer(centroid_config)
-    trainer2.train(
-        chunks_dir_path=(trainer1.train_input_dir).split("train_chunks")[0],
-    )
-
-
 def test_topdown_centered_instance_model(config, tmp_path: str):
 
     # unet