Update src/ai_processor.py

src/ai_processor.py

@@ -210,61 +210,10 @@ def load_yolo_model():
         model = YOLO(YOLO_MODEL_PATH)
     return model
     
-def load_segmentation_model():
-    """
-    Robust TF/Keras loader that:
-      - Hides GPUs
-      - Tries plain load
-      - Falls back to patched InputLayer (drops batch_shape args)
-      - Tries SavedModel/.keras fallback
-    """
-    import tensorflow as tf
-    tf.config.set_visible_devices([], "GPU")
-
-    def _plain_load():
-        return tf.keras.models.load_model(SEG_MODEL_PATH, compile=False)
-
-    def _patched_load():
-        from tensorflow.keras.layers import InputLayer as KInputLayer
-
-        def _InputLayerPatched(*args, **kwargs):
-            kwargs.pop("batch_shape", None)
-            kwargs.pop("batch_input_shape", None)
-            return KInputLayer(**kwargs)
-
-        return tf.keras.models.load_model(
-            SEG_MODEL_PATH,
-            compile=False,
-            custom_objects={"InputLayer": _InputLayerPatched},
-        )
-
-    # 1) Try plain
-    try:
-        m = _plain_load()
-        logging.info("✅ Segmentation model loaded (plain).")
-        return m
-    except Exception as e1:
-        logging.warning(f"Plain load failed: {e1}")
-
-    # 2) Try patched InputLayer
-    try:
-        m = _patched_load()
-        logging.info("✅ Segmentation model loaded (patched InputLayer).")
-        return m
-    except Exception as e2:
-        logging.warning(f"Patched load failed: {e2}")
-
-    # 3) Try SavedModel/.keras folder fallback (strip extension)
-    try:
-        base, _ = os.path.splitext(SEG_MODEL_PATH)
-        m = tf.keras.models.load_model(base, compile=False)
-        logging.info("✅ Segmentation model loaded (SavedModel/.keras fallback).")
-        return m
-    except Exception as e3:
-        raise RuntimeError(f"Segmentation model could not be loaded: {e3}")
-
-
-
+def load_segmentation_model(seg_model_path):
+    """Lazy import and load segmentation model."""
+    from tensorflow.keras.models import load_model
+    return load_model(seg_model_path, compile=False)
 
 def load_classification_pipeline():
     pipe = _import_hf_cls()