Merge pull request #1439 from fdeguire03/dev

Improve get_file_size efficiency for .npy files
flatironinstitute · Dec 13, 2024 · 2b29d49 · 2b29d49
2 parents 0d9c159 + 5ce0ce5
commit 2b29d49
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/caiman/base/movies.py b/caiman/base/movies.py
@@ -2026,7 +2026,14 @@ def get_file_size(file_name, var_name_hdf5:str='mov') -> tuple[tuple, Union[int,
                     raise Exception('Variable not found. Use one of the above')
                 T, dims = siz[0], siz[1:]
             elif extension in ('.npy', ):
-                shape = np.load(file_name, allow_pickle=False).shape
+                with open(file_name, 'rb') as f:
+                    version = np.lib.format.read_magic(f)
+                    if version == (1, 0):
+                        shape, _, _ = np.lib.format.read_array_header_1_0(f)
+                    elif version == (2, 0):
+                        shape, _, _ = np.lib.format.read_array_header_2_0(f)
+                    else:
+                        raise ValueError(f"Unsupported .npy file version: {version}. Update caiman.base.movies.get_file_size() to handle it.")
                 T = shape[0]
                 dims = shape[1:]
             elif extension in ('.sbx'):