tensorimmune 0.1.0__tar.gz

tensorimmune-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Developer
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tensorimmune-0.1.0/PKG-INFO

@@ -0,0 +1,129 @@
+Metadata-Version: 2.4
+Name: tensorimmune
+Version: 0.1.0
+Summary: A lightweight Keras wrapper that injects Out-of-Distribution detection natively into a TensorFlow/Keras model's computational graph.
+Author: Developer
+License: MIT
+Project-URL: Homepage, https://github.com/prysthedev/TensorImmune
+Project-URL: Bug Tracker, https://github.com/prysthedev/TensorImmune/issues
+Project-URL: Source Code, https://github.com/prysthedev/TensorImmune
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: tensorflow>=2.10.0
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Dynamic: license-file
+
+# TensorImmune
+
+**TensorImmune** is a lightweight Keras wrapper that injects Out-of-Distribution (OOD) detection natively into a TensorFlow/Keras model's computational graph. A single exported artifact can flag anomalous inputs at inference time without any external monitoring infrastructure!
+
+## Overview
+
+Deploying machine learning models to the real world involves dealing with unexpected or out-of-distribution data. Normally, OOD detection is handled by external monitoring infrastructure or complex multi-model setups. TensorImmune simplifies this by weaving a small, symbiotic autoencoder directly into your model's computational graph. 
+
+By observing an intermediate layer's activations (the "monitor layer"), TensorImmune simultaneously trains your primary task and the autoencoder. The model calibrates an anomaly threshold during training. When deployed, the model returns both the primary prediction and an immunity score (or anomaly flag), even when converted to edge formats like TensorFlow Lite!
+
+## Installation
+
+Install via pip:
+
+```bash
+pip install tensorimmune
+```
+
+## Quickstart
+
+```python
+import tensorflow as tf
+from tensorimmune import ImmuneModel
+
+# 1. Define your base Keras model (Sequential or Functional)
+base_model = tf.keras.Sequential([
+    tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
+    tf.keras.layers.Conv2D(32, 3, activation='relu'),
+    tf.keras.layers.MaxPooling2D(),
+    tf.keras.layers.Conv2D(64, 3, activation='relu'),
+    tf.keras.layers.GlobalAveragePooling2D(),
+    tf.keras.layers.Dense(64, activation='relu', name='feature_bottleneck'),
+    tf.keras.layers.Dense(10, activation='softmax')
+])
+
+# 2. Wrap it with ImmuneModel
+# We monitor the 'feature_bottleneck' layer. 
+# sensitivity=0.95 means we set the anomaly threshold to the 95th percentile 
+# of the reconstruction errors seen during training.
+model = ImmuneModel(
+    base_model=base_model,
+    monitor_layer='feature_bottleneck',
+    sensitivity=0.95 
+)
+
+# 3. Compile and train normally!
+model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+model.fit(x_train, y_train, epochs=5)
+
+# 4. Predict returns task predictions AND anomaly flags
+predictions, is_anomaly = model.predict(x_test)
+
+# You can also get the raw continuous immunity score
+predictions, immunity_scores = model.predict(x_test, return_score=True)
+```
+
+## How Symbiotic Training Works
+
+Under the hood, `ImmuneModel` creates a `feature_extractor` that splits the graph at the `monitor_layer`. It then dynamically builds a small autoencoder to reconstruct the activations of that layer. 
+
+By subclassing `tf.keras.Model` and overriding `train_step`, TensorImmune trains both the original task loss and the autoencoder's mean squared error (MSE) reconstruction loss simultaneously in a single backward pass. The losses are combined via the `immune_loss_weight` parameter (which defaults to `1.0`).
+
+If the monitor layer outputs a spatial tensor (e.g. from a Conv2D layer), TensorImmune automatically inserts a GlobalAveragePooling layer before the autoencoder to flatten the representations intelligently.
+
+## Immunity Score & Sensitivity Calibration
+
+During training, TensorImmune maintains a running record of reconstruction error statistics (mean, variance, count) inside the model's graph. 
+
+The `sensitivity` parameter (between 0 and 1) dictates how strictly the model should flag anomalies. It is treated as a percentile of the training distribution. Using the running statistics and the inverse error function, TensorImmune continuously calibrates a numeric anomaly threshold (e.g. `sensitivity=0.95` maps to ~1.64 standard deviations above the mean). 
+
+At the end of training, this threshold is automatically stored as a non-trainable `tf.Variable` directly inside the model. When you save and export the model, the threshold is serialized with it!
+
+## Edge Deployment with TFLite
+
+Because TensorImmune builds its architecture and threshold logic using pure `tf.keras.layers` and `tf` operations, the resulting model can be converted seamlessly to TensorFlow Lite for edge devices.
+
+```python
+import tensorflow as tf
+
+# Get concrete function for TFLite
+run_model = tf.function(lambda x: model(x, return_score=False))
+concrete_func = run_model.get_concrete_function(
+    tf.TensorSpec([1, 28, 28, 1], tf.float32)
+)
+
+# Convert! No custom ops required.
+converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+tflite_model = converter.convert()
+
+with open("immune_model.tflite", "wb") as f:
+    f.write(tflite_model)
+```
+
+The exported TFLite model natively returns the boolean anomaly flag as its second output, making edge integration trivial.
+
+## Limitations
+
+It is important to state a known limitation: **autoencoder reconstruction error on intermediate features is generally a weaker OOD signal than more rigorous methods** (like Mahalanobis distance, Energy-based scoring, or evidential deep learning). 
+
+TensorImmune is intended as a **lightweight, drop-in first line of defense** for edge deployment where you cannot afford external monitoring infrastructure or complex multi-model pipelines. It should not be used as a replacement for rigorous safety-critical OOD research techniques.
+
+## Contributing
+
+Contributions are welcome! Please open an issue or submit a pull request if you have ideas for improvements, bug fixes, or new features.
+
+When contributing:
+1. Ensure your code passes all existing tests.
+2. Add tests for new features.
+3. Update the documentation as necessary.
+
+License: MIT

tensorimmune-0.1.0/README.md

@@ -0,0 +1,111 @@
+# TensorImmune
+
+**TensorImmune** is a lightweight Keras wrapper that injects Out-of-Distribution (OOD) detection natively into a TensorFlow/Keras model's computational graph. A single exported artifact can flag anomalous inputs at inference time without any external monitoring infrastructure!
+
+## Overview
+
+Deploying machine learning models to the real world involves dealing with unexpected or out-of-distribution data. Normally, OOD detection is handled by external monitoring infrastructure or complex multi-model setups. TensorImmune simplifies this by weaving a small, symbiotic autoencoder directly into your model's computational graph. 
+
+By observing an intermediate layer's activations (the "monitor layer"), TensorImmune simultaneously trains your primary task and the autoencoder. The model calibrates an anomaly threshold during training. When deployed, the model returns both the primary prediction and an immunity score (or anomaly flag), even when converted to edge formats like TensorFlow Lite!
+
+## Installation
+
+Install via pip:
+
+```bash
+pip install tensorimmune
+```
+
+## Quickstart
+
+```python
+import tensorflow as tf
+from tensorimmune import ImmuneModel
+
+# 1. Define your base Keras model (Sequential or Functional)
+base_model = tf.keras.Sequential([
+    tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
+    tf.keras.layers.Conv2D(32, 3, activation='relu'),
+    tf.keras.layers.MaxPooling2D(),
+    tf.keras.layers.Conv2D(64, 3, activation='relu'),
+    tf.keras.layers.GlobalAveragePooling2D(),
+    tf.keras.layers.Dense(64, activation='relu', name='feature_bottleneck'),
+    tf.keras.layers.Dense(10, activation='softmax')
+])
+
+# 2. Wrap it with ImmuneModel
+# We monitor the 'feature_bottleneck' layer. 
+# sensitivity=0.95 means we set the anomaly threshold to the 95th percentile 
+# of the reconstruction errors seen during training.
+model = ImmuneModel(
+    base_model=base_model,
+    monitor_layer='feature_bottleneck',
+    sensitivity=0.95 
+)
+
+# 3. Compile and train normally!
+model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+model.fit(x_train, y_train, epochs=5)
+
+# 4. Predict returns task predictions AND anomaly flags
+predictions, is_anomaly = model.predict(x_test)
+
+# You can also get the raw continuous immunity score
+predictions, immunity_scores = model.predict(x_test, return_score=True)
+```
+
+## How Symbiotic Training Works
+
+Under the hood, `ImmuneModel` creates a `feature_extractor` that splits the graph at the `monitor_layer`. It then dynamically builds a small autoencoder to reconstruct the activations of that layer. 
+
+By subclassing `tf.keras.Model` and overriding `train_step`, TensorImmune trains both the original task loss and the autoencoder's mean squared error (MSE) reconstruction loss simultaneously in a single backward pass. The losses are combined via the `immune_loss_weight` parameter (which defaults to `1.0`).
+
+If the monitor layer outputs a spatial tensor (e.g. from a Conv2D layer), TensorImmune automatically inserts a GlobalAveragePooling layer before the autoencoder to flatten the representations intelligently.
+
+## Immunity Score & Sensitivity Calibration
+
+During training, TensorImmune maintains a running record of reconstruction error statistics (mean, variance, count) inside the model's graph. 
+
+The `sensitivity` parameter (between 0 and 1) dictates how strictly the model should flag anomalies. It is treated as a percentile of the training distribution. Using the running statistics and the inverse error function, TensorImmune continuously calibrates a numeric anomaly threshold (e.g. `sensitivity=0.95` maps to ~1.64 standard deviations above the mean). 
+
+At the end of training, this threshold is automatically stored as a non-trainable `tf.Variable` directly inside the model. When you save and export the model, the threshold is serialized with it!
+
+## Edge Deployment with TFLite
+
+Because TensorImmune builds its architecture and threshold logic using pure `tf.keras.layers` and `tf` operations, the resulting model can be converted seamlessly to TensorFlow Lite for edge devices.
+
+```python
+import tensorflow as tf
+
+# Get concrete function for TFLite
+run_model = tf.function(lambda x: model(x, return_score=False))
+concrete_func = run_model.get_concrete_function(
+    tf.TensorSpec([1, 28, 28, 1], tf.float32)
+)
+
+# Convert! No custom ops required.
+converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+tflite_model = converter.convert()
+
+with open("immune_model.tflite", "wb") as f:
+    f.write(tflite_model)
+```
+
+The exported TFLite model natively returns the boolean anomaly flag as its second output, making edge integration trivial.
+
+## Limitations
+
+It is important to state a known limitation: **autoencoder reconstruction error on intermediate features is generally a weaker OOD signal than more rigorous methods** (like Mahalanobis distance, Energy-based scoring, or evidential deep learning). 
+
+TensorImmune is intended as a **lightweight, drop-in first line of defense** for edge deployment where you cannot afford external monitoring infrastructure or complex multi-model pipelines. It should not be used as a replacement for rigorous safety-critical OOD research techniques.
+
+## Contributing
+
+Contributions are welcome! Please open an issue or submit a pull request if you have ideas for improvements, bug fixes, or new features.
+
+When contributing:
+1. Ensure your code passes all existing tests.
+2. Add tests for new features.
+3. Update the documentation as necessary.
+
+License: MIT

tensorimmune-0.1.0/pyproject.toml

@@ -0,0 +1,31 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "tensorimmune"
+version = "0.1.0"
+description = "A lightweight Keras wrapper that injects Out-of-Distribution detection natively into a TensorFlow/Keras model's computational graph."
+readme = "README.md"
+requires-python = ">=3.8"
+license = { text = "MIT" }
+authors = [
+    { name = "Developer" }
+]
+dependencies = [
+    "tensorflow>=2.10.0",
+    "numpy"
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[project.urls]
+"Homepage" = "https://github.com/prysthedev/TensorImmune"
+"Bug Tracker" = "https://github.com/prysthedev/TensorImmune/issues"
+"Source Code" = "https://github.com/prysthedev/TensorImmune"

tensorimmune-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tensorimmune-0.1.0/src/tensorimmune/__init__.py

@@ -0,0 +1,3 @@
+from .model import ImmuneModel
+
+__all__ = ["ImmuneModel"]

tensorimmune-0.1.0/src/tensorimmune/autoencoder.py

@@ -0,0 +1,25 @@
+import tensorflow as tf
+
+def build_autoencoder(input_dim, bottleneck_size=None):
+    """
+    Builds a small symbiotic autoencoder to reconstruct layer activations.
+    
+    Args:
+        input_dim (int): The dimension of the flattened activation vector.
+        bottleneck_size (int, optional): The size of the bottleneck layer. Defaults to roughly input_dim // 4.
+        
+    Returns:
+        tf.keras.Model: The autoencoder model.
+    """
+    if bottleneck_size is None:
+        bottleneck_size = max(1, input_dim // 4)
+        
+    # Use standard Keras API to build the autoencoder
+    inputs = tf.keras.Input(shape=(input_dim,))
+    x = tf.keras.layers.Dense(max(bottleneck_size * 2, 2), activation='relu')(inputs)
+    x = tf.keras.layers.Dense(bottleneck_size, activation='relu', name='bottleneck')(x)
+    x = tf.keras.layers.Dense(max(bottleneck_size * 2, 2), activation='relu')(x)
+    outputs = tf.keras.layers.Dense(input_dim, activation='linear')(x)
+    
+    model = tf.keras.Model(inputs=inputs, outputs=outputs, name='symbiotic_autoencoder')
+    return model

tensorimmune-0.1.0/src/tensorimmune/calibration.py

@@ -0,0 +1,23 @@
+import tensorflow as tf
+
+def get_threshold_from_stats(mean, variance, sensitivity):
+    """
+    Calculates the numeric anomaly threshold based on running mean and variance and sensitivity.
+    
+    Args:
+        mean (tf.Tensor): EMA of reconstruction error.
+        variance (tf.Tensor): EMA of variance of reconstruction error.
+        sensitivity (float or tf.Tensor): Percentile sensitivity (0 to 1).
+        
+    Returns:
+        tf.Tensor: The calibrated threshold.
+    """
+    std = tf.sqrt(tf.maximum(variance, 0.0))
+    
+    # Map sensitivity (percentile) to number of standard deviations for a normal distribution
+    p = tf.cast(sensitivity, tf.float32)
+    p = tf.clip_by_value(p, 0.0001, 0.9999)
+    z = tf.math.sqrt(2.0) * tf.math.erfinv(2.0 * p - 1.0)
+    
+    threshold = mean + z * std
+    return threshold

tensorimmune-0.1.0/src/tensorimmune/model.py

@@ -0,0 +1,215 @@
+import tensorflow as tf
+from .shape_utils import get_pooling_layer
+from .autoencoder import build_autoencoder
+from .calibration import get_threshold_from_stats
+
+@tf.keras.utils.register_keras_serializable(package="TensorImmune")
+class ImmuneModel(tf.keras.Model):
+    """
+    A Keras wrapper that injects Out-of-Distribution detection natively into a 
+    TensorFlow/Keras model's computational graph.
+    """
+    def __init__(self, base_model, monitor_layer, sensitivity=0.95, bottleneck_size=None, immune_loss_weight=1.0, **kwargs):
+        super(ImmuneModel, self).__init__(**kwargs)
+        
+        # Validate base_model
+        if not hasattr(base_model, "inputs") or not base_model.inputs:
+            raise ValueError("The base model must be a Keras Sequential or Functional model. Subclassed models without explicitly defined inputs are not supported because their computational graph cannot be introspected.")
+            
+        try:
+            target_layer = base_model.get_layer(monitor_layer)
+        except ValueError:
+            available_layers = [l.name for l in base_model.layers]
+            raise ValueError(f"Monitor layer '{monitor_layer}' not found in the base model. Available layers: {available_layers}")
+            
+        self.base_model_ref = base_model
+        self.monitor_layer = monitor_layer
+        self.sensitivity = float(sensitivity)
+        self.bottleneck_size = bottleneck_size
+        self.immune_loss_weight = float(immune_loss_weight)
+        
+        # Safely get base model outputs
+        base_outputs = base_model.outputs if hasattr(base_model, 'outputs') and base_model.outputs else [base_model.output]
+        self.base_output_count = len(base_outputs)
+        
+        # Build feature extractor capturing both task output(s) and monitor layer
+        self.feature_extractor = tf.keras.Model(
+            inputs=base_model.inputs,
+            outputs=base_outputs + [target_layer.output],
+            name="feature_extractor"
+        )
+        
+        # Determine pool layer based on shape
+        monitor_shape = target_layer.output.shape
+        rank = len(monitor_shape)
+        self.pool_layer = get_pooling_layer(rank)
+        
+        # Calculate flattened dimension
+        dummy_input = tf.keras.Input(shape=monitor_shape[1:])
+        pooled_dummy = self.pool_layer(dummy_input)
+        flat_dim = pooled_dummy.shape[-1]
+        
+        # Build symbiotic autoencoder
+        self.autoencoder = build_autoencoder(flat_dim, bottleneck_size)
+        
+        # EMA Statistics variables for threshold calibration
+        self.calib_mean = self.add_weight(name='calib_mean', shape=(), initializer='zeros', trainable=False, dtype=tf.float32)
+        self.calib_var = self.add_weight(name='calib_var', shape=(), initializer='zeros', trainable=False, dtype=tf.float32)
+        self.calib_count = self.add_weight(name='calib_count', shape=(), initializer='zeros', trainable=False, dtype=tf.float32)
+        
+        # Calibrated anomaly threshold (serialized with model)
+        self.anomaly_threshold = self.add_weight(name='anomaly_threshold', shape=(), initializer='zeros', trainable=False, dtype=tf.float32)
+
+    def _update_calibration_stats(self, batch_mse):
+        batch_mean = tf.reduce_mean(batch_mse)
+        # Prevent variance calculation errors if batch size is 1
+        batch_size = tf.cast(tf.shape(batch_mse)[0], tf.float32)
+        batch_var = tf.where(batch_size > 1.0, tf.math.reduce_variance(batch_mse), 0.0)
+        
+        momentum = tf.constant(0.99, dtype=tf.float32)
+        
+        is_first = tf.equal(self.calib_count, 0.0)
+        
+        def first_step():
+            self.calib_mean.assign(batch_mean)
+            self.calib_var.assign(batch_var)
+            return tf.constant(0.0)
+            
+        def next_steps():
+            delta = batch_mean - self.calib_mean
+            new_mean = self.calib_mean + (1.0 - momentum) * delta
+            new_var = momentum * self.calib_var + (1.0 - momentum) * batch_var + momentum * (1.0 - momentum) * tf.square(delta)
+            self.calib_mean.assign(new_mean)
+            self.calib_var.assign(new_var)
+            return tf.constant(0.0)
+            
+        tf.cond(is_first, first_step, next_steps)
+        self.calib_count.assign_add(1.0)
+        
+        new_threshold = get_threshold_from_stats(self.calib_mean, self.calib_var, self.sensitivity)
+        self.anomaly_threshold.assign(new_threshold)
+
+    @tf.function
+    def call(self, inputs, training=False, return_score=False):
+        extractor_outputs = self.feature_extractor(inputs, training=training)
+        
+        task_output = extractor_outputs[:-1]
+        if self.base_output_count == 1:
+            task_output = task_output[0]
+            
+        monitor_acts = extractor_outputs[-1]
+        
+        pooled_acts = self.pool_layer(monitor_acts)
+        reconstructed = self.autoencoder(pooled_acts, training=training)
+        
+        mse = tf.reduce_mean(tf.square(pooled_acts - reconstructed), axis=-1)
+        
+        if return_score:
+            return task_output, mse
+        else:
+            return task_output, mse > self.anomaly_threshold
+
+    def train_step(self, data):
+        if len(data) == 3:
+            x, y, sample_weight = data
+        else:
+            x, y = data
+            sample_weight = None
+
+        with tf.GradientTape() as tape:
+            extractor_outputs = self.feature_extractor(x, training=True)
+            task_output = extractor_outputs[:-1]
+            if self.base_output_count == 1:
+                task_output = task_output[0]
+            monitor_acts = extractor_outputs[-1]
+            
+            pooled_acts = self.pool_layer(monitor_acts)
+            reconstructed = self.autoencoder(pooled_acts, training=True)
+            
+            if hasattr(self, "compute_loss"):
+                loss = self.compute_loss(x=x, y=y, y_pred=task_output, sample_weight=sample_weight)
+            else:
+                loss = self.compiled_loss(y, task_output, sample_weight=sample_weight, regularization_losses=self.losses)
+            
+            batch_mse = tf.reduce_mean(tf.square(pooled_acts - reconstructed), axis=-1)
+            ae_loss = tf.reduce_mean(batch_mse)
+            
+            total_loss = loss + self.immune_loss_weight * ae_loss
+
+        trainable_vars = self.trainable_variables
+        gradients = tape.gradient(total_loss, trainable_vars)
+        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
+        
+        self._update_calibration_stats(batch_mse)
+        
+        if hasattr(self, "compute_metrics"):
+            self.compute_metrics(x=x, y=y, y_pred=task_output, sample_weight=sample_weight)
+            results = {m.name: m.result() for m in self.metrics}
+        else:
+            self.compiled_metrics.update_state(y, task_output, sample_weight=sample_weight)
+            results = {m.name: m.result() for m in self.metrics}
+            
+        results["ae_loss"] = ae_loss
+        results["threshold"] = self.anomaly_threshold
+        return results
+
+    def predict_step(self, data):
+        if isinstance(data, tuple):
+            x = data[0]
+        else:
+            x = data
+        return self(x, training=False, return_score=True)
+
+    def predict(self, x, *args, **kwargs):
+        return_score = kwargs.pop('return_score', False)
+        outputs = super().predict(x, *args, **kwargs)
+        
+        # Keras predict unrolls outputs. The last element is always the MSE score.
+        if isinstance(outputs, list) and len(outputs) > 1:
+            scores = outputs[-1]
+            task_preds = outputs[:-1]
+            if len(task_preds) == 1:
+                task_preds = task_preds[0]
+        else:
+            # Fallback if outputs is somehow not unrolled
+            task_preds, scores = outputs[0], outputs[1]
+            
+        if return_score:
+            return task_preds, scores
+        else:
+            flags = scores > self.anomaly_threshold.numpy()
+            return task_preds, flags
+
+    def get_config(self):
+        config = super().get_config()
+        try:
+            base_model_config = tf.keras.utils.serialize_keras_object(self.base_model_ref)
+        except AttributeError:
+            if hasattr(tf.keras.saving, 'serialize_keras_object'):
+                base_model_config = tf.keras.saving.serialize_keras_object(self.base_model_ref)
+            else:
+                base_model_config = self.base_model_ref.get_config()
+                
+        config.update({
+            "base_model": base_model_config,
+            "monitor_layer": self.monitor_layer,
+            "sensitivity": self.sensitivity,
+            "bottleneck_size": self.bottleneck_size,
+            "immune_loss_weight": self.immune_loss_weight,
+        })
+        return config
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        base_model_config = config.pop("base_model")
+        try:
+            base_model = tf.keras.utils.deserialize_keras_object(base_model_config, custom_objects=custom_objects)
+        except AttributeError:
+            if hasattr(tf.keras.saving, 'deserialize_keras_object'):
+                base_model = tf.keras.saving.deserialize_keras_object(base_model_config, custom_objects=custom_objects)
+            else:
+                if "layers" in base_model_config:
+                    base_model = tf.keras.Sequential.from_config(base_model_config, custom_objects=custom_objects)
+                else:
+                    base_model = tf.keras.Model.from_config(base_model_config, custom_objects=custom_objects)
+        return cls(base_model=base_model, **config)

tensorimmune-0.1.0/src/tensorimmune/shape_utils.py

@@ -0,0 +1,22 @@
+import tensorflow as tf
+
+def get_pooling_layer(rank):
+    """
+    Returns the appropriate pooling or flattening layer based on the input tensor's rank.
+    
+    Args:
+        rank (int): The rank (number of dimensions) of the tensor.
+        
+    Returns:
+        tf.keras.layers.Layer: The pooling or flatten layer.
+    """
+    if rank == 2:
+        return tf.keras.layers.Flatten()
+    elif rank == 3:
+        return tf.keras.layers.GlobalAveragePooling1D()
+    elif rank == 4:
+        return tf.keras.layers.GlobalAveragePooling2D()
+    elif rank == 5:
+        return tf.keras.layers.GlobalAveragePooling3D()
+    else:
+        return tf.keras.layers.Flatten()

tensorimmune-0.1.0/src/tensorimmune.egg-info/PKG-INFO

@@ -0,0 +1,129 @@
+Metadata-Version: 2.4
+Name: tensorimmune
+Version: 0.1.0
+Summary: A lightweight Keras wrapper that injects Out-of-Distribution detection natively into a TensorFlow/Keras model's computational graph.
+Author: Developer
+License: MIT
+Project-URL: Homepage, https://github.com/prysthedev/TensorImmune
+Project-URL: Bug Tracker, https://github.com/prysthedev/TensorImmune/issues
+Project-URL: Source Code, https://github.com/prysthedev/TensorImmune
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: tensorflow>=2.10.0
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Dynamic: license-file
+
+# TensorImmune
+
+**TensorImmune** is a lightweight Keras wrapper that injects Out-of-Distribution (OOD) detection natively into a TensorFlow/Keras model's computational graph. A single exported artifact can flag anomalous inputs at inference time without any external monitoring infrastructure!
+
+## Overview
+
+Deploying machine learning models to the real world involves dealing with unexpected or out-of-distribution data. Normally, OOD detection is handled by external monitoring infrastructure or complex multi-model setups. TensorImmune simplifies this by weaving a small, symbiotic autoencoder directly into your model's computational graph. 
+
+By observing an intermediate layer's activations (the "monitor layer"), TensorImmune simultaneously trains your primary task and the autoencoder. The model calibrates an anomaly threshold during training. When deployed, the model returns both the primary prediction and an immunity score (or anomaly flag), even when converted to edge formats like TensorFlow Lite!
+
+## Installation
+
+Install via pip:
+
+```bash
+pip install tensorimmune
+```
+
+## Quickstart
+
+```python
+import tensorflow as tf
+from tensorimmune import ImmuneModel
+
+# 1. Define your base Keras model (Sequential or Functional)
+base_model = tf.keras.Sequential([
+    tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
+    tf.keras.layers.Conv2D(32, 3, activation='relu'),
+    tf.keras.layers.MaxPooling2D(),
+    tf.keras.layers.Conv2D(64, 3, activation='relu'),
+    tf.keras.layers.GlobalAveragePooling2D(),
+    tf.keras.layers.Dense(64, activation='relu', name='feature_bottleneck'),
+    tf.keras.layers.Dense(10, activation='softmax')
+])
+
+# 2. Wrap it with ImmuneModel
+# We monitor the 'feature_bottleneck' layer. 
+# sensitivity=0.95 means we set the anomaly threshold to the 95th percentile 
+# of the reconstruction errors seen during training.
+model = ImmuneModel(
+    base_model=base_model,
+    monitor_layer='feature_bottleneck',
+    sensitivity=0.95 
+)
+
+# 3. Compile and train normally!
+model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+model.fit(x_train, y_train, epochs=5)
+
+# 4. Predict returns task predictions AND anomaly flags
+predictions, is_anomaly = model.predict(x_test)
+
+# You can also get the raw continuous immunity score
+predictions, immunity_scores = model.predict(x_test, return_score=True)
+```
+
+## How Symbiotic Training Works
+
+Under the hood, `ImmuneModel` creates a `feature_extractor` that splits the graph at the `monitor_layer`. It then dynamically builds a small autoencoder to reconstruct the activations of that layer. 
+
+By subclassing `tf.keras.Model` and overriding `train_step`, TensorImmune trains both the original task loss and the autoencoder's mean squared error (MSE) reconstruction loss simultaneously in a single backward pass. The losses are combined via the `immune_loss_weight` parameter (which defaults to `1.0`).
+
+If the monitor layer outputs a spatial tensor (e.g. from a Conv2D layer), TensorImmune automatically inserts a GlobalAveragePooling layer before the autoencoder to flatten the representations intelligently.
+
+## Immunity Score & Sensitivity Calibration
+
+During training, TensorImmune maintains a running record of reconstruction error statistics (mean, variance, count) inside the model's graph. 
+
+The `sensitivity` parameter (between 0 and 1) dictates how strictly the model should flag anomalies. It is treated as a percentile of the training distribution. Using the running statistics and the inverse error function, TensorImmune continuously calibrates a numeric anomaly threshold (e.g. `sensitivity=0.95` maps to ~1.64 standard deviations above the mean). 
+
+At the end of training, this threshold is automatically stored as a non-trainable `tf.Variable` directly inside the model. When you save and export the model, the threshold is serialized with it!
+
+## Edge Deployment with TFLite
+
+Because TensorImmune builds its architecture and threshold logic using pure `tf.keras.layers` and `tf` operations, the resulting model can be converted seamlessly to TensorFlow Lite for edge devices.
+
+```python
+import tensorflow as tf
+
+# Get concrete function for TFLite
+run_model = tf.function(lambda x: model(x, return_score=False))
+concrete_func = run_model.get_concrete_function(
+    tf.TensorSpec([1, 28, 28, 1], tf.float32)
+)
+
+# Convert! No custom ops required.
+converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+tflite_model = converter.convert()
+
+with open("immune_model.tflite", "wb") as f:
+    f.write(tflite_model)
+```
+
+The exported TFLite model natively returns the boolean anomaly flag as its second output, making edge integration trivial.
+
+## Limitations
+
+It is important to state a known limitation: **autoencoder reconstruction error on intermediate features is generally a weaker OOD signal than more rigorous methods** (like Mahalanobis distance, Energy-based scoring, or evidential deep learning). 
+
+TensorImmune is intended as a **lightweight, drop-in first line of defense** for edge deployment where you cannot afford external monitoring infrastructure or complex multi-model pipelines. It should not be used as a replacement for rigorous safety-critical OOD research techniques.
+
+## Contributing
+
+Contributions are welcome! Please open an issue or submit a pull request if you have ideas for improvements, bug fixes, or new features.
+
+When contributing:
+1. Ensure your code passes all existing tests.
+2. Add tests for new features.
+3. Update the documentation as necessary.
+
+License: MIT

tensorimmune-0.1.0/src/tensorimmune.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+LICENSE
+README.md
+pyproject.toml
+src/tensorimmune/__init__.py
+src/tensorimmune/autoencoder.py
+src/tensorimmune/calibration.py
+src/tensorimmune/model.py
+src/tensorimmune/shape_utils.py
+src/tensorimmune.egg-info/PKG-INFO
+src/tensorimmune.egg-info/SOURCES.txt
+src/tensorimmune.egg-info/dependency_links.txt
+src/tensorimmune.egg-info/requires.txt
+src/tensorimmune.egg-info/top_level.txt
+tests/test_model.py

tensorimmune-0.1.0/src/tensorimmune.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

tensorimmune-0.1.0/src/tensorimmune.egg-info/requires.txt

@@ -0,0 +1,5 @@
+tensorflow>=2.10.0
+numpy
+
+[dev]
+pytest

tensorimmune-0.1.0/src/tensorimmune.egg-info/top_level.txt

@@ -0,0 +1 @@
+tensorimmune

tensorimmune-0.1.0/tests/test_model.py

@@ -0,0 +1,131 @@
+import pytest
+import numpy as np
+import tensorflow as tf
+from tensorimmune import ImmuneModel
+import os
+import tempfile
+
+def create_sequential_base_model():
+    model = tf.keras.Sequential([
+        tf.keras.layers.InputLayer(input_shape=(10,)),
+        tf.keras.layers.Dense(16, activation='relu', name='feature_bottleneck'),
+        tf.keras.layers.Dense(2, activation='softmax', name='output')
+    ])
+    return model
+
+def create_functional_base_model():
+    inputs = tf.keras.Input(shape=(10,))
+    x = tf.keras.layers.Dense(16, activation='relu', name='feature_bottleneck')(inputs)
+    outputs = tf.keras.layers.Dense(2, activation='softmax', name='output')(x)
+    return tf.keras.Model(inputs=inputs, outputs=outputs)
+
+def test_build_from_sequential():
+    base_model = create_sequential_base_model()
+    immune_model = ImmuneModel(base_model, monitor_layer='feature_bottleneck', sensitivity=0.95)
+    assert immune_model.monitor_layer == 'feature_bottleneck'
+    assert immune_model.anomaly_threshold is not None
+
+def test_build_from_functional():
+    base_model = create_functional_base_model()
+    immune_model = ImmuneModel(base_model, monitor_layer='feature_bottleneck', sensitivity=0.95)
+    assert immune_model.monitor_layer == 'feature_bottleneck'
+
+def test_reject_invalid_layer():
+    base_model = create_sequential_base_model()
+    with pytest.raises(ValueError, match="not found in the base model"):
+        ImmuneModel(base_model, monitor_layer='nonexistent_layer')
+
+def test_training_and_prediction():
+    base_model = create_sequential_base_model()
+    immune_model = ImmuneModel(base_model, monitor_layer='feature_bottleneck', sensitivity=0.95)
+    
+    immune_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+    
+    # Synthetic data (in-distribution)
+    x_train = np.random.normal(0, 1, size=(100, 10)).astype(np.float32)
+    y_train = np.random.randint(0, 2, size=(100,)).astype(np.float32)
+    
+    immune_model.fit(x_train, y_train, epochs=2, batch_size=16, verbose=0)
+    
+    # Test predict shapes and types
+    preds, flags = immune_model.predict(x_train[:5], return_score=False)
+    assert preds.shape == (5, 2)
+    assert flags.shape == (5,)
+    assert flags.dtype == bool
+    
+    preds, scores = immune_model.predict(x_train[:5], return_score=True)
+    assert scores.shape == (5,)
+    assert scores.dtype == np.float32
+
+    # Check threshold is a float variable
+    threshold = immune_model.anomaly_threshold.numpy()
+    assert isinstance(threshold, (float, np.float32))
+    assert threshold >= 0.0
+
+def test_ood_detection():
+    base_model = create_sequential_base_model()
+    immune_model = ImmuneModel(base_model, monitor_layer='feature_bottleneck', sensitivity=0.5)
+    
+    immune_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
+    
+    # Train on specific distribution
+    x_train = np.random.normal(0, 1, size=(500, 10)).astype(np.float32)
+    y_train = np.random.randint(0, 2, size=(500,)).astype(np.float32)
+    
+    immune_model.fit(x_train, y_train, epochs=5, batch_size=32, verbose=0)
+    
+    # OOD data (different distribution)
+    x_ood = np.random.normal(10, 5, size=(100, 10)).astype(np.float32)
+    
+    _, scores_id = immune_model.predict(x_train[:100], return_score=True)
+    _, scores_ood = immune_model.predict(x_ood, return_score=True)
+    
+    assert np.mean(scores_ood) > np.mean(scores_id)
+
+def test_save_and_load():
+    base_model = create_functional_base_model()
+    immune_model = ImmuneModel(base_model, monitor_layer='feature_bottleneck', sensitivity=0.95)
+    immune_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
+    
+    x_train = np.random.normal(0, 1, size=(50, 10)).astype(np.float32)
+    y_train = np.random.randint(0, 2, size=(50,)).astype(np.float32)
+    immune_model.fit(x_train, y_train, epochs=1, verbose=0)
+    
+    preds_orig, scores_orig = immune_model.predict(x_train[:5], return_score=True)
+    
+    with tempfile.TemporaryDirectory() as tmpdir:
+        save_path = os.path.join(tmpdir, "model.keras")
+        # Save model
+        immune_model.save(save_path)
+        
+        # Load model
+        # We need custom_objects to load ImmuneModel
+        custom_objects = {"ImmuneModel": ImmuneModel}
+        loaded_model = tf.keras.models.load_model(save_path, custom_objects=custom_objects)
+        
+        preds_loaded, scores_loaded = loaded_model.predict(x_train[:5], return_score=True)
+        
+        np.testing.assert_allclose(preds_orig, preds_loaded, rtol=1e-5)
+        np.testing.assert_allclose(scores_orig, scores_loaded, rtol=1e-5)
+        assert loaded_model.anomaly_threshold.numpy() == immune_model.anomaly_threshold.numpy()
+
+def test_tflite_conversion():
+    base_model = create_sequential_base_model()
+    immune_model = ImmuneModel(base_model, monitor_layer='feature_bottleneck', sensitivity=0.95)
+    immune_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
+    
+    x_train = np.random.normal(0, 1, size=(10, 10)).astype(np.float32)
+    y_train = np.random.randint(0, 2, size=(10,)).astype(np.float32)
+    immune_model.fit(x_train, y_train, epochs=1, verbose=0)
+    
+    # Get concrete function for TFLite
+    run_model = tf.function(lambda x: immune_model(x, return_score=False))
+    concrete_func = run_model.get_concrete_function(
+        tf.TensorSpec(immune_model.base_model_ref.inputs[0].shape, immune_model.base_model_ref.inputs[0].dtype)
+    )
+    
+    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+    tflite_model = converter.convert()
+    
+    assert tflite_model is not None
+    assert len(tflite_model) > 0