Zernike3Deep compute distance matrix: Fixes

README.md

@@ -28,7 +28,7 @@ This package provides a neural network environment integrating the deep learning
 
 The Flexutils-Toolkit relies on Conda and Pip for its correct installation.
 
-Before installing, Conda must be set in the PATH variable so it can be discovered during the installation. Once this requirement is met, the package can be either installed with ``pip install -e git+https://github.com/I2PC/Flexutils-Toolkit.git@master#egg=flexutils-toolkit`` or with ``pip install -e path/to/cloned/flexutils-toolkit`` after cloning this repository.
+Before installing, Conda must be set in the PATH variable so it can be discovered during the installation. Once this requirement is met, the package can be either installed with ``pip install -e git+https://github.com/I2PC/Flexutils-Toolkit.git@master#egg=tensorflow_toolkit`` or with ``pip install -e path/to/cloned/flexutils-toolkit`` after cloning this repository.
 
 We recommend adding the flag `-v` to pip installation command to have a better tracking of the installation.
 

tensorflow_toolkit/networks/zernike3deep.py

@@ -728,7 +728,7 @@ def predict_step(self, data):
             encoded = [self.z_space_x(x), self.z_space_y(x), self.z_space_z(x), self.delta_euler(x),
                        self.delta_shifts(x)]
         elif self.mode == "tomo":
-            x, latent = self.encoder(inputs)
+            x, latent = self.encoder_exp(inputs)
             encoded = [self.z_space_x(latent), self.z_space_y(latent), self.z_space_z(latent),
                        self.delta_euler(x),
                        self.delta_shifts(x)]

tensorflow_toolkit/scripts/compute_distance_matrix_zernike3deep.py

@@ -40,9 +40,9 @@
 
 
 def compute_distance_matrix(outPath, references_file, targets_file, L1, L2, batch_size, epochs, cost,
-                            architecture="convnn", lr=1e-5, jit_compile=True, regNorm=1e-4, gpu=0, thr=8):
+                            architecture="convnn", lr=1e-5, jit_compile=True, regNorm=1e-4, numProjections=20, gpu=0, thr=8):
 
-    # Load references and targers (MMap to save memory)
+    # Load references and targets (MMap to save memory)
     references = np.load(references_file, mmap_mode="r")
     targets = np.load(targets_file, mmap_mode="r")
 
@@ -57,9 +57,10 @@ def compute_distance_matrix(outPath, references_file, targets_file, L1, L2, batc
     volIds = np.repeat(np.arange(targets.shape[0]), 20)
     if not os.path.isfile(md_file):
         for volume in tqdm(targets, desc="Projecting volumes: "):
-            volume = np.reshape(volume, (target_bx, target_bx, target_bx))
+            if volume.ndim == 1:
+                volume = np.reshape(volume, (target_bx, target_bx, target_bx))
             volume = ih.scaleSplines(data=volume, finalDimension=64)
-            proj, angles = ih.generateProjections(20, volume=volume,
+            proj, angles = ih.generateProjections(numProjections, volume=volume,
                                                   n_jobs=thr, useFourier=True)
             proj_all.append(proj)
             angles_all.append(angles)
@@ -78,7 +79,8 @@ def compute_distance_matrix(outPath, references_file, targets_file, L1, L2, batc
     for volume in references:
         # Prepare data
         md.write(filename=md_file, overwrite=True)
-        volume = np.reshape(volume, (reference_bx, reference_bx, reference_bx))
+        if volume.ndim == 1:
+              volume = np.reshape(volume, (reference_bx, reference_bx, reference_bx))
         volume = ih.scaleSplines(data=volume, finalDimension=64)
         ImageHandler().write(volume, os.path.join(outPath, "volume.mrc"), overwrite=True)
         mask = ih.generateMask(inputFn=os.path.join(outPath, "volume.mrc"),
@@ -148,6 +150,7 @@ def main():
     parser.add_argument('--architecture', type=str, default="mlpnn")
     parser.add_argument('--jit_compile', action='store_true')
     parser.add_argument('--regNorm', type=float, default=0.001)
+    parser.add_argument('--num_projections', type=int, default=20)
     parser.add_argument('--gpu', type=str)
     parser.add_argument('--thr', type=int)
 
@@ -157,7 +160,7 @@ def main():
               "outPath": args.out_path, "L1": args.L1,
               "L2": args.L2, "batch_size": args.batch_size, "epochs": args.epochs,
               "cost": args.cost, "architecture": args.architecture,
-              "lr": args.lr, "jit_compile": args.jit_compile,
+              "lr": args.lr, "jit_compile": args.jit_compile, "numProjections": args.numProjections,
               "regNorm": args.regNorm, "gpu": args.gpu, "thr": args.thr}
 
     # Initialize volume slicer