From 05535db5f77bdc9c46df36855fe4064ded89d7cb Mon Sep 17 00:00:00 2001
From: Sebastian Schmidt <mrschmidt@google.com>
Date: Thu, 8 Dec 2022 15:01:34 -0800
Subject: [PATCH] Fix assertion failure in Hair Segmentation demo

PiperOrigin-RevId: 494004801
---
 mediapipe/graphs/hair_segmentation/BUILD                 | 1 +
 .../hair_segmentation_desktop_live.pbtxt                 | 9 ++++++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/mediapipe/graphs/hair_segmentation/BUILD b/mediapipe/graphs/hair_segmentation/BUILD
index b177726bf..945f02c62 100644
--- a/mediapipe/graphs/hair_segmentation/BUILD
+++ b/mediapipe/graphs/hair_segmentation/BUILD
@@ -43,6 +43,7 @@ cc_library(
     deps = [
         "//mediapipe/calculators/core:flow_limiter_calculator",
         "//mediapipe/calculators/core:previous_loopback_calculator",
+        "//mediapipe/calculators/image:color_convert_calculator",
         "//mediapipe/calculators/image:image_transformation_calculator",
         "//mediapipe/calculators/image:recolor_calculator",
         "//mediapipe/calculators/image:set_alpha_calculator",
diff --git a/mediapipe/graphs/hair_segmentation/hair_segmentation_desktop_live.pbtxt b/mediapipe/graphs/hair_segmentation/hair_segmentation_desktop_live.pbtxt
index 36c6970e1..f48b26be0 100644
--- a/mediapipe/graphs/hair_segmentation/hair_segmentation_desktop_live.pbtxt
+++ b/mediapipe/graphs/hair_segmentation/hair_segmentation_desktop_live.pbtxt
@@ -60,7 +60,14 @@ node {
     tag_index: "LOOP"
     back_edge: true
   }
-  output_stream: "PREV_LOOP:previous_hair_mask"
+  output_stream: "PREV_LOOP:previous_hair_mask_rgb"
+}
+
+# Converts the 4 channel hair mask to a single channel mask
+node {
+  calculator: "ColorConvertCalculator"
+  input_stream: "RGB_IN:previous_hair_mask_rgb"
+  output_stream: "GRAY_OUT:previous_hair_mask"
 }
 
 # Embeds the hair mask generated from the previous round of hair segmentation