Directly generate videos

zebranoise/__init__.py

@@ -1,4 +1,5 @@
 from .perlin_stimulus import PerlinStimulus
 from .util import generate_frames
+from .easy import zebra_noise
 
 Perl = PerlinStimulus # Backward compatibility

zebranoise/easy.py

@@ -0,0 +1,19 @@
+import imageio
+import warnings
+from .util import generate_frames, filter_frames_index_function, apply_filters, discretize
+
+def zebra_noise(output_file, xsize, ysize, tdur, levels=10, xyscale=.2, tscale=50, fps=30, xscale=1.0, yscale=1.0, seed=0, filters=[("comb", 0.08)]):
+    tsize = int(tdur*fps)
+    textra = (tscale - (tsize % tscale)) % tscale
+    if textra > 0:
+        warnings.warn(f"Adding {textra} extra timepoints to make tscale a multiple of tdur")
+    tsize += textra
+    get_index = filter_frames_index_function(filters, tsize)
+    writer = imageio.get_writer(output_file, fps=fps)
+    for _i in range(0, tsize):
+        i = get_index(_i)
+        frame = generate_frames(xsize, ysize, tsize, [i], levels=levels, xyscale=xyscale, tscale=tscale, xscale=xscale, yscale=yscale, seed=seed)
+        filtered = apply_filters(frame[None], filters)[0] # TODO I don't think this will work with the photodiode filter
+        disc = discretize(filtered[:,:,0])
+        writer.append_data(disc)
+    writer.close()

zebranoise/perlin_stimulus.py

@@ -9,7 +9,7 @@
 from imageio_ffmpeg import get_ffmpeg_exe
 
 from . import _perlin
-from .util import generate_frames, filter_frames, filter_frames_index_function, XYSCALEBASE
+from .util import generate_frames, filter_frames, filter_frames_index_function, XYSCALEBASE, discretize
 
 
 class PerlinStimulus:
@@ -109,23 +109,6 @@ def cache_filename(self, batch=None):
         if batch == "stats":
             return str(self.cachedir.joinpath(f"perlcache_{h}_stats.npz"))
         return str(self.cachedir.joinpath(f"perlcache_{h}_{batch}.npy"))
-    @classmethod
-    def discretize(cls, im):
-        """Convert movie to an unsigned 8-bit integer
-
-        Parameters
-        ----------
-        im : 3D float ndarray, values ∈ [0,1]
-            Noise movie
-
-        Returns
-        -------
-        3D int ndarray, values ∈ [0,255]
-            Noise movie
-        """
-        im *= 255
-        ret = im.astype(np.uint8)
-        return ret
     def generate_frame(self, t=0, filters=[]):
         """Generate and return a single image of noise.
 
@@ -144,19 +127,12 @@ def generate_frame(self, t=0, filters=[]):
         2-dimensional ndarray
             An image of the noise
         """
-        arr = generate_frames(self.size[0], self.size[1], self.size[2], timepoints=[t], levels=self.levels,
+        arr = generate_frames(self.size[0], self.size[1], self.size[2], timepoints=[t] if not hasattr(t, "__iter__") else t, levels=self.levels,
                               xyscale=self.xyscale, tscale=self.tscale, xscale=self.xscale,
                               yscale=self.yscale, seed=self.seed)
         if self.demean in ["both", "time"]:
             arr -= np.mean(arr, axis=(0,1), keepdims=True)
-        for f in filters:
-            if isinstance(f, str):
-                n = f
-                args = []
-            else:
-                n = f[0]
-                args = f[1:]
-            arr = filter_frames(arr, n, *args)
+        arr = apply_filters(arr, filters)
         return arr.squeeze()
 
     def generate_batch(self):
@@ -272,7 +248,7 @@ def save_video(self, fn, loop=1, filters=[], bitrate=20):
                     n = f[0]
                     args = f[1:]
                 data = filter_frames(data, n, *args)
-            data = self.discretize(data)
+            data = discretize(data)
             assert data.dtype == 'uint8'
             for j in range(0, data.shape[2]):
                 imageio.imsave(self.tmpdir.joinpath(f"_frame{filtind(i):05}.tif"), data[:,:,j], format="pillow", compression="tiff_adobe_deflate")

zebranoise/util.py

@@ -1,5 +1,5 @@
 import numpy as np
-
+import scipy.ndimage
 from . import _perlin
 
 XYSCALEBASE = 100
@@ -73,6 +73,18 @@ def filter_frames(im, filt, *args):
         return filt(im)
     raise ValueError("Invalid filter specified")
 
+def apply_filters(arr, filters):
+    for f in filters:
+        if isinstance(f, str):
+            n = f
+            args = []
+        else:
+            n = f[0]
+            args = f[1:]
+        arr = filter_frames(arr, n, *args)
+    return arr
+
+
 def filter_frames_index_function(filters, nframes):
     """Reordering frames in the video based on the filter.
 
@@ -102,6 +114,23 @@ def filter_frames_index_function(filters, nframes):
         return lambda x : nframes - x - 1
     return lambda x : x
 
+def discretize(im):
+    """Convert movie to an unsigned 8-bit integer
+
+    Parameters
+    ----------
+    im : 3D float ndarray, values ∈ [0,1]
+        Noise movie
+
+    Returns
+    -------
+    3D int ndarray, values ∈ [0,255]
+        Noise movie
+    """
+    im *= 255
+    ret = im.astype(np.uint8)
+    return ret
+
 def generate_frames(xsize, ysize, tsize, timepoints, levels=10, xyscale=.5, tscale=1, xscale=1.0, yscale=1.0, seed=0):
     """Preprocess arguments before passing to the C implementation of Perlin noise.
     """