Update src/ai_processor.py

src/ai_processor.py

@@ -210,13 +210,10 @@ def load_yolo_model():
         model = YOLO(YOLO_MODEL_PATH)
     return model
     
-def load_segmentation_model(SEG_MODEL_PATH):
-    """Lazy import and load segmentation model."""
-    # Use the lazy loader to ensure TF is imported correctly
+def load_segmentation_model():
+    import tensorflow as tf
     load_model = _import_tf_loader()
-    # The compile=False argument is crucial for avoiding optimizer/loss issues
-    # and can sometimes help with layer deserialization.
-    return load_model(SEG_MODEL_PATH, compile=False)
+    return load_model(SEG_MODEL_PATH, compile=False, custom_objects={'InputLayer': tf.keras.layers.InputLayer})
 
 
 def load_classification_pipeline():