Implement SafeDense interceptor to bypass Keras 3 h5 loading bug

requirements.txt

@@ -4,4 +4,4 @@ python-multipart
 opencv-python
 numpy
 matplotlib
-tensorflow-cpu==2.15.0
+tensorflow-cpu

src/api/api.py

@@ -29,11 +29,26 @@ app.add_middleware(
 
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 
+class SafeDense(tf.keras.layers.Dense):
+    def __init__(self, **kwargs):
+        kwargs.pop('quantization_config', None)
+        super().__init__(**kwargs)
+
+class SafeConv2D(tf.keras.layers.Conv2D):
+    def __init__(self, **kwargs):
+        kwargs.pop('quantization_config', None)
+        super().__init__(**kwargs)
+
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 MODEL_PATH = os.path.join(BASE_DIR, 'alexnet_cifar10_keras.h5')
 
-print("Loading AlexNet from scratch...")
-model = tf.keras.models.load_model(MODEL_PATH)
+model = tf.keras.models.load_model(
+    MODEL_PATH, 
+    custom_objects={
+        'Dense': SafeDense, 
+        'Conv2D': SafeConv2D
+    }
+)
 
 conv_layers = [layer for layer in model.layers if isinstance(layer, tf.keras.layers.Conv2D)]
 feature_extractor = tf.keras.Model(inputs=model.inputs, outputs=[layer.output for layer in conv_layers])