Release DeepTrack2 1.6.0 (#191)

giovannivolpe · BenjaminMidtvedt · github-actions[bot] · web-flow · commit c751e5dccc51 · 2023-09-13T09:58:34.000+02:00
* chore: autopublish 2022-07-26T13:54:44Z * Remove create-badges job * Delete test.py * Add multi-head masked attention * Update multi-head gated attention to match parent layer * Update documentation * Test multi-head masked attention * allow gated attention layers to use bias * test bias in gated attention layers * set return_attention_weights to False to avoid multi-outputs Use MultiHeadSelfAttention and MultiHeadGatedSelfAttention if want to return the attention weights * reformat gnns/layers.py This commit adds new message-passing graph layers (MPN) and graph convolutional layers to dt, including vanilla MPN, GRUMPN, Masked-attention FGNN, and GraphTransformer. * Update layers.py * Update test_layers.py * Update models.py * Update test_models.py * Update test_models.py * Fix indexing problems related to tf.gather * Allow multi-inputs in ContinuousGenerator * Fix bad conversion to integer * version bump * Fix phase correction at focus and offset calculation * Fix phase correction in propagation * Fix mie phase out of foucs * Fix mie phase out of foucs * Update README.md * Bm/version 1.4.0 (#137) * Update layers.py * Update convolutional.py Transformer-based models can now be reused and expanded quickly and easily * Update documentation * Update Transformer-based models * Delete classifying_MNIST_vit_tutorial.ipynb * Create classifying_MNIST_vit_tutorial.ipynb * Update datasets.py * Allows kwargs as inputs in single_layer_call * Update embeddings.py * masked transformers * reformat transformer models * Create trajectory_analysis_tutorial.ipynb * Add Variational autoencoders * Add variational autoencoders * Update vae.py * Create MNIST_VAE_tutorial.ipynb * Update MNIST_VAE_tutorial.ipynb * Create folder for course examples * Update README.md * Update README.md * Update examples * Update README.md * Update README.md * Update MNIST VAE examples * Added MLP regression example * Update README.md * Create image_segmentation_Unet.ipynb * Update README.md * Documented and tested cell_counting_tutorial.ipynb * improve dnn example * Shift variant mie * Position mie scatterer correctly * implement set z * implement mnist v1 * implement z dependence * remove logging * Implement flattening methods * Implement pooling and resizing * Implement TensorflowDataset * Finalize MNIST * Implement Malaria classification * alpha0 release * fix batchsize in fit * implement dataset.take * Implement datasets * fix phase in mie * Fix mie positioning and focusing * Commit to new branch * add tensorflow datasets dependence * remove test Co-authored-by: Jesús Pineda <jesus.pineda@physics.gu.se> Co-authored-by: Jesús Pineda <36273229+JesusPinedaC@users.noreply.github.com> Co-authored-by: Benjamin Midtvedt <benmid@student.chalmers.se> Co-authored-by: Ccx55 <ccx555@gmail.com> * Add tensorflow datasets to the list of dependencies. * Read requirements.txt into setup.py * remove sphinx from build * remove create badges * Create CITATION.cff * Create .zenodo.json * Update transformer models * Update pint_definition.py * Update requirements.txt * create TimeDistributed CNN * small fixes to lodestar * Update layers.py * Update test_layers.py * remove direct getter of properties * Update scatterers.py Coherence length fix for MieScatterer * Update scatterers.py Added coherence length to the conversion table * mie phase fix * removed pydeepimagej from deps * Change loss input order of CGAN and PCGAN * Create dmdataset (dataset for graph-level regression tasks) * Update gnns/__init__.py * Add detection_linking_hela dataset * Update dmdataset.py * Create the regression_diffusion_landscape * Update scatterers.py CuPy fix for coherence length * Update test_scatterers.py Added a new method for testing MieSphere when coherence length parameter is provided. * Update augmentations.py * Update test_scatterers.py * Update test_scatterers.py * Create endothelial_vs dataset * Update layers.py * Update utils.py * Update docs link * Update README.md * version bump * version bump * Update README.md * Update README.md * Update graphs.py * Update test_generators.py * Update generators.py * fix test * Update vae.py * Bugfix in endothelial_vs dataset * Fix issue with repeated oneof-features (#169) * Update cell_migration_analysis.ipynb * Update features.py (#189) * added WAE (both MMD and GAN) (#185) * added deterministic wae_gan, slightly modified vae (input size)) * changes to VAE, WAE_GAN, and GAN Implemented proposed changes to GAN, VAE, WAE_GAN: - optimizers of WAE_GAN as input parameters by overriding the compile method. - included documentation of input parameters for VAE and WAE_GAN. - WAE_GAN: different learning rates for the autoencoder and the discriminator, 1e-3 and 5e-4, respectively as in the original paper. - included a unit test for WAE_GAN, VAE, and GAN. - removed @as_KerasModel from GAN. - formatted the code for compatibility. * Update test_models.py * fixed typo * extra fixes * WAE generalized allows WAE-GAN and WAE-MMD * fixes * check unit test * compatibility issue removed match for compatibility, replaced with if * test unit corrected assert in test unit for GAN, VAE. and WAE * fixed dimension in test_models * added compile for GAN and WAE in test_model * defined call in gan * fixed input size in GAN --------- Co-authored-by: BenjaminMidtvedt <41636530+BenjaminMidtvedt@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Jesús Pineda <jesus.pineda@physics.gu.se> Co-authored-by: Benjamin Midtvedt <benmid@student.chalmers.se> Co-authored-by: Jesús Pineda <36273229+JesusPinedaC@users.noreply.github.com> Co-authored-by: Ccx55 <ccx555@gmail.com> Co-authored-by: Harshith Bachimanchi <62615092+HarshithBachimanchi@users.noreply.github.com> Co-authored-by: gideon <gideon.jagenstedt@gmail.com> Co-authored-by: Benjamin Midtvedt <benjamin.midtvedt@physics.gu.se> Co-authored-by: Carlo <carlo.manzo@gmail.com>
diff --git a/deeptrack/features.py b/deeptrack/features.py
@@ -1637,7 +1637,7 @@ def _process_and_get(self, images, **kwargs):
                 p0[0] = np.max([p0[0], 0])
                 p0[1] = np.max([p0[1], 0])
 
-                p0 = p0.astype(np.int)
+                p0 = p0.astype(int)
 
                 output_slice = output[
                     p0[0] : p0[0] + labelarg.shape[0],
diff --git a/deeptrack/models/__init__.py b/deeptrack/models/__init__.py
@@ -7,6 +7,7 @@
 from .gans import *
 from .gnns import *
 from .vaes import *
+from .waes import *
 
 # from .mrcnn import *
 # from .yolov1 import *
diff --git a/deeptrack/models/gans/gan.py b/deeptrack/models/gans/gan.py
@@ -3,15 +3,24 @@
 
 layers = tf.keras.layers
 
-
-@as_KerasModel
 class GAN(tf.keras.Model):
+    """Generative Adversarial Network (GAN) model.
+
+    Parameters:
+    discriminator: keras model, optional
+        The discriminator network.
+    generator: keras model, optional
+        The generator network.
+    latent_dim: int, optional
+        Dimension of the latent space for random vectors.
+    """
+
     def __init__(self, discriminator=None, generator=None, latent_dim=128):
         super(GAN, self).__init__()
 
+        # Initialize discriminator and generator, or use default if not provided
         if discriminator is None:
             discriminator = self.default_discriminator()
-
         if generator is None:
             generator = self.default_generator()
 
@@ -21,9 +30,13 @@ def __init__(self, discriminator=None, generator=None, latent_dim=128):
 
     def compile(self, d_optimizer, g_optimizer, loss_fn):
         super(GAN, self).compile()
+
+        # Set optimizers and loss function for training
         self.d_optimizer = d_optimizer
         self.g_optimizer = g_optimizer
         self.loss_fn = loss_fn
+
+        # Define metrics to track during training
         self.d_loss_metric = tf.keras.metrics.Mean(name="d_loss")
         self.g_loss_metric = tf.keras.metrics.Mean(name="g_loss")
 
@@ -36,17 +49,18 @@ def train_step(self, real_images):
         batch_size = tf.shape(real_images)[0]
         random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
 
-        # Decode them to fake images
+        # Generate fake images using the generator
         generated_images = self.generator(random_latent_vectors)
 
-        # Combine them with real images
+        # Combine real and fake images
         combined_images = tf.concat([generated_images, real_images], axis=0)
 
-        # Assemble labels discriminating real from fake images
+        # Create labels for real and fake images
         labels = tf.concat(
             [tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0
         )
-        # Add random noise to the labels - important trick!
+
+        # Add random noise to labels to improve stability
         labels += 0.05 * tf.random.uniform(tf.shape(labels))
 
         # Train the discriminator
@@ -58,14 +72,13 @@ def train_step(self, real_images):
             zip(grads, self.discriminator.trainable_weights)
         )
 
-        # Sample random points in the latent space
+        # Generate new random latent vectors
         random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
 
-        # Assemble labels that say "all real images"
+        # Create labels indicating "all real images" for generator training
         misleading_labels = tf.zeros((batch_size, 1))
 
-        # Train the generator (note that we should *not* update the weights
-        # of the discriminator)!
+        # Train the generator while keeping discriminator weights fixed
         with tf.GradientTape() as tape:
             predictions = self.discriminator(self.generator(random_latent_vectors))
             g_loss = self.loss_fn(misleading_labels, predictions)
@@ -75,12 +88,19 @@ def train_step(self, real_images):
         # Update metrics
         self.d_loss_metric.update_state(d_loss)
         self.g_loss_metric.update_state(g_loss)
+
+        # Return updated loss metrics
         return {
             "d_loss": self.d_loss_metric.result(),
             "g_loss": self.g_loss_metric.result(),
         }
 
+    def call(self, inputs):
+        # Run generator
+        return self.generator(inputs)
+
     def default_generator(self, latent_dim=128):
+        # Define the default generator architecture
         return tf.keras.Sequential(
             [
                 tf.keras.Input(shape=(latent_dim,)),
@@ -98,6 +118,7 @@ def default_generator(self, latent_dim=128):
         )
 
     def default_discriminator(self):
+        # Define the default discriminator architecture
         return tf.keras.Sequential(
             [
                 tf.keras.Input(shape=(64, 64, 3)),
@@ -112,4 +133,4 @@ def default_discriminator(self):
                 layers.Dense(1, activation="sigmoid"),
             ],
             name="discriminator",
-        )
+        )
diff --git a/deeptrack/models/vaes/vae.py b/deeptrack/models/vaes/vae.py
@@ -1,28 +1,39 @@
 import tensorflow as tf
 from tensorflow.keras import layers
-
 from ..utils import as_KerasModel
 
 
-@as_KerasModel
 class VAE(tf.keras.Model):
+    """Variational Autoencoder (VAE) model.
+
+    Parameters:
+    encoder: keras model, optional
+        The encoder network.
+    decoder: keras model, optional
+        The decoder network.
+    latent_dim: int, optional
+        Dimension of the latent space.
+    """
+
     def __init__(self, encoder=None, decoder=None, latent_dim=2, **kwargs):
-        super().__init__(**kwargs)
+        super(VAE, self).__init__(**kwargs)
 
-        # Dimensionality of the latent space
+        # Define encoder latent dimension
         self.latent_dim = latent_dim
 
+        # Initialize encoder and decoder, or use defaults
         if encoder is None:
             self.encoder = self.default_encoder()
 
         if decoder is None:
             self.decoder = self.default_decoder()
 
     def train_step(self, data):
-
         data, _ = data
 
+        # Gradient tape for automatic differentiation
         with tf.GradientTape() as tape:
+            # Encode input data and sample from latent space.
             # The encoder outputs the mean and log of the variance of the
             # Gaussian distribution. The log of the variance is computed
             # instead of the variance for numerical stability.
@@ -32,42 +43,37 @@ def train_step(self, data):
             epsilon = tf.random.normal(shape=tf.shape(z_mean))
             z = z_mean + tf.exp(0.5 * z_log_var) * epsilon
 
-            # Reconstruct the input image
+            # Decode latent samples and compute reconstruction loss
             rdata = self.decoder(z)
-
-            # Reconstruction loss
             rloss = self.loss(data, rdata)
 
-            # KL divergence loss
-            kl_loss = -0.5 * (
-                1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
-            )
+            # Compute KL divergence loss
+            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
             kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
 
-            # Total loss
+            # Compute total loss
             loss = rloss + kl_loss
 
-        # Compute gradients
+        # Compute gradients and update model weights
         grads = tape.gradient(loss, self.trainable_weights)
+        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
 
-        # Update weights
-        self.optimizer.apply_gradients(
-            zip(grads, self.trainable_weights),
-        )
-
-        # Update metrics
+        # Update metrics for monitoring
         self.compiled_metrics.update_state(data, rdata)
 
+        # Return loss values for visualization
         return {
             "loss": loss,
             "reconstruction_loss": rloss,
             "kl_loss": kl_loss,
         }
 
     def call(self, inputs):
+        # Use encoder to obtain latent representation
         return self.encoder(inputs)
 
     def default_encoder(self):
+        # Define the default encoder architecture
         return tf.keras.Sequential(
             [
                 tf.keras.Input(shape=(28, 28, 1)),
@@ -88,14 +94,13 @@ def default_encoder(self):
                 layers.Flatten(),
                 layers.Dense(16),
                 layers.LeakyReLU(alpha=0.2),
-                layers.Dense(
-                    self.latent_dim + self.latent_dim, name="z_mean_log_var"
-                ),
+                layers.Dense(self.latent_dim + self.latent_dim, name="z_mean_log_var"),
             ],
             name="encoder",
         )
 
     def default_decoder(self):
+        # Define the default decoder architecture
         return tf.keras.Sequential(
             [
                 tf.keras.Input(shape=(self.latent_dim,)),
diff --git a/deeptrack/models/waes/__init__.py b/deeptrack/models/waes/__init__.py
@@ -0,0 +1 @@
+from .wae import *
diff --git a/deeptrack/models/waes/wae.py b/deeptrack/models/waes/wae.py
diff --git a/deeptrack/test/test_models.py b/deeptrack/test/test_models.py
