fixing preprocess imports

defs.py

@@ -0,0 +1,213 @@
+import os
+from enum import Enum
+from pprint import pprint
+from pathlib import Path
+
+
+
+def initDataPaths(project_dir=os.path.dirname(os.path.dirname(__file__))):
+    # initializes datapaths
+
+    global PROJECT_DIR
+    global METADATA_DIR
+    global IMAGE_DIRECTORIES
+    global TRAINIG_DATA_DIR_254_IMG_ORIENTATION
+    global TRULY_ROTATED_IMG_224
+    global TRAINIG_DATA_DIR_256_MASKED_IMBALANCED
+    global TRAINIG_DATA_DIR_256_MASKED_BALANCED
+    global MLRUNS_URI
+    global MLRUNS_DIR
+    global MODELS_DIR
+    global EXTERNAL_RAW_DEMO_DATA_DIR
+    global EXTERNAL_MASKED_DEMO_DATA_DIR
+
+    PROJECT_DIR = project_dir
+    MODELS_DIR = os.path.join(PROJECT_DIR, "models")
+    METADATA_DIR = os.path.join(PROJECT_DIR, r"metadata")
+    TRAINIG_DATA_DIR_254_IMG_ORIENTATION = os.path.join(
+        PROJECT_DIR, r"data_224x224\train_val_224x224"
+    )
+    TRULY_ROTATED_IMG_224 = os.path.join(PROJECT_DIR, r"224x224_truly_rotated")
+    TRAINIG_DATA_DIR_256_MASKED_IMBALANCED = os.path.join(
+        PROJECT_DIR, r"256x256_masked_images_imbalanced"
+    )
+    TRAINIG_DATA_DIR_256_MASKED_BALANCED = os.path.join(
+        PROJECT_DIR, r"256x256_masked_images_balanced"
+    )
+
+    MLRUNS_URI = Path(os.path.abspath(os.path.join(PROJECT_DIR, "mlruns_vst"))).as_uri()
+    MLRUNS_DIR = os.path.abspath(os.path.join(PROJECT_DIR, "mlruns_vst"))
+    
+    EXTERNAL_RAW_DEMO_DATA_DIR = os.path.join(PROJECT_DIR, "src", "streamlit", "data", "data_for_product_demo", "unlabeled", "external", "external_raw_299x299")
+    EXTERNAL_MASKED_DEMO_DATA_DIR = os.path.join(PROJECT_DIR, "src", "streamlit", "data", "data_for_product_demo", "unlabeled", "external", "external_masked_256x256")
+
+    IMAGE_DIRECTORIES = {
+        "COVID": {
+            "images": os.path.join(
+                PROJECT_DIR, "data", "COVID-19_Radiography_Dataset", "COVID", "images"
+            ),
+            "masks": os.path.join(
+                PROJECT_DIR, "data", "COVID-19_Radiography_Dataset", "COVID", "masks"
+            ),
+        },
+        "Lung_Opacity": {
+            "images": os.path.join(
+                PROJECT_DIR,
+                "data",
+                "COVID-19_Radiography_Dataset",
+                "Lung_Opacity",
+                "images",
+            ),
+            "masks": os.path.join(
+                PROJECT_DIR,
+                "data",
+                "COVID-19_Radiography_Dataset",
+                "Lung_Opacity",
+                "masks",
+            ),
+        },
+        "Normal": {
+            "images": os.path.join(
+                PROJECT_DIR, "data", "COVID-19_Radiography_Dataset", "Normal", "images"
+            ),
+            "masks": os.path.join(
+                PROJECT_DIR, "data", "COVID-19_Radiography_Dataset", "Normal", "masks"
+            ),
+        },
+        "Viral Pneumonia": {
+            "images": os.path.join(
+                PROJECT_DIR,
+                "data",
+                "COVID-19_Radiography_Dataset",
+                "Viral Pneumonia",
+                "images",
+            ),
+            "masks": os.path.join(
+                PROJECT_DIR,
+                "data",
+                "COVID-19_Radiography_Dataset",
+                "Viral Pneumonia",
+                "masks",
+            ),
+        },
+    }
+
+
+def checkPaths():
+    print(
+        "\nPROJECT_DIR: ",
+        PROJECT_DIR,
+        "\nMETADATA_DIR: ",
+        METADATA_DIR,
+        "\nIMAGE_DIRECTORIES: ",
+        IMAGE_DIRECTORIES,
+        "\nTRAINIG_DATA_DIR_254_IMG_ORIENTATION: ",
+        TRAINIG_DATA_DIR_254_IMG_ORIENTATION,
+        "\nTRULY_ROTATED_IMG_224: ",
+        TRULY_ROTATED_IMG_224,
+        "\nTRAINIG_DATA_DIR_256_MASKED_IMBALANCED: ",
+        TRAINIG_DATA_DIR_256_MASKED_IMBALANCED,
+        "\nTRAINIG_DATA_DIR_256_MASKED_BALANCED: ",
+        TRAINIG_DATA_DIR_256_MASKED_BALANCED,
+        "\nMLRUNS_URI: ",
+        MLRUNS_URI,
+        "\nMODELS_DIR: ",
+        MODELS_DIR,
+    )
+    
+    
+#----setting paths----
+initDataPaths()
+#---and checking them----
+checkPaths()
+    
+
+
+class _Base(str, Enum):
+    def __str__(self):
+        return self.value
+
+
+class ModelPath(_Base):
+    CAPSNET = os.path.join(
+        MODELS_DIR, "capsnet-4class-disease-classifier", "model.keras"
+    )
+    COVID19 = os.path.join(MODELS_DIR, "ds-crx-covid19", "model.keras")
+    GAN = os.path.join(MODELS_DIR, "lung-segmentation-gan", "model.keras")
+    UNET = os.path.join(MODELS_DIR, "lung-segmentation-unet", "model.keras")
+    MOBNET = os.path.join(
+        MODELS_DIR, "orientation-classifier-224x224-aug-head1-mobnet", "model.keras"
+    )
+    RESNET = os.path.join(
+        MODELS_DIR, "orientation-classifier-224x224-aug-head2-resnet50", "model.keras"
+    )
+
+
+class DiseaseCategory(_Base):
+    # Enum for the different disease categories
+    # alligned to file names without extension .png
+    VIRAL_PNEUMONIA = "Viral Pneumonia"
+    COVID = "COVID"
+    LUNG_OPACITY = "Lung_Opacity"
+    NORMAL = "Normal"
+
+
+class ImageType(_Base):
+    IMAGES = "images"
+    MASKS = "masks"
+    MASKED = "masked"
+
+
+class ModelType(_Base):
+    # Enum for the different model types
+    RESNET50 = "resnet50"
+    MOBILENET = "mobnet"
+    GAN = "gan"
+    UNET = "unet"
+    CUST_COVID_CNN = "cust_covid_cnn"
+    CAPSNET = "capsnet"
+
+
+class ExperimentName(_Base):
+    # mlflow experiment names
+    ORIENTATION_CLASSIFIER = "orientation_classifier"
+    DESEASE_CLASSIFIER = "desease_classifier"
+
+
+# >>>>>IMPORTANT: the mapping must be the same as for the training dataset!!!!<<<<<
+# check loaded dataset
+class_to_orientation_map = {
+    "long": {0: "rotated_0", 1: "rotated_180", 2: "rotated_90", 3: "rotated_minus_90"},
+    "short": {
+        0: "0°",
+        1: "180°",
+        2: "90°",
+        3: "-90°",
+    },
+}
+
+orientation_labels = {
+    "short": [
+        "0°",
+        "180°",
+        "90°",
+        "-90°",
+    ],
+    "long": ["rotated_0", "rotated_180", "rotated_90", "rotated_minus_90"],
+}
+
+class_to_disease_map = {
+    0: "COVID",
+    1: "Lung_Opacity",
+    2: "Normal",
+    3: "Viral Pneumonia",
+}
+
+disease_labels = ["COVID", "Lung_Opacity", "Normal", "Viral Pneumonia"]
+
+
+class DatasetType(_Base):
+    TRAIN = "train"
+    TEST = "test"
+    PREDICT = "predict"
+

modelbuilder.py

@@ -6,21 +6,17 @@ Class to construct the different type of models
 import tensorflow as tf
 from tensorflow import keras
 from tensorflow.keras import layers, Sequential
-from tensorflow.keras.layers import (
-    Dense,
-    Input,
-    Rescaling
-)
+from tensorflow.keras.layers import Dense, Input, Rescaling
 from tensorflow.keras.applications import MobileNet, ResNet50
 
 # --- CapsNet-specific
 from keras.saving import register_keras_serializable  # For custom layer serialization
 
 # --- Project-specific
-from src.defs import ModelType as mt
+from defs import ModelType as mt
 
 
-class ModelBuilder():
+class ModelBuilder:
     # builds the models
 
     def __init__(self, model_type, **model_params):
@@ -32,40 +28,49 @@ class ModelBuilder():
 
         # config extractor and attributes adding by model type
         if self.model_type in (mt.MOBILENET, mt.RESNET50):
-                self.base_model_params = self.model_params.pop("base_model")
-                self.model_name = self.base_model_params["name"]
-                self.input_shape = self.base_model_params["input_shape"]
-                self.base_trainable = self.model_params.pop("base_trainable")
-                self.base_model = None
-                
+            self.base_model_params = self.model_params.pop("base_model")
+            self.model_name = self.base_model_params["name"]
+            self.input_shape = self.base_model_params["input_shape"]
+            self.base_trainable = self.model_params.pop("base_trainable")
+            self.base_model = None
+
         elif self.model_type == mt.CAPSNET:
-                self.model_name = model_params.pop("name")
-                self.input_shape = model_params.pop("input_shape")
-                self.prim_caps_params = model_params.pop("prim_caps")
-                self.digit_caps_params = model_params.pop("digit_caps")
-                self.routing_algo = model_params.pop("routing_algo") # informative only
-                
+            self.model_name = model_params.pop("name")
+            self.input_shape = model_params.pop("input_shape")
+            self.prim_caps_params = model_params.pop("prim_caps")
+            self.digit_caps_params = model_params.pop("digit_caps")
+            self.routing_algo = model_params.pop("routing_algo")  # informative only
+
         # model_type vs input shape validation
-        if self.model_type in (mt.MOBILENET, mt.RESNET50,):
-            if self.input_shape != (224,224,3):
-                raise Exception(f"input shape for {self.model_name} model must be (224,224,3)") 
+        if self.model_type in (
+            mt.MOBILENET,
+            mt.RESNET50,
+        ):
+            if self.input_shape != (224, 224, 3):
+                raise Exception(
+                    f"input shape for {self.model_name} model must be (224,224,3)"
+                )
         elif self.model_type == mt.CAPSNET:
-            if self.input_shape != (256,256,3):
-                raise Exception(f"input shape for {self.model_name} model must be (256,256,3)")
+            if self.input_shape != (256, 256, 3):
+                raise Exception(
+                    f"input shape for {self.model_name} model must be (256,256,3)"
+                )
         else:
-            raise Exception(f"Model not supported: {self.model_name}. The model name must contain one substring from {mt.MOBILENET, mt.RESNET50, mt.CAPSNET}")
-        
-        
-        
+            raise Exception(
+                f"Model not supported: {self.model_name}. The model name must contain one substring from {mt.MOBILENET, mt.RESNET50, mt.CAPSNET}"
+            )
+
     def get_augmentation_pipe(self):
         # Random-* layers are stochastic only when training=True
         # disabled during inference/evaluation
-        return Sequential([
-            layers.RandomRotation(0.1),
-            layers.RandomTranslation(height_factor=0.1, width_factor=0.1),
-            layers.RandomZoom(0.1),
-        ], name="augmentation")
-                    
+        return Sequential(
+            [
+                layers.RandomRotation(0.1),
+                layers.RandomTranslation(height_factor=0.1, width_factor=0.1),
+                layers.RandomZoom(0.1),
+            ],
+            name="augmentation",
+        )
 
     def get_compiled_model(self):
         # Extract config
@@ -74,9 +79,11 @@ class ModelBuilder():
         # Define input layer
         inputs = Input(shape=self.input_shape, name="inputs")
         # Random-* layers are stochastic only when training=True
-        
-        x_aug  = self.get_augmentation_pipe()(inputs)      # stochastic only when training=True
-        x = Rescaling(1./255)(x_aug)                       # disabled during inference/evaluation
+
+        x_aug = self.get_augmentation_pipe()(
+            inputs
+        )  # stochastic only when training=True
+        x = Rescaling(1.0 / 255)(x_aug)  # disabled during inference/evaluation
 
         # Model selector
         match self.model_type:
@@ -85,21 +92,25 @@ class ModelBuilder():
                 self.base_model.trainable = self.base_trainable
 
             case mt.MOBILENET:
-                self.base_model = MobileNet(input_tensor=x_aug, **self.base_model_params)
+                self.base_model = MobileNet(
+                    input_tensor=x_aug, **self.base_model_params
+                )
                 self.base_model.trainable = self.base_trainable
 
             case mt.CAPSNET:
                 self.base_model = None
-                x = Rescaling(1./255)(x)
-                outputs = self.build_capsnet(inputs = x_aug, **self.model_params)
+                x = Rescaling(1.0 / 255)(x)
+                outputs = self.build_capsnet(inputs=x_aug, **self.model_params)
 
             case _:
-                raise Exception(f"Model type {self.model_type} not supported: {self.model_name}")
+                raise Exception(
+                    f"Model type {self.model_type} not supported: {self.model_name}"
+                )
 
         # Classification head
         if self.model_type in (mt.RESNET50, mt.MOBILENET):
             x = self.base_model.output
-            outputs = Dense(4, activation='softmax')(x)
+            outputs = Dense(4, activation="softmax")(x)
         elif self.model_type == mt.CAPSNET:
             pass
         else:
@@ -110,12 +121,9 @@ class ModelBuilder():
         self.model.compile(**compile_params)
 
         print(f"The {self.model_name} model has been compiled successfully")
-        
-        return self.base_model, self.model
 
+        return self.base_model, self.model
 
-            
-  
     def build_capsnet(self, inputs, **params):
         """
         Build a Capsule Network model for four class lung iseases classification: COVID, Normal, Pneumonia and Opacity.
@@ -130,77 +138,100 @@ class ModelBuilder():
         Returns:
             model: to be compiled
         """
-                
-        first_Conv2DKernel_size =  params.pop("first_Conv2DKernel_size")
-                
+
+        first_Conv2DKernel_size = params.pop("first_Conv2DKernel_size")
+
         # --- Preprocessing Layers ---
         x = inputs
 
         # --- Feature Extraction ---
         # learns 64 different 3x3 filters
-        x = layers.Conv2D(filters = 64, kernel_size=first_Conv2DKernel_size, strides=2, padding='valid', activation='relu')(x) # downsampling  strides=2, no padding because only exposed lung area matters/contains features
+        x = layers.Conv2D(
+            filters=64,
+            kernel_size=first_Conv2DKernel_size,
+            strides=2,
+            padding="valid",
+            activation="relu",
+        )(
+            x
+        )  # downsampling  strides=2, no padding because only exposed lung area matters/contains features
         x = layers.BatchNormalization()(x)
 
-        x = layers.Conv2D(128, 5, strides=2, padding='same', activation='relu')(x)    # padding="same" because of transformed output of the 1rst conv2D-layer (None, 125, 125, 64) to not lose the spatial info
+        x = layers.Conv2D(128, 5, strides=2, padding="same", activation="relu")(
+            x
+        )  # padding="same" because of transformed output of the 1rst conv2D-layer (None, 125, 125, 64) to not lose the spatial info
         x = layers.BatchNormalization()(x)
         x = layers.Dropout(0.25)(x)  # Dropout after second block (early regularization)
 
-        x = layers.Conv2D(128, 3, strides=1, padding='same', activation='relu')(x)
+        x = layers.Conv2D(128, 3, strides=1, padding="same", activation="relu")(x)
         x = layers.BatchNormalization()(x)
 
-        x = layers.Conv2D(256, 3, strides=1, padding='same', activation='relu')(x)
+        x = layers.Conv2D(256, 3, strides=1, padding="same", activation="relu")(x)
         x = layers.BatchNormalization()(x)
         x = layers.Dropout(0.3)(x)  # Deeper regularization after more feature maps
 
-        x = layers.Conv2D(512, 3, strides=1, padding='same', activation='relu')(x) # out : (None, 64, 64, 512)
-        x = layers.BatchNormalization()(x) # out: (None, 64, 64, 512)
+        x = layers.Conv2D(512, 3, strides=1, padding="same", activation="relu")(
+            x
+        )  # out : (None, 64, 64, 512)
+        x = layers.BatchNormalization()(x)  # out: (None, 64, 64, 512)
 
-        x = layers.Dropout(0.3)(x)  # Final dropout before capsules, out : (None, 64, 64, 512)
+        x = layers.Dropout(0.3)(
+            x
+        )  # Final dropout before capsules, out : (None, 64, 64, 512)
 
         # --- Capsule Layers for classification---
-        primary_caps = PrimaryCaps(**self.prim_caps_params)(x)      #dim_capsule=8, # Each capsule is an 8D vector (i.e. each capsule outputs a vector of length 8)
-                                                                #n_channels=32, # There are 32 capsule "types" per spatial location (like 32 different filters)
-                                                                #kernel_size=9,
-                                                                #strides=2,     # Moves the 3×3 kernel with stride x → if x > 1 it reduces spatial size by x (downsampling)
-                                                                #                # stride=1 This means the kernel moves 1 pixel at a time, covering every possible position in the input.
-                                                                #padding='same') # same: No padding → output size shrinks (no border pixels used)
-
-        digit_caps = DigitCaps( **self.digit_caps_params)(primary_caps)   #num_capsule=n_class, # 1 capsule per class (e.g. 4 diseases = 4 capsules)
-                                                            #dim_capsule=16,      # Each output capsule is a 16D vector → captures pose info
-                                                            #routing_iters=routing_iters # Use 3 iterations of dynamic routing (or EM routing) to refine capsule agreement
-                                                            #) # out: (None, 4, 1, 16)
+        primary_caps = PrimaryCaps(**self.prim_caps_params)(
+            x
+        )  # dim_capsule=8, # Each capsule is an 8D vector (i.e. each capsule outputs a vector of length 8)
+        # n_channels=32, # There are 32 capsule "types" per spatial location (like 32 different filters)
+        # kernel_size=9,
+        # strides=2,     # Moves the 3×3 kernel with stride x → if x > 1 it reduces spatial size by x (downsampling)
+        #                # stride=1 This means the kernel moves 1 pixel at a time, covering every possible position in the input.
+        # padding='same') # same: No padding → output size shrinks (no border pixels used)
+
+        digit_caps = DigitCaps(**self.digit_caps_params)(
+            primary_caps
+        )  # num_capsule=n_class, # 1 capsule per class (e.g. 4 diseases = 4 capsules)
+        # dim_capsule=16,      # Each output capsule is a 16D vector → captures pose info
+        # routing_iters=routing_iters # Use 3 iterations of dynamic routing (or EM routing) to refine capsule agreement
+        # ) # out: (None, 4, 1, 16)
 
         outputs = Length()(digit_caps)
-        
-        return outputs  
-
 
+        return outputs
 
 
 # Squash function: This function shrinks small vectors to zero and large vectors to unit vectors.
 def squash(vectors, axis=-1):
     s_squared_norm = tf.reduce_sum(tf.square(vectors), axis, keepdims=True)
     # tf.keras.backend.epsilon() on google coalb with A100 GPU = 1e-07
-    scale = s_squared_norm / (1 + s_squared_norm) / tf.sqrt(s_squared_norm + tf.keras.backend.epsilon())
+    scale = (
+        s_squared_norm
+        / (1 + s_squared_norm)
+        / tf.sqrt(s_squared_norm + tf.keras.backend.epsilon())
+    )
     return scale * vectors
 
 
-
 # PrimaryCaps Layer/ Lower-level capsules (e.g. detecting edges or textures)
-@register_keras_serializable() #make it serializable to .keras format
+@register_keras_serializable()  # make it serializable to .keras format
 class PrimaryCaps(layers.Layer):
 
-    def __init__(self, dim_capsule, n_channels, kernel_size, strides, padding, **kwargs):
+    def __init__(
+        self, dim_capsule, n_channels, kernel_size, strides, padding, **kwargs
+    ):
         super(PrimaryCaps, self).__init__(**kwargs)
-        self.conv = layers.Conv2D(filters=dim_capsule * n_channels,
-                                  kernel_size=kernel_size,
-                                  strides=strides,
-                                  padding=padding,
-                                  activation='relu')
+        self.conv = layers.Conv2D(
+            filters=dim_capsule * n_channels,
+            kernel_size=kernel_size,
+            strides=strides,
+            padding=padding,
+            activation="relu",
+        )
         self.dim_capsule = dim_capsule
         self.n_channels = n_channels
-        self.kernel_size = kernel_size   #
-        self.strides = strides           #
+        self.kernel_size = kernel_size  #
+        self.strides = strides  #
         self.padding = padding
 
     def build(self, input_shape):
@@ -208,27 +239,33 @@ class PrimaryCaps(layers.Layer):
         self.conv.build(input_shape)
         super().build(input_shape)  # Let Keras know the layer is built
 
-
     def call(self, inputs):
         outputs = self.conv(inputs)
-        outputs = tf.reshape(outputs, (-1, outputs.shape[1] * outputs.shape[2] * self.n_channels, self.dim_capsule))
+        outputs = tf.reshape(
+            outputs,
+            (
+                -1,
+                outputs.shape[1] * outputs.shape[2] * self.n_channels,
+                self.dim_capsule,
+            ),
+        )
         return squash(outputs)
 
-
     def get_config(self):
         # hook in to keras Layer to modify layer's config on reload
         config = super().get_config()
-        config.update({
-            "dim_capsule": self.dim_capsule,
-            "n_channels": self.n_channels,
-            "kernel_size": self.kernel_size,
-            "strides": self.strides,
-            "padding": self.padding
-        })
+        config.update(
+            {
+                "dim_capsule": self.dim_capsule,
+                "n_channels": self.n_channels,
+                "kernel_size": self.kernel_size,
+                "strides": self.strides,
+                "padding": self.padding,
+            }
+        )
         return config
 
 
-
 @register_keras_serializable()
 class DigitCaps(layers.Layer):
     # DigitCaps Layer / Higher-level capsules (e.g. detecting objects like digits or lungs)
@@ -242,10 +279,16 @@ class DigitCaps(layers.Layer):
     def build(self, input_shape):
         self.input_num_capsule = input_shape[1]
         self.input_dim_capsule = input_shape[2]
-        self.W = self.add_weight(shape=[self.input_num_capsule, self.num_capsule,
-                                        self.input_dim_capsule, self.dim_capsule],
-                                 initializer='glorot_uniform',
-                                 trainable=True)
+        self.W = self.add_weight(
+            shape=[
+                self.input_num_capsule,
+                self.num_capsule,
+                self.input_dim_capsule,
+                self.dim_capsule,
+            ],
+            initializer="glorot_uniform",
+            trainable=True,
+        )
 
     def call(self, inputs):
         inputs_expand = tf.expand_dims(inputs, 2)
@@ -253,31 +296,39 @@ class DigitCaps(layers.Layer):
         inputs_tiled = tf.tile(inputs_tiled, [1, 1, self.num_capsule, 1, 1])
         inputs_hat = tf.matmul(inputs_tiled, self.W)
 
-        b = tf.zeros(shape=[tf.shape(inputs)[0], self.input_num_capsule, self.num_capsule, 1, 1])
+        b = tf.zeros(
+            shape=[tf.shape(inputs)[0], self.input_num_capsule, self.num_capsule, 1, 1]
+        )
 
         # Dynamic Routing by Agreement algo
         for i in range(self.routing_iters):
-            c = tf.nn.softmax(b, axis=2)  # coupling coefficient, beacause of softmax(...) all c's connected to a single higher capsule sum to 1.
-            s = tf.reduce_sum(c * inputs_hat, axis=1, keepdims=True)  # weighted sum along axis=1
-            v = squash(s, axis=-2)    # shrinks small vectors to zero and large vectors to unit vectors
+            c = tf.nn.softmax(
+                b, axis=2
+            )  # coupling coefficient, beacause of softmax(...) all c's connected to a single higher capsule sum to 1.
+            s = tf.reduce_sum(
+                c * inputs_hat, axis=1, keepdims=True
+            )  # weighted sum along axis=1
+            v = squash(
+                s, axis=-2
+            )  # shrinks small vectors to zero and large vectors to unit vectors
             if i < self.routing_iters - 1:
                 b += tf.reduce_sum(inputs_hat * v, axis=-1, keepdims=True)
 
         return tf.squeeze(v, axis=1)
 
-
     def get_config(self):
         # hook in to keras Layer to modify layer's config on reload
         config = super().get_config()
-        config.update({
-            "num_capsule": self.num_capsule,
-            "dim_capsule": self.dim_capsule,
-            "routing_iters": self.routing_iters
-        })
+        config.update(
+            {
+                "num_capsule": self.num_capsule,
+                "dim_capsule": self.dim_capsule,
+                "routing_iters": self.routing_iters,
+            }
+        )
         return config
 
 
-
 # Length Layer
 @register_keras_serializable()
 class Length(layers.Layer):
@@ -285,7 +336,6 @@ class Length(layers.Layer):
         return tf.sqrt(tf.reduce_sum(tf.square(inputs), -1))
 
 
-
 # Margin Loss for Capsule Networks
 def margin_loss(y_true, y_pred):
     # y_true is a one-hot vector
@@ -293,21 +343,15 @@ def margin_loss(y_true, y_pred):
     m_plus = 0.9
     m_minus = 0.1
     lambda_val = 0.5
-    L = y_true * tf.square(tf.maximum(0., m_plus - y_pred)) + \
-        lambda_val * (1 - y_true) * tf.square(tf.maximum(0., y_pred - m_minus))
+    L = y_true * tf.square(tf.maximum(0.0, m_plus - y_pred)) + lambda_val * (
+        1 - y_true
+    ) * tf.square(tf.maximum(0.0, y_pred - m_minus))
     return tf.reduce_mean(tf.reduce_sum(L, axis=1))
 
 
 capsnet_custom_objects = {
-    'PrimaryCaps': PrimaryCaps,
-    'DigitCaps': DigitCaps,
-    'Length': Length,
-    'margin_loss': margin_loss
+    "PrimaryCaps": PrimaryCaps,
+    "DigitCaps": DigitCaps,
+    "Length": Length,
+    "margin_loss": margin_loss,
 }
-
-    
-            
-        
-        
-        
-