Fix loading nested-Sequential .h5 exports under Keras 3

src/teachable_machine.py

@@ -1,3 +1,5 @@
+import json
+
 from tensorflow.keras.models import load_model
 from tensorflow.keras.layers import DepthwiseConv2D
 from PIL import Image, ImageOps, ImageDraw, ImageFont
@@ -27,7 +29,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
 
-class TeachableMachine(object):
+class TeachableMachine:
     """
     Create a TeachableMachine object to run pre-trained AI models.
     """
@@ -59,13 +61,162 @@ def _load_model(self, model_path: str):
                 compile=False,
                 custom_objects={"DepthwiseConv2D": _CompatDepthwiseConv2D},
             )
+            return
         except IOError as e:
             print("LoadingModelError: Error while loading Teachable Machine model")
             raise IOError("Error loading model") from e
         except Exception as e:
+            # Some Teachable Machine .h5 exports save the model as a Sequential
+            # that wraps nested Sequential/Functional submodels (Sequential >
+            # [InputLayer, Sequential > [InputLayer, Functional backbone,
+            # GlobalAveragePooling2D], Sequential > [InputLayer, Dense, Dense]]).
+            # Keras 3's legacy H5 loader mis-rebuilds that shape: when the
+            # Functional backbone is called as a plain layer inside
+            # Sequential.build(), its single output comes back wrapped in a
+            # length-1 list, which the next layer then treats as two separate
+            # inputs, failing with e.g. "expects 1 input(s), but it received 2
+            # input tensors". Fall back to a manual reconstruction for that case.
+            try:
+                self._model = self._load_nested_sequential_h5(model_path)
+                return
+            except Exception:
+                pass
             print("LoadingModelError: Error while loading Teachable Machine model")
             raise FileNotFoundError("Model file not found") from e
 
+    def _load_nested_sequential_h5(self, model_path: str):
+        """
+        Rebuild a Teachable Machine .h5 export whose Sequential wraps nested
+        Sequential/Functional submodels, working around the Keras 3 legacy H5
+        loader bug described in `_load_model`.
+
+        Flattens the nested Sequential wrappers (keeping only the first
+        InputLayer and leaving any Functional submodels intact), rebuilds the
+        layers in order via the Functional API -- unwrapping the length-1
+        list a Functional submodel's call returns -- and then restores
+        weights by matching each leaf layer's name against the H5 file's
+        saved weight paths, sidestepping Keras's built-in weight loader
+        (which expects the now-discarded nested-Sequential grouping).
+        """
+        import h5py
+        from keras import Model
+        from keras.layers import Input
+        from keras.src.legacy.saving import saving_utils as legacy_saving_utils
+
+        custom_objects = {"DepthwiseConv2D": _CompatDepthwiseConv2D}
+
+        with h5py.File(model_path, "r") as f:
+            raw_config = f.attrs.get("model_config")
+            if raw_config is None:
+                raise ValueError("No model_config found in H5 file")
+            if isinstance(raw_config, bytes):
+                raw_config = raw_config.decode("utf-8")
+            full_config = json.loads(raw_config)
+
+        if full_config.get("class_name") != "Sequential":
+            raise ValueError("Not a nested-Sequential export; nothing to flatten")
+
+        flat_layers = self._flatten_sequential_layers(full_config)
+        if not flat_layers or flat_layers[0]["class_name"] != "InputLayer":
+            raise ValueError("Unexpected model shape after flattening")
+
+        input_shape = flat_layers[0]["config"]["batch_input_shape"][1:]
+        x = Input(shape=input_shape)
+        model_input = x
+        leaf_layers = []
+        for layer_config in flat_layers[1:]:
+            if layer_config["class_name"] == "Functional":
+                sub_model = Model.from_config(
+                    layer_config["config"], custom_objects=custom_objects
+                )
+                x = sub_model(x)
+                if isinstance(x, (list, tuple)) and len(x) == 1:
+                    x = x[0]
+                leaf_layers.extend(sub_model.layers)
+            else:
+                layer = legacy_saving_utils.model_from_config(
+                    layer_config, custom_objects=custom_objects
+                )
+                x = layer(x)
+                leaf_layers.append(layer)
+
+        model = Model(inputs=model_input, outputs=x)
+        self._restore_weights_by_leaf_name(model_path, leaf_layers)
+        return model
+
+    @staticmethod
+    def _flatten_sequential_layers(layer_config, _seen_input=None):
+        """
+        Inline nested Sequential layer lists, keeping only the first
+        InputLayer seen overall and leaving Functional submodels intact.
+        """
+        if _seen_input is None:
+            _seen_input = [False]
+        if layer_config["class_name"] != "Sequential":
+            return [layer_config]
+        flattened = []
+        for sub_layer_config in layer_config["config"]["layers"]:
+            if sub_layer_config["class_name"] == "InputLayer":
+                if not _seen_input[0]:
+                    flattened.append(sub_layer_config)
+                    _seen_input[0] = True
+                continue
+            flattened.extend(
+                TeachableMachine._flatten_sequential_layers(
+                    sub_layer_config, _seen_input
+                )
+            )
+        return flattened
+
+    @staticmethod
+    def _restore_weights_by_leaf_name(model_path, leaf_layers):
+        """
+        Load weights for `leaf_layers` by matching each layer's name against
+        the H5 file's saved weight paths, regardless of the (now-discarded)
+        nested-Sequential grouping those paths were originally saved under.
+        """
+        import h5py
+
+        with h5py.File(model_path, "r") as f:
+            root = f["model_weights"]
+            leaf_paths = {}
+
+            def walk(group):
+                """Recurse into `group`, indexing every leaf layer's saved
+                weight paths (from its 'weight_names' attribute) by leaf
+                layer name."""
+                weight_names = group.attrs.get("weight_names")
+                if weight_names is not None:
+                    for weight_name in weight_names:
+                        if isinstance(weight_name, bytes):
+                            weight_name = weight_name.decode("utf-8")
+                        full_path = f"{group.name}/{weight_name}"
+                        leaf_name = (
+                            weight_name.rsplit("/", 1)[0]
+                            if "/" in weight_name
+                            else group.name.rsplit("/", 1)[-1]
+                        )
+                        leaf_paths.setdefault(leaf_name, []).append(full_path)
+                for key in group:
+                    item = group[key]
+                    if isinstance(item, h5py.Group):
+                        walk(item)
+
+            walk(root)
+
+            missing = []
+            for layer in leaf_layers:
+                if not layer.weights:
+                    continue
+                paths = leaf_paths.get(layer.name)
+                if paths is None:
+                    missing.append(layer.name)
+                    continue
+                layer.set_weights([np.asarray(f[path]) for path in paths])
+
+            if missing:
+                raise ValueError(f"Could not find saved weights for layers: {missing}")
+
     def _load_labels(self, labels_file_path):
         try:
             with open(labels_file_path, "r") as file: