added data generation and particle detection code

Training_Data/Particle_Tracking_Training_Data.py

@@ -0,0 +1,141 @@
+import numpy as np
+from numpy import pi
+import tensorflow as tf
+
+class Particle_Tracking_Training_Data(tf.Module):
+    def __init__(self, Nt, rings=True):
+        self.Nt = int(Nt)
+        self.Ny = self.Nx = 256
+        self.d = 3
+        ximg = [[[i, j] for i in np.arange(self.Ny)]
+            for j in np.arange(self.Nx)]
+        self.ximg = np.float32(ximg)
+
+        x = np.arange(self.Nx) - self.Nx//2
+        y = np.arange(self.Ny) - self.Ny//2
+        X0, Y0 = np.meshgrid(x, y)
+        self.X = np.float32(X0)
+        self.Y = np.float32(Y0)
+
+        if rings:
+            self.ring_indicator = 1.
+        else:
+            self.ring_indicator = 0.
+
+        self._gen_video = tf.function(
+            input_signature=(
+                tf.TensorSpec(
+                    shape=[self.Ny, self.Nx, self.Nt, None], dtype=tf.float32),
+                tf.TensorSpec(shape=[self.Nt, None], dtype=tf.float32),
+                tf.TensorSpec(shape=[], dtype=tf.float32),
+                tf.TensorSpec(shape=[], dtype=tf.float32),
+                tf.TensorSpec(shape=[], dtype=tf.float32),)
+        )(self._gen_video)
+
+        self._gen_labels = tf.function(
+            input_signature=(
+                tf.TensorSpec(
+                    shape=[self.Ny, self.Nx, self.Nt, None], dtype=tf.float32),)
+        )(self._gen_labels)
+
+    def __call__(self, kappa, a, IbackLevel, Nparticles, sigma_motion):
+        """a: spot radius scale factor (1.5-4 is a reasonable range)
+        kappa: noise level (around 0.1 or so)
+        IbackLevel: intensity level of the random background relative to maximum (must be between zero and one)
+        Nparticles: the number of particles in the video (larger numbers means slower run time)
+        sigma_motion: the standard deviation of the random Brownian motion per video frame"""
+        ## random brownian motion paths
+        ## Nt, Nparticles, 3
+        xi = self._sample_motion(Nparticles, sigma_motion)
+
+        #### translate track positions to img coords
+        ## Ny, Nx, Nt, Np, 2
+        XALL = (self.ximg[:, :, None, None, :]
+                - xi[None, None, :, :, :2])
+        ## Ny, Nx, Nt, Np
+        r = tf.math.sqrt(XALL[..., 0]**2 + XALL[..., 1]**2)
+        z = xi[..., 2]
+
+        ### generate video
+        I = self._gen_video(r, z, kappa, a, IbackLevel)
+
+        ### generate labels
+        labels = self._gen_labels(r)
+
+        return I, labels, xi
+
+    @staticmethod
+    def rand(n):
+        return tf.random.uniform([n], dtype=tf.float32)
+
+    @tf.function(
+        input_signature=(
+            tf.TensorSpec(shape=[], dtype=tf.int32),
+            tf.TensorSpec(shape=[], dtype=tf.float32),))
+    def _sample_motion(self, Nparticles, sigma_motion):
+        #### boundaries
+        b_lower = tf.constant(
+            [-10, -10, -30.], tf.float32)
+        b_upper = tf.constant(
+            [self.Nx+10, self.Ny+10, 30.], tf.float32)
+        #### uniform random initial possitions
+        U = tf.random.uniform(
+            [1, Nparticles, self.d],
+            dtype=tf.float32)
+        X0 = b_lower + (b_upper - b_lower)*U
+        #### normal increments
+        dX = tf.random.normal(
+            [self.Nt, Nparticles, self.d],
+            stddev=sigma_motion,
+            dtype=tf.float32)
+        #### unbounded Brownian motion
+        X = X0 + tf.math.cumsum(dX, axis=0)
+        #### reflected brownian motion
+        ## note that this is imperfect,
+        ## if increments are very large it wont work
+        X = tf.math.abs(X - b_lower) + b_lower
+        X = -tf.math.abs(b_upper - X) + b_upper
+        return X
+
+    def _gen_video(self, r, z, kappa, a, IbackLevel):
+        uw = (0.5 + self.rand(1))/2.
+        un = tf.floor(3*self.rand(1))
+        uampRing = 0.2 + 0.8*self.rand(1)
+        ufade = 15 + 10*self.rand(1)
+        rmax = ufade*(un/uw)**(2./3.)
+        ufadeMax = 0.85
+        fade = (1. - ufadeMax*tf.abs(tf.tanh(z/ufade)))
+        core = tf.exp(-(r**2/(8.*a))**2)
+        ring = fade*(tf.exp(-(r - z)**4/(a)**4)
+                + 0.5*uampRing*tf.cast(r<z, tf.float32))
+        I = tf.transpose(
+            tf.reduce_sum(
+                fade*(core + self.ring_indicator*ring),
+                axis=3),
+            [2, 0, 1]) # Nt, Ny, Nx
+        I += IbackLevel*tf.sin(
+            self.rand(1)*6*pi/512*tf.sqrt(
+                self.rand(1)*(self.X - self.rand(1)*512)**2
+                    + self.rand(1)*(self.Y - self.rand(1)*512)**2))
+        I += tf.random.normal(
+            [self.Nt, self.Ny, self.Nx],
+            stddev=kappa,
+            dtype=tf.float32)
+        Imin = tf.reduce_min(I)
+        Imax = tf.reduce_max(I)
+        I = (I - Imin)/(Imax - Imin)
+        I = tf.round(I*tf.maximum(256., tf.round(2**16*self.rand(1))))
+        return I
+
+    def _gen_labels(self, r):
+        R_detect = 3.
+        ## (Ny, Nx, Nt)
+        detectors = tf.reduce_sum(
+            tf.cast(r[::2, ::2, :, :] < R_detect, tf.int32),
+            axis=3)
+        ## (Nt, Ny, Nx)
+        P = tf.transpose(
+            tf.cast(detectors > 0, tf.int32), [2, 0, 1])
+        ## (Nt, Ny, Nx, 2)
+        labels = tf.stack([1-P, P], 3)
+        return labels