custom_objects.py

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.saving import register_keras_serializable

@register_keras_serializable()
class StressAmplifier(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def call(self, inputs):
        stress = inputs[:, 2]
        slip = inputs[:, 4]
        stress_boost = tf.sigmoid((stress - 400) * 0.01)
        slip_boost = tf.sigmoid((slip - 8) * 0.5)
        modulation = 1.0 + 0.4 * stress_boost * slip_boost
        return tf.expand_dims(modulation, axis=-1)

@register_keras_serializable()
class DepthSuppressor(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def call(self, inputs):
        depth = inputs[:, 3]
        suppression = tf.sigmoid((depth - 25) * 0.15)
        modulation = 1.0 - 0.3 * suppression
        return tf.expand_dims(modulation, axis=-1)

@register_keras_serializable()
class DisplacementActivator(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def call(self, inputs):
        displacement = inputs[:, 1]
        activation = tf.sigmoid((displacement - 30) * 0.08)
        modulation = 1.0 + 0.3 * activation
        return tf.expand_dims(modulation, axis=-1)

@register_keras_serializable()
class SoftScale(tf.keras.layers.Layer):
    def __init__(self, factor=0.25, **kwargs):
        super().__init__(**kwargs)
        self.factor = factor

    def call(self, inputs):
        return 1.0 + self.factor * tf.tanh(inputs - 1.0)

CUSTOM_OBJECTS = {
    'StressAmplifier': StressAmplifier,
    'DepthSuppressor': DepthSuppressor,
    'DisplacementActivator': DisplacementActivator,
    'SoftScale': SoftScale
}