Merge master in devel

.gitignore

@@ -7,3 +7,6 @@ toolkit/*.egg-info
 
 # IDE files
 .idea
+
+# Installation files
+flexutils_tensorflow_installed

install_open3d.sh

@@ -72,11 +72,11 @@ done
 
 # Inform the user about missing packages
 if [ ${#missing_packages[@]} -ne 0 ]; then
-    colored_echo "red" "The following packages are missing:"
+    colored_echo "yellow" "The following packages are missing:"
     for pkg in "${missing_packages[@]}"; do
-        colored_echo "red" " - $pkg"
+        colored_echo "yellow" " - $pkg"
     done
-    colored_echo "red" "Open3D functionalities will not be available. If you want them to be used, please, install the listed
+    colored_echo "yellow" "Open3D functionalities will not be available. If you want them to be used, please, install the listed
     packages as sudo and rerun the scipion-em-flexutils plugin installation."
     exit 0
 fi
@@ -122,9 +122,19 @@ colored_echo "green" "##### Done! #####"
 # Check Cuda is installed in the system
 colored_echo "green" "##### Checking Cuda... #####"
 if command -v nvcc > /dev/null 2>&1; then
-    colored_echo "green" "Cuda not installed in the system. Please, install Cuda."
+    cuda_version=$(nvcc --version | grep "release" | awk '{print $6}' | cut -d',' -f1)
+    cleaned_version=$(echo $cuda_version | sed 's/V//;s/\([0-9]*\.[0-9]*\).*/\1/')
+    cuda_version_number=$(echo $cleaned_version | awk -F. '{printf "%d%02d", $1, $2}')
+    if [ "$cuda_version_number" -ge 1202 ]; then
+        colored_echo "green" "Cuda found in the system."
+    else
+        colored_echo "yellow" "Cuda has been found in your system, but does not fulfill versio requirements. Currently,
+        Open3D can only be installed if the Cuda version in the system is at least 12.2. If you want to use Open3D,
+        please update your Cuda version."
+        exit 0
+
 else
-    colored_echo "red" "CUDA not found, exiting. To installed Open3D capabilities, please, install Cuda in your system
+    colored_echo "yellow" "CUDA not found, exiting. To installed Open3D capabilities, please, install Cuda in your system
     and retry the installation. If Cuda is already installed and you are seeing this message, you might need to
     manually add Cuda to the bashrc file so it can be found."
     exit 0

tensorflow_toolkit/build.py

@@ -127,7 +127,7 @@ def runCondaInstallation(self):
                             "Install Conda and/or add it to the PATH variable and try to install again "
                             "this package with 'pip install tensorflow-toolkit'")
 
-        req_file, condabin_path, install_conda_command, cuda_version, tensorlfow = self.condaInstallationCommands()
+        req_file, condabin_path, install_conda_command, cuda_version, tensorflow = self.condaInstallationCommands()
 
         # Install flexutils-tensorflow conda environment
         self.print_flush("Installing Tensorflow conda env...")
@@ -140,7 +140,7 @@ def runCondaInstallation(self):
         # Get command to install envrionment and pip dependencies
         self.print_flush("Getting env pip...")
         if install_conda_command is not None:
-            if tensorlfow == "2.12":
+            if tensorflow == "2.12":
                 install_toolkit_command = 'eval "$(%s shell.bash hook)" && conda activate flexutils-tensorflow && ' \
                                           'conda install -y -c nvidia cuda-nvcc=11.3.58 && ' \
                                           'pip install nvidia-cudnn-cu11==8.6.0.163 && ' \
@@ -163,17 +163,17 @@ def runCondaInstallation(self):
 
         # Set Tensorflow env variables when env is activated
         self.print_flush("Set environment variables in conda env...")
-        if tensorlfow == "2.15" or "None":
+        if tensorflow == "2.15" or tensorflow == "None":
             commands = []
-        elif tensorlfow == "2.12":
+        elif tensorflow == "2.12":
             commands = ['eval "$(%s shell.bash hook) "' % condabin_path,
                         'conda activate flexutils-tensorflow ',
                         'mkdir -p $CONDA_PREFIX/etc/conda/activate.d ',
                         'echo \'CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))\''
                         ' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh ',
-                        'echo \'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib\\n\' '
+                        'echo \'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib\' '
                         '>> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh ',
-                        'echo \'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/lib/\\n\' >> '
+                        'echo \'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/lib/\' >> '
                         '$CONDA_PREFIX/etc/conda/activate.d/env_vars.sh ',
                         'mkdir -p $CONDA_PREFIX/lib/nvvm/libdevice ',
                         'cp $CONDA_PREFIX/lib/libdevice.10.bc $CONDA_PREFIX/lib/nvvm/libdevice/'
@@ -184,7 +184,7 @@ def runCondaInstallation(self):
                         'mkdir -p $CONDA_PREFIX/etc/conda/activate.d ',
                         'echo \'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/\' '
                         '>> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh ',
-                        'echo \'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/lib/\\n\' >> '
+                        'echo \'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/lib/\' >> '
                         '$CONDA_PREFIX/etc/conda/activate.d/env_vars.sh ',
                         'mkdir -p $CONDA_PREFIX/lib/nvvm/libdevice ',
                         'cp $CONDA_PREFIX/lib/libdevice.10.bc $CONDA_PREFIX/lib/nvvm/libdevice/'

tensorflow_toolkit/networks/zernike3deep.py

@@ -320,7 +320,7 @@ def __call__(self, x):
 
 class AutoEncoder(tf.keras.Model):
     def __init__(self, generator, architecture="convnn", CTF="apply", mode=None, l_bond=0.01, l_angle=0.01,
-                 l_clashes=None, jit_compile=True, **kwargs):
+                 l_clashes=None, l_norm=1e-4, jit_compile=True, **kwargs):
         super(AutoEncoder, self).__init__(**kwargs)
         self.generator = generator
         self.CTF = CTF
@@ -329,6 +329,8 @@ def __init__(self, generator, architecture="convnn", CTF="apply", mode=None, l_b
         self.l_bond = l_bond
         self.l_angle = l_angle
         self.l_clashes = l_clashes if l_clashes is not None else 0.0
+        self.l_norm = l_norm
+        # self.architecture = architecture
         self.encoder = Encoder(generator.zernike_size.shape[0], generator.xsize,
                                generator.refinePose, architecture=architecture,
                                mode=self.mode, jit_compile=jit_compile)
@@ -338,6 +340,7 @@ def __init__(self, generator, architecture="convnn", CTF="apply", mode=None, l_b
         self.bond_loss_tracker = tf.keras.metrics.Mean(name="bond_loss")
         self.angle_loss_tracker = tf.keras.metrics.Mean(name="angle_loss")
         self.clash_loss_tracker = tf.keras.metrics.Mean(name="clash_loss")
+        self.norm_loss_tracker = tf.keras.metrics.Mean(name="norm_loss")
 
         # XLA compilation of cost function
         if jit_compile:
@@ -370,7 +373,8 @@ def metrics(self):
             self.img_loss_tracker,
             self.bond_loss_tracker,
             self.angle_loss_tracker,
-            self.clash_loss_tracker
+            self.clash_loss_tracker,
+            self.norm_loss_tracker
         ]
 
     def train_step(self, data):
@@ -452,7 +456,20 @@ def train_step(self, data):
             else:
                 bond_loss, angle_loss = tf.constant(0.0, tf.float32), tf.constant(0.0, tf.float32)
 
-            total_loss = img_loss + self.l_bond * bond_loss + self.l_angle * angle_loss + self.l_clashes * clashes
+            # Coefficent norm loss
+            # if self.architecture == "deepconv":
+            #     norm_x = tf.reduce_mean(tf.square(encoded[0]), axis=-1)
+            #     norm_y = tf.reduce_mean(tf.square(encoded[1]), axis=-1)
+            #     norm_z = tf.reduce_mean(tf.square(encoded[2]), axis=-1)
+            #     norm_loss = (norm_x + norm_y + norm_z) / 3.
+            # else:
+            #     norm_loss = tf.constant(0.0, tf.float32)
+            norm_x = tf.reduce_mean(tf.square(encoded[0]), axis=-1)
+            norm_y = tf.reduce_mean(tf.square(encoded[1]), axis=-1)
+            norm_z = tf.reduce_mean(tf.square(encoded[2]), axis=-1)
+            norm_loss = (norm_x + norm_y + norm_z) / 3.
+
+            total_loss = img_loss + self.l_bond * bond_loss + self.l_angle * angle_loss + self.l_clashes * clashes + self.l_norm * norm_loss
 
         grads = tape.gradient(total_loss, self.trainable_weights)
         self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
@@ -461,12 +478,14 @@ def train_step(self, data):
         self.angle_loss_tracker.update_state(angle_loss)
         self.bond_loss_tracker.update_state(bond_loss)
         self.clash_loss_tracker.update_state(clashes)
+        self.norm_loss_tracker.update_state(norm_loss)
         return {
             "loss": self.total_loss_tracker.result(),
             "img_loss": self.img_loss_tracker.result(),
             "bond": self.bond_loss_tracker.result(),
             "angle": self.angle_loss_tracker.result(),
             "clashes": self.clash_loss_tracker.result(),
+            "norm": self.norm_loss_tracker.result(),
         }
 
     def test_step(self, data):

tensorflow_toolkit/scripts/predict_zernike3deep.py

@@ -44,6 +44,8 @@
 def predict(md_file, weigths_file, L1, L2, refinePose, architecture, ctfType, pad=2,
             sr=1.0, applyCTF=1):
 
+    # We need to import network and generators here instead of at the beginning of the script to allow Tensorflow
+    # get the right GPUs set in CUDA_VISIBLE_DEVICES
     from tensorflow_toolkit.generators.generator_zernike3deep import Generator
     from tensorflow_toolkit.networks.zernike3deep import AutoEncoder
 

tensorflow_toolkit/scripts/train_zernike3deep.py

@@ -45,10 +45,10 @@
 
 def train(outPath, md_file, L1, L2, batch_size, shuffle, step, splitTrain, epochs, cost,
           radius_mask, smooth_mask, refinePose, architecture="convnn", ctfType="apply", pad=2,
-          sr=1.0, applyCTF=1, lr=1e-5, jit_compile=True, regBond=0.01, regAngle=0.01, regClashes=None,
+          sr=1.0, applyCTF=1, lr=1e-5, jit_compile=True, regNorm=1e-4, regBond=0.01, regAngle=0.01, regClashes=None,
           tensorboard=True, weigths_file=None):
 
-    # We need to import network and generators here instead of at the begining of the script to allow Tensorflow
+    # We need to import network and generators here instead of at the beginning of the script to allow Tensorflow
     # get the right GPUs set in CUDA_VISIBLE_DEVICES
     from tensorflow_toolkit.generators.generator_zernike3deep import Generator
     from tensorflow_toolkit.networks.zernike3deep import AutoEncoder
@@ -82,7 +82,7 @@ def train(outPath, md_file, L1, L2, batch_size, shuffle, step, splitTrain, epoch
             jit_compile = False
         else:
             autoencoder = AutoEncoder(generator, architecture=architecture, CTF=ctfType, l_bond=regBond,
-                                      l_angle=regAngle, l_clashes=regClashes, jit_compile=False)
+                                      l_angle=regAngle, l_clashes=regClashes, l_norm=regNorm, jit_compile=False)
 
         # Fine tune a previous model
         if weigths_file:
@@ -183,6 +183,7 @@ def main():
     parser.add_argument('--sr', type=float, required=True)
     parser.add_argument('--apply_ctf', type=int, required=True)
     parser.add_argument('--jit_compile', action='store_true')
+    parser.add_argument('--regNorm', type=float, default=0.0001)
     parser.add_argument('--regBond', type=float, default=0.01)
     parser.add_argument('--regAngle', type=float, default=0.01)
     parser.add_argument('--regClashes', type=float, default=None)
@@ -213,8 +214,8 @@ def main():
               "refinePose": args.refine_pose, "architecture": args.architecture,
               "ctfType": args.ctf_type, "pad": args.pad, "sr": args.sr,
               "applyCTF": args.apply_ctf, "lr": args.lr, "jit_compile": args.jit_compile,
-              "regBond": args.regBond, "regAngle": args.regAngle, "regClashes": args.regClashes,
-              "tensorboard": args.tensorboard, "weigths_file": args.weigths_file}
+              "regNorm": args.regNorm, "regBond": args.regBond, "regAngle": args.regAngle,
+              "regClashes": args.regClashes, "tensorboard": args.tensorboard, "weigths_file": args.weigths_file}
 
     # Initialize volume slicer
     train(**inputs)