module refactor; add more input test

brainscore_vision/model_helpers/activations/temporal/core/executor.py

@@ -25,9 +25,16 @@ class JoblibMapper:
     def __init__(self, num_threads: int):
         self._num_threads = num_threads
         self._pool = Parallel(n_jobs=num_threads, verbose=False, backend="loky")
+        self._not_supported = False
 
     def map(self, func, *data):
-        return self._pool(delayed(func)(*x) for x in zip(*data))
+        from joblib.externals.loky.process_executor import TerminatedWorkerError
+        if not self._not_supported:
+            try:
+                return self._pool(delayed(func)(*x) for x in zip(*data))
+            except TerminatedWorkerError:
+                self._not_supported = True
+        return [func(*x) for x in zip(*data)]
 
 
 class BatchExecutor:

brainscore_vision/model_helpers/activations/temporal/inputs/image.py

@@ -27,10 +27,13 @@ def from_path(path):
 
     def to_img(self):
         return PILImage.fromarray(self.to_numpy())
+
+    def get_frame(self):
+        return np.array(PILImage.open(self._path).convert('RGB'))
 
     # return (H, W, C[RGB])
     def to_numpy(self):
-        arr = np.array(PILImage.open(self._path).convert('RGB'))
+        arr = self.get_frame()
 
         if arr.shape[:2][::-1] != self._size:
             arr = batch_2d_resize(arr[None,:], self._size, "bilinear")[0]

brainscore_vision/model_helpers/activations/temporal/inputs/video.py

@@ -120,7 +120,6 @@ def get_frames(self, indices):
 
     ### I/O
     def from_path(path):
-        path = path
         fps, end, size = get_video_stats(path)
         start = 0
         return Video(path, fps, start, end, size)

brainscore_vision/model_helpers/activations/temporal/utils.py

@@ -43,7 +43,7 @@ def stack_with_nan_padding_(arr_list, axis=0, dtype=np.float16):
     return result
 
 
-def stack_with_nan_padding(arr_list, axis=0, dtype=np.float16):
+def stack_with_nan_padding(arr_list, axis=0, dtype=None):
     # Get shapes of all arrays
     shapes = [np.array(arr.shape) for arr in arr_list]
     max_shape = np.max(shapes, axis=0)
@@ -58,7 +58,7 @@ def stack_with_nan_padding(arr_list, axis=0, dtype=np.float16):
 
     result = np.stack(results, axis=axis)
     result = np.swapaxes(result, 0, axis)
-    if result.dtype != dtype:
+    if dtype is not None and result.dtype != dtype:
         result = result.astype(dtype)
 
     return result

tests/test_model_helpers/temporal/activations/test_inputs.py

@@ -40,6 +40,7 @@ def test_video():
     video3 = video1.set_window(-10, 0, padding="repeat")
     video4 = video1.set_window(-20, -10, padding="repeat")
     assert (video3.to_numpy() == video4.to_numpy()).all()
+    assert (video3.to_numpy()[0] == video1.to_numpy()[0]).all()
 
     assert video2.fps == 30
     assert video2.set_fps(1).to_numpy().shape[0] == 1