Update src/ai_processor.py

src/ai_processor.py

@@ -255,57 +255,10 @@ def load_yolo_model():
         model = YOLO(YOLO_MODEL_PATH)
     return model
 
-def load_segmentation_model(path: Optional[str] = None):
-    """
-    Robust loader for legacy .h5 models across TF/Keras versions.
-    Uses global SEG_MODEL_PATH by default.
-    """
-    import ast
+def load_segmentation_model():
     import tensorflow as tf
-    tf.config.set_visible_devices([], "GPU")
-    model_path = path or SEG_MODEL_PATH
-
-    # Attempt 1: tf.keras with safe_mode=False
-    try:
-        m = tf.keras.models.load_model(model_path, compile=False, safe_mode=False)
-        logging.info("✅ Segmentation model loaded (tf.keras, safe_mode=False).")
-        return m
-    except Exception as e1:
-        logging.warning(f"tf.keras load (safe_mode=False) failed: {e1}")
-
-    # Attempt 2: patched InputLayer (drop legacy args; coerce string shapes)
-    try:
-        from tensorflow.keras.layers import InputLayer as _KInputLayer
-        def _InputLayerPatched(*args, **kwargs):
-            kwargs.pop("batch_shape", None)
-            kwargs.pop("batch_input_shape", None)
-            if "shape" in kwargs and isinstance(kwargs["shape"], str):
-                try:
-                    kwargs["shape"] = tuple(ast.literal_eval(kwargs["shape"]))
-                except Exception:
-                    kwargs.pop("shape", None)
-            return _KInputLayer(**kwargs)
-        m = tf.keras.models.load_model(
-            model_path,
-            compile=False,
-            custom_objects={"InputLayer": _InputLayerPatched},
-            safe_mode=False,
-        )
-        logging.info("✅ Segmentation model loaded (patched InputLayer).")
-        return m
-    except Exception as e2:
-        logging.warning(f"Patched InputLayer load failed: {e2}")
-
-    # Attempt 3: keras 2 shim (tf_keras) if present
-    try:
-        import tf_keras
-        m = tf_keras.models.load_model(model_path, compile=False)
-        logging.info("✅ Segmentation model loaded (tf_keras compat).")
-        return m
-    except Exception as e3:
-        logging.warning(f"tf_keras load failed or not installed: {e3}")
-
-    raise RuntimeError("Segmentation model could not be loaded; please convert/resave the model.")
+    load_model = _import_tf_loader()
+    return load_model(SEG_MODEL_PATH, compile=False, custom_objects={'InputLayer': tf.keras.layers.InputLayer})
 
 def load_classification_pipeline():
     pipe = _import_hf_cls()