obj detection seems ok, while segm is off

app.py

@@ -100,7 +100,7 @@ def run_prediction(image, confidence_threshold, resolution, model_name):
     image_input = image.resize((resolution, resolution))
     
     if model_name in SEGMENTATION_MODELS:
-        return run_segmentation(model, image_input)
+        return run_segmentation(model, image_input, image)
     else:
         return run_detection(model, image_input, image, confidence_threshold)
 
@@ -149,7 +149,6 @@ def run_detection(model, image_input, original_image, confidence_threshold):
     return annotated, analytics_text, {"count": len(boxes), "objects": class_counts}
 
 # Segm code opt 1
-'''
 def run_segmentation(model, image_input, original_image):
     mask_tensor = model.predict(image_input)
     mask_np = mask_tensor.cpu().numpy().astype(np.uint8)
@@ -196,8 +195,8 @@ def run_segmentation(model, image_input, original_image):
 
 
     return Image.fromarray(blended), analytics_text, {"classes_found": list(found_classes)}
-'''
 
+'''
 # Segm code opt 2
 def run_segmentation(model, image):
     """
@@ -225,7 +224,7 @@ def run_segmentation(model, image):
     blended = cv2.addWeighted(image_np, 0.6, colored_mask, 0.4, 0)
     return Image.fromarray(blended)
 
-
+'''
 # --- GRADIO UI ---
 
 theme = gr.themes.Soft(