PyTorch ver 1.0: keras 3.+ with pytorch backend for nn_models.py and …

caiman/__init__.py

@@ -7,5 +7,6 @@
 from caiman.cluster import start_server, stop_server
 from caiman.mmapping import load_memmap, save_memmap, save_memmap_each, save_memmap_join
 from caiman.summary_images import local_correlations
+from caiman.pytorch_model_arch import PyTorchCNN
 
 __version__ = importlib.metadata.version('caiman')

caiman/components_evaluation.py

@@ -6,17 +6,17 @@
 import numpy as np
 import os
 import peakutils
-import tensorflow as tf
 import scipy
 from scipy.sparse import csc_matrix
 from scipy.stats import norm
+import torch
 from typing import Any, Union
 import warnings
 
 import caiman
 from caiman.paths import caiman_datadir
+from caiman.pytorch_model_arch import PyTorchCNN
 import caiman.utils.stats
-import caiman.utils.utils
 
 try:
     cv2.setNumThreads(0)
@@ -268,45 +268,22 @@ def evaluate_components_CNN(A,
         then this code will try not to use a GPU. Otherwise it will use one if it finds it.
     """
     logger = logging.getLogger("caiman")
-
-    # TODO: Find a less ugly way to do this
     if not isGPU and 'CAIMAN_ALLOW_GPU' not in os.environ:
         print("GPU run not requested, disabling use of GPUs")
         os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
-    try:
-        os.environ["KERAS_BACKEND"] = "tensorflow"
-        from tensorflow.keras.models import model_from_json
-        use_keras = True
-        logger.info('Using Keras')
-    except (ModuleNotFoundError):
-        use_keras = False
-        logger.info('Using Tensorflow')
 
-    if loaded_model is None:
-        if use_keras:
-            if os.path.isfile(os.path.join(caiman_datadir(), model_name + ".json")):
-                model_file = os.path.join(caiman_datadir(), model_name + ".json")
-                model_weights = os.path.join(caiman_datadir(), model_name + ".h5")
-            elif os.path.isfile(model_name + ".json"):
-                model_file = model_name + ".json"
-                model_weights = model_name + ".h5"
-            else:
-                raise FileNotFoundError(f"File for requested model {model_name} not found")
-            with open(model_file, 'r') as json_file:
-                print(f"USING MODEL (keras API): {model_file}")
-                loaded_model_json = json_file.read()
+    logger.info('Using Torch')
 
-            loaded_model = model_from_json(loaded_model_json)
-            loaded_model.load_weights(model_name + '.h5')
+    if loaded_model is None:
+        if os.path.isfile(os.path.join(caiman_datadir(), 'model', 'pytorch-models', model_name + ".pt")):
+            model_file = os.path.join(caiman_datadir(), 'model', 'pytorch-models', model_name + ".pt")
+        elif os.path.isfile(model_name + ".pt"):
+            model_file = model_name + ".pt"
         else:
-            if os.path.isfile(os.path.join(caiman_datadir(), model_name + ".h5.pb")):
-                model_file = os.path.join(caiman_datadir(), model_name + ".h5.pb")
-            elif os.path.isfile(model_name + ".h5.pb"):
-                model_file = model_name + ".h5.pb"
-            else:
-                raise FileNotFoundError(f"File for requested model {model_name} not found")
-            print(f"USING MODEL (tensorflow API): {model_file}")
-            loaded_model = caiman.utils.utils.load_graph(model_file)
+            raise FileNotFoundError(f"File for requested model {model_name} not found")
+        print(f"USING MODEL (torch API): {model_file}")
+        loaded_model = PyTorchCNN()
+        loaded_model.load_state_dict(torch.load(model_file))
 
         logger.debug("Loaded model from disk")
 
@@ -320,16 +297,15 @@ def evaluate_components_CNN(A,
                                               half_crop[1]:com[1] + half_crop[1]] for mm, com in zip(A.tocsc().T, coms)
     ]
     final_crops = np.array([cv2.resize(im / np.linalg.norm(im), (patch_size, patch_size)) for im in crop_imgs])
-    if use_keras:
-        predictions = loaded_model.predict(final_crops[:, :, :, np.newaxis], batch_size=32, verbose=1)
-    else:
-        tf_in = loaded_model.get_tensor_by_name('prefix/conv2d_20_input:0')
-        tf_out = loaded_model.get_tensor_by_name('prefix/output_node0:0')
-        with tf.Session(graph=loaded_model) as sess:
-            predictions = sess.run(tf_out, feed_dict={tf_in: final_crops[:, :, :, np.newaxis]})
-            sess.close()
-
-    return predictions, final_crops
+
+    final_crops = torch.tensor(final_crops, dtype=torch.float32)
+    final_crops = final_crops.unsqueeze(1)
+
+    with torch.no_grad():
+        predictions = loaded_model(final_crops)
+
+    predictions_numpy = predictions.cpu().numpy()
+    return predictions_numpy, final_crops
 
 def evaluate_components(Y: np.ndarray,
                         traces: np.ndarray,

caiman/pytorch_model_arch.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class PyTorchCNN(nn.Module):
+    def __init__(self):
+        super(PyTorchCNN, self).__init__()
+        # First convolutional block
+        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
+        self.conv2 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3)
+        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.dropout1 = nn.Dropout(p=0.25)
+
+        # Second convolutional block
+        self.conv3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding='same')
+        self.conv4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3)
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.dropout2 = nn.Dropout(p=0.25)
+
+        # Flattening and fully connected layers
+        self.flatten = nn.Flatten()
+        self.fc1 = nn.Linear(in_features=6400, out_features=512)
+        self.dropout3 = nn.Dropout(p=0.5)
+        self.fc2 = nn.Linear(in_features=512, out_features=2)
+
+    def forward(self, x):
+        # Convolutional Block 1 
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = self.pool1(x)
+        x = self.dropout1(x)
+
+        # Convolutional block 2 
+        x = F.relu(self.conv3(x))
+        x = F.relu(self.conv4(x))
+        x = self.pool2(x)
+        x = self.dropout2(x)
+
+        # Flattening and in_features layers
+        x = self.flatten(x)
+        x = F.relu(self.fc1(x))
+        x = self.dropout3(x)
+        x = self.fc2(x)
+        return F.softmax(x, dim=1)

caiman/source_extraction/cnmf/online_cnmf.py

@@ -26,7 +26,8 @@
 from scipy.stats import norm
 from sklearn.decomposition import NMF
 from sklearn.preprocessing import normalize
-import tensorflow as tf
+import torch
+from torch.utils.data import DataLoader, TensorDataset
 from time import time
 
 import caiman
@@ -38,6 +39,7 @@
                                   high_pass_filter_space, sliding_window,
                                   register_translation_3d, apply_shifts_dft)
 import caiman.paths
+from caiman.pytorch_model_arch import PyTorchCNN
 from caiman.source_extraction.cnmf.cnmf import CNMF
 from caiman.source_extraction.cnmf.estimates import Estimates
 from caiman.source_extraction.cnmf.initialization import imblur, initialize_components, hals, downscale
@@ -47,7 +49,7 @@
 from caiman.source_extraction.cnmf.utilities import (update_order, peak_local_max, decimation_matrix,
                         gaussian_filter, uniform_filter)
 import caiman.summary_images
-from caiman.utils.utils import save_dict_to_hdf5, load_dict_from_hdf5, parmap, load_graph
+from caiman.utils.utils import save_dict_to_hdf5, load_dict_from_hdf5, parmap
 from caiman.utils.stats import pd_solve
 from caiman.utils.nn_models import (fit_NL_model, create_LN_model, quantile_loss, rate_scheduler)
 
@@ -56,7 +58,6 @@
 except():
     pass
 
-#FIXME ???
 try:
     profile
 except:
@@ -122,6 +123,31 @@ def __init__(self, params=None, estimates=None, path=None, dview=None, Ain=None)
             # change things for code/notebooks that've worked for a long time; we should save such changes for a major rewrite
             # (if someone takes a particular interest in that).
 
+    def __str__(self):
+        ret = f"Caiman OnACID Object. subfields:{list(self.__dict__.keys()) }"
+        if hasattr(self.estimates, 'A') and self.estimates.A is not None:
+            ret += f" A.shape={self.estimates.A.shape}"
+        if hasattr(self.estimates, 'b') and self.estimates.b is not None:
+            ret += f" b.shape={self.estimates.b.shape}"
+        if hasattr(self.estimates, 'C') and self.estimates.C is not None:
+            ret += f" C.shape={self.estimates.C.shape}"
+        return ret
+
+    def __repr__(self):
+        ret = f"Caiman OnACID Object"
+        if hasattr(self.estimates, 'A') and self.estimates.A is not None:
+            ret += f" A.shape={self.estimates.A.shape}"
+        if hasattr(self.estimates, 'b') and self.estimates.b is not None and len(self.estimates.b.shape) > 1:
+            ret += f" bg components={self.estimates.b.shape[1]}"
+        if hasattr(self.estimates, 'C') and self.estimates.C is not None:
+            ret += f" C.shape={self.estimates.C.shape}"
+        ret += " Use str() for more details"
+        return ret
+
+    def __getitem__(self, idx):
+        return getattr(self, idx)
+    # We want subscripting to be read-only so we do not define a __setitem__ method
+
     @profile
     def _prepare_object(self, Yr, T, new_dims=None, idx_components=None):
 
@@ -155,7 +181,7 @@ def _prepare_object(self, Yr, T, new_dims=None, idx_components=None):
 
         if Yr.shape[-1] != self.params.get('online', 'init_batch'):
             raise Exception(
-                'The movie size used for initialization does not match with the minibatch size')
+                'The movie size used for initialization does not match the minibatch size')
 
         if new_dims is not None:
 
@@ -309,6 +335,7 @@ def _prepare_object(self, Yr, T, new_dims=None, idx_components=None):
             self.estimates.rho_buf = RingBuffer(self.estimates.rho_buf, self.params.get('online', 'minibatch_shape'))
             self.estimates.sv = np.sum(self.estimates.rho_buf.get_last_frames(
                 min(self.params.get('online', 'init_batch'), self.params.get('online', 'minibatch_shape')) - 1), 0)
+
         self.estimates.AtA = (self.estimates.Ab.T.dot(self.estimates.Ab)).toarray()
         self.estimates.AtY_buf = self.estimates.Ab.T.dot(self.estimates.Yr_buf.T)
         self.estimates.groups = list(map(list, update_order(self.estimates.Ab)[0]))
@@ -326,34 +353,13 @@ def _prepare_object(self, Yr, T, new_dims=None, idx_components=None):
         if self.params.get('online', 'path_to_model') is None or self.params.get('online', 'sniper_mode') is False:
             loaded_model = None
             self.params.set('online', {'sniper_mode': False})
-            self.tf_in = None
-            self.tf_out = None
         else:
-            try:
-                from tensorflow.keras.models import model_from_json
-                logger.info('Using Keras')
-                use_keras = True
-            except(ModuleNotFoundError):
-                use_keras = False
-                logger.info('Using Tensorflow')
-            if use_keras:
-                path = self.params.get('online', 'path_to_model').split(".")[:-1]
-                json_path = ".".join(path + ["json"])
-                model_path = ".".join(path + ["h5"])
-                json_file = open(json_path, 'r')
-                loaded_model_json = json_file.read()
-                json_file.close()
-                loaded_model = model_from_json(loaded_model_json)
-                loaded_model.load_weights(model_path)
-                self.tf_in = None
-                self.tf_out = None
-            else:
-                path = self.params.get('online', 'path_to_model').split(".")[:-1]
-                model_path = '.'.join(path + ['h5', 'pb'])
-                loaded_model = load_graph(model_path)
-                self.tf_in = loaded_model.get_tensor_by_name('prefix/conv2d_1_input:0')
-                self.tf_out = loaded_model.get_tensor_by_name('prefix/output_node0:0')
-                loaded_model = tf.Session(graph=loaded_model)
+            logger.info('Using Torch')
+            path = self.params.get('online', 'path_to_model').split(".")[:-1]
+            model_path = '.'.join(path + ['pt'])
+            loaded_model = PyTorchCNN()
+            loaded_model.load_state_dict(torch.load(model_path))
+
         self.loaded_model = loaded_model
 
         if self.is1p:
@@ -555,7 +561,6 @@ def fit_next(self, t, frame_in, num_iters_hals=3):
                 sniper_mode=self.params.get('online', 'sniper_mode'),
                 use_peak_max=self.params.get('online', 'use_peak_max'),
                 mean_buff=self.estimates.mean_buff,
-                tf_in=self.tf_in, tf_out=self.tf_out,
                 ssub_B=ssub_B, W=self.estimates.W if self.is1p else None,
                 b0=self.estimates.b0 if self.is1p else None,
                 corr_img=self.estimates.corr_img if use_corr else None,
@@ -1004,7 +1009,7 @@ def initialize_online(self, model_LN=None, T=None):
             self.bnd_AC = np.percentile(np.ravel(self.estimates.A.dot(self.estimates.C)),
                                         (0.001, 100-0.005))
         return self
-
+  
     def save(self,filename):
         """save object in hdf5 file format
 
@@ -1155,6 +1160,7 @@ def fit_online(self, **kwargs):
                 self.params.set('ring_CNN', {'path_to_model': path_to_model})
         else:
             model_LN = None
+
         epochs = self.params.get('online', 'epochs')
         self.initialize_online(model_LN=model_LN)
         self.t_init += time()
@@ -1221,6 +1227,13 @@ def fit_online(self, **kwargs):
                                          + str(self.estimates.Ab.shape[-1] - self.params.get('init', 'nb')))
                             old_comps = self.N
 
+                        if np.isnan(np.sum(frame)):
+                            raise Exception('Frame {frame_count} contains NaN')
+                        if t % 500 == 0:
+                            logger.info(f'Epoch: {iter + 1}. {t} frames have been processed.'
+                                         f'{self.N - old_comps} new components were added. Total: {self.N}')
+                            old_comps = self.N
+
                         # Downsample and normalize
                         frame_ = frame.copy().astype(np.float32)
                         if self.params.get('online', 'ds_factor') > 1:
@@ -1299,9 +1312,7 @@ def fit_online(self, **kwargs):
             cv2.destroyAllWindows()
         self.t_online = t_online
         self.estimates.C_on = self.estimates.C_on[:self.M]
-        self.estimates.noisyC = self.estimates.noisyC[:self.M]
-
-        return self
+        self.estimates.noisyC = self.estimates.noisyC[:self.M]  
 
     def create_frame(self, frame_cor, show_residuals=True, resize_fact=3, transpose=True):
         if show_residuals:
@@ -2009,8 +2020,7 @@ def get_candidate_components(sv, dims, Yres_buf, min_num_trial=3, gSig=(5, 5),
                              gHalf=(5, 5), sniper_mode=True, rval_thr=0.85,
                              patch_size=50, loaded_model=None, test_both=False,
                              thresh_CNN_noisy=0.5, use_peak_max=False,
-                             thresh_std_peak_resid = 1, mean_buff=None,
-                             tf_in=None, tf_out=None):
+                             thresh_std_peak_resid = 1, mean_buff=None):
     """
     Extract new candidate components from the residual buffer and test them
     using space correlation or the CNN classifier. The function runs the CNN
@@ -2091,11 +2101,23 @@ def get_candidate_components(sv, dims, Yres_buf, min_num_trial=3, gSig=(5, 5),
         Ain2 /= np.std(Ain2,axis=1)[:,None]
         Ain2 = np.reshape(Ain2,(-1,) + tuple(np.diff(ijSig_cnn).squeeze()),order= 'F')
         Ain2 = np.stack([cv2.resize(ain,(patch_size ,patch_size)) for ain in Ain2])
-        if tf_in is None:
-            predictions = loaded_model.predict(Ain2[:,:,:,np.newaxis], batch_size=min_num_trial, verbose=0)
-        else:
-            predictions = loaded_model.run(tf_out, feed_dict={tf_in: Ain2[:, :, :, np.newaxis]})
-        keep_cnn = list(np.where(predictions[:, 0] > thresh_CNN_noisy)[0])
+
+        final_crops = Ain2[:, :, :, np.newaxis]
+        final_crops_tensor = torch.tensor(final_crops, dtype=torch.float32).permute(0, 3, 1, 2)
+
+        #Create DataLoader for batching 
+        dataset = TensorDataset(final_crops_tensor)
+        loader = DataLoader(dataset, batch_size=int(min_num_trial), shuffle=False)
+
+        loaded_model.eval()
+        all_predictions = []
+        with torch.no_grad():
+            for batch in loader:
+                outputs = loaded_model(batch[0])
+                all_predictions.append(outputs)   
+
+        predictions = torch.cat(all_predictions).cpu().numpy()
+        keep_cnn = list(np.where(predictions[:,0] > thresh_CNN_noisy)[0])
         cnn_pos = Ain2[keep_cnn]
     else:
         keep_cnn = []  # list(range(len(Ain_cnn)))
@@ -2144,8 +2166,7 @@ def update_num_components(t, sv, Ab, Cf, Yres_buf, Y_buf, rho_buf,
                           mean_buff=None, ssub_B=1, W=None, b0=None,
                           corr_img=None, first_moment=None, second_moment=None,
                           crosscorr=None, col_ind=None, row_ind=None, corr_img_mode=None,
-                          max_img=None, downscale_matrix=None, upscale_matrix=None,
-                          tf_in=None, tf_out=None):
+                          max_img=None, downscale_matrix=None, upscale_matrix=None):
     """
     Checks for new components in the residual buffer and incorporates them if they pass the acceptance tests
     """
@@ -2174,8 +2195,7 @@ def update_num_components(t, sv, Ab, Cf, Yres_buf, Y_buf, rho_buf,
         min_num_trial=min_num_trial, gSig=gSig, gHalf=gHalf,
         sniper_mode=sniper_mode, rval_thr=rval_thr, patch_size=50,
         loaded_model=loaded_model, thresh_CNN_noisy=thresh_CNN_noisy,
-        use_peak_max=use_peak_max, test_both=test_both, mean_buff=mean_buff,
-        tf_in=tf_in, tf_out=tf_out)
+        use_peak_max=use_peak_max, test_both=test_both, mean_buff=mean_buff)
 
     ind_new_all = ijsig_all
 
@@ -2551,4 +2571,4 @@ def load_OnlineCNMF(filename, dview = None):
     return new_obj
 
 def inv_mat_vec(A):
-    return np.linalg.solve(A[0], A[1])
+    return np.linalg.solve(A[0], A[1])