Update neural_network.py to account for rat bug

Author: wjguan

DeepSlice/neural_network/neural_network.py

@@ -55,32 +55,40 @@ def initialise_network(xception_weights: str, weights: str, species: str) -> Seq
         model.add(Dense(9, activation="linear"))
 
     if weights != None:
-        model = load_xception_weights(model, weights)
+        model = load_xception_weights(model, weights, species)
     return model
 
 
-def load_xception_weights(model, weights):
+def load_xception_weights(model, weights, species = "mouse"):
     with h5py.File(weights, "r") as new:
         # set weight of each layer manually
-        model.layers[1].set_weights([new["dense"]["dense"]["kernel:0"], new["dense"]["dense"]["bias:0"]])
-        model.layers[2].set_weights([new["dense_1"]["dense_1"]["kernel:0"], new["dense_1"]["dense_1"]["bias:0"]])
-        model.layers[3].set_weights([new["dense_2"]["dense_2"]["kernel:0"], new["dense_2"]["dense_2"]["bias:0"]])
+        if species == "mouse":
+            xception_idx = 0
+            dense_idx = 1
+        elif species == "rat":
+            # RatModelInProgress.h5 has an "input_2" layer at index 0, so we need to adjust the indices<
+            xception_idx = 1
+            dense_idx = 2
+
+        model.layers[dense_idx].set_weights([new["dense"]["dense"]["kernel:0"], new["dense"]["dense"]["bias:0"]])
+        model.layers[dense_idx+1].set_weights([new["dense_1"]["dense_1"]["kernel:0"], new["dense_1"]["dense_1"]["bias:0"]])
+        model.layers[dense_idx+2].set_weights([new["dense_2"]["dense_2"]["kernel:0"], new["dense_2"]["dense_2"]["bias:0"]])
 
         # Set the weights of the xception model 
         weight_names = new["xception"].attrs["weight_names"].tolist()
         weight_names_layers = [name.decode("utf-8").split("/")[0] for name in weight_names]
 
-        for i in range(len(model.layers[0].layers)):
-            name_of_layer = model.layers[0].layers[i].name
+        for i in range(len(model.layers[xception_idx].layers)):
+            name_of_layer = model.layers[xception_idx].layers[i].name
             # if layer name is in the weight names, then we will set weights
             if name_of_layer in weight_names_layers:
                 # Get name of weights in the layer
                 layer_weight_names = []
-                for weight in model.layers[0].layers[i].weights:
+                for weight in model.layers[xception_idx].layers[i].weights:
                     layer_weight_names.append(weight.name.split("/")[1])
                 h5_group = new["xception"][name_of_layer]
                 weights_list = [np.array(h5_group[kk]) for kk in layer_weight_names]
-                model.layers[0].layers[i].set_weights(weights_list)
+                model.layers[xception_idx].layers[i].set_weights(weights_list)
     return model 
 
 def load_images_from_path(image_path: str) -> np.ndarray:
@@ -163,6 +171,7 @@ def predictions_util(
     primary_weights: str,
     secondary_weights: str,
     ensemble: bool = False,
+    species : str = "mouse"
 ):
     """
     Predict the image alignments
@@ -174,7 +183,7 @@ def predictions_util(
     :return: The predicted alignments
     :rtype: list
     """
-    model = load_xception_weights(model, primary_weights)
+    model = load_xception_weights(model, primary_weights, species)
     predictions = model.predict(
         image_generator,
         steps=image_generator.n // image_generator.batch_size,
@@ -183,14 +192,14 @@ def predictions_util(
     predictions = predictions.astype(np.float64)
     if ensemble:
         image_generator.reset()
-        model = load_xception_weights(model, secondary_weights)
+        model = load_xception_weights(model, secondary_weights, species)
         secondary_predictions = model.predict(
             image_generator,
             steps=image_generator.n // image_generator.batch_size,
             verbose=1,
         )
         predictions = np.mean([predictions, secondary_predictions], axis=0)
-        model = load_xception_weights(model, primary_weights)
+        model = load_xception_weights(model, primary_weights, species)
     filenames = image_generator.filenames
     filenames = [os.path.basename(i) for i in filenames]
     predictions_df = pd.DataFrame(