Merge pull request #195 from talmolab/divya/close-videos-before-multiprocessing

Close videos before creating data loaders

Author: gitttt-1234

environment.yml

@@ -9,7 +9,7 @@ dependencies:
   - python=3.9
   - pytorch-cuda=11.8
   - numpy
-  - sleap-io
+  - sleap-io>=0.2.0
   - pydantic
   - lightning
   - cudnn

environment_cpu.yml

@@ -7,7 +7,7 @@ channels:
 dependencies:
   - python=3.9
   - numpy
-  - sleap-io
+  - sleap-io>=0.2.0
   - pytorch
   - pydantic
   - lightning

environment_mac.yml

@@ -7,7 +7,7 @@ channels:
 dependencies:
   - python=3.9
   - numpy
-  - sleap-io
+  - sleap-io>=0.2.0
   - pydantic
   - lightning
   - pytorch

sleap_nn/data/custom_datasets.py

@@ -161,7 +161,8 @@ def _fill_cache(self):
                 self.cache[lf_idx] = img
 
         for video in self.labels.videos:
-            video.close()
+            if video.is_open:
+                video.close()
 
     def _get_video_idx(self, lf):
         """Return indsample of `lf.video` in `labels.videos`."""

sleap_nn/training/model_trainer.py

@@ -475,6 +475,16 @@ def _create_data_loaders_torch_dataset(self):
             else True
         )
 
+        # If using caching, close the videos to prevent `h5py objects can't be pickled error` when num_workers > 0.
+        if "cache_img" in self.data_pipeline_fw:
+            for video in self.train_labels.videos:
+                if video.is_open:
+                    video.close()
+
+            for video in self.val_labels.videos:
+                if video.is_open:
+                    video.close()
+
         # train
         self.train_data_loader = DataLoader(
             dataset=self.train_dataset,