Hair segmentation issue in python

[Rameshv]

I am trying to use this hair segmentation model https://storage.googleapis.com/mediapipe-assets/hair_segmentation.tflite in python and I was following the python guide from this https://developers.google.com/mediapipe/solutions/vision/image_segmenter/python

And using the same notebook mentioned in the link https://colab.research.google.com/github/googlesamples/mediapipe/blob/main/examples/image_segmentation/python/image_segmentation.ipynb

I am able to download the model with

!wget -O hair.tflite -q https://storage.googleapis.com/mediapipe-assets/hair_segmentation.tflite

but when trying to run the inference, using the below code

base_options = BaseOptions(model_asset_path='hair.tflite')
options = vision.ImageSegmenterOptions(base_options=base_options,running_mode=VisionRunningMode.IMAGE,
                                              output_type=OutputType.CATEGORY_MASK)
 x = vision.ImageSegmenter.create_from_options(options)

I am getting the error on line vision.ImageSegmenter.create_from_options(options)

The input tensor should have dimensions 1 x height x width x 3. Got 1 x 512 x 512 x 4.

I suspect this is due to tensor size mismatch between the hair model and Vision image segmenter. But I could't find any references to that and stuck.

heres the complete stack trace

[kuaashish]

Hi @Rameshv,
Thank you for raising a fresh issue, this is already known to us and we are working towards fix. As workaround, you can load the file manually and use model_asset_buffer: https://github.com/google/mediapipe/blob/master/mediapipe/tasks/python/core/base_options.py#L47

Please let us know if the above resolves the issue. Thank you

[kuaashish]

Hi @Rameshv,
Could you please let us know the status of the above if the issue has been resolved or further need help from our end. Thank you!

[Rameshv]

Thanks @kuaashish for the response

But I still got the same issue. I have loaded hair_segmentation.tflite and used it in model_asset_buffer, but its still the same issue

[kuaashish]

@Rameshv,
Would you please let us know the information about the OS using, Python Version to reproduce the issue from our end? We suggest to fill the issue template to receive maximum relevant info to understand the issue better and share the possible solution. Thank you

[Rameshv]

Hi @kuaashish
I am running this notebook https://colab.research.google.com/github/googlesamples/mediapipe/blob/main/examples/image_segmentation/python/image_segmentation.ipynb in google colab. so the system is Ubuntu 20.04 LTS and python version is 3.9.16

And here are the dependencies as per the notebook

!pip install -q flatbuffers==2.0.0
!pip install -q mediapipe==0.9.1

and the code

import numpy as np
import mediapipe as mp

from mediapipe.python._framework_bindings import image
from mediapipe.python._framework_bindings import image_frame
from mediapipe.tasks.python import vision
from mediapipe import tasks

BG_COLOR = (192, 192, 192) # gray
MASK_COLOR = (255, 255, 255) # white

OutputType = vision.ImageSegmenterOptions.OutputType
Activation = vision.ImageSegmenterOptions.Activation
VisionRunningMode = vision.RunningMode
BaseOptions =tasks.BaseOptions

IMAGE_FILENAMES = ['13147680-fd2e-4d67-b29a-89b9a6752703.jpeg']

in_file = open("hair.tflite", "rb")
data = in_file.read()

# Create the options that will be used for ImageSegmenter
base_options = BaseOptions(model_asset_buffer=data)
options = vision.ImageSegmenterOptions(base_options=base_options,running_mode=VisionRunningMode.IMAGE,
                                              output_type=OutputType.CATEGORY_MASK)

# Create the image segmenter
with vision.ImageSegmenter.create_from_options(options) as segmenter:

  # Loop through demo image(s)
  for image_file_name in IMAGE_FILENAMES:

    # Create the MediaPipe image file that will be segmented
    image = mp.Image.create_from_file(image_file_name)

    # Retrieve the masks for the segmented image
    category_masks = segmenter.segment(image)
    # print(category_masks)

    # Generate solid color images for showing the output segmentation mask.
    image_data = image.numpy_view()
    fg_image = np.zeros(image_data.shape, dtype=np.uint8)
    fg_image[:] = MASK_COLOR
    bg_image = np.zeros(image_data.shape, dtype=np.uint8)
    bg_image[:] = BG_COLOR

    for name in category_masks:
      condition = np.stack((name.numpy_view(),) * 3, axis=-1) > 0.2
      output_image = np.where(condition, fg_image, bg_image)

      print(f'Segmentation mask of {name}:')
      resize_and_show(output_image)

its the same code I used deeplabv3.tflite and selfie_segm_128_128_3.tflite. Both works fine. Only the hair_segmentation.tflite model throws the error. So I dont think this is related to a platform.

hope this helps

[ayushgdev]

Hello @Rameshv Your suspicion is correct. If you check hair segmentation solution graph here, it needs number of channels as 4. That is why the error shows the actual shape it got is 1x512x512x4. It infers this from the tflite file.

However, the ImageSegmenter API explicitly states the channels should be 1 or 3. And the batch can only be 1. Hence the error states the expected input was 1 x any height x any width x 3. If you check metadata for DeepLab v3 model, it also expects input to be 257x257x3. Hence it works.

[Rameshv]

Thanks @ayushgdev for the clarification.

So what do you think is better way to approach this? Is it all possible to do with ImageSegmenter API.??

any pointers to solve this would be really appreciated. I am new to python but I can start digging if I am shown a path..

thanks

[ayushgdev]

@Rameshv There is a possible solution for it as per analysis. Please allow us some more time to investigate on the issue. Apologies for the delay.

[ayushgdev]

Hello @Rameshv This issue has been fixed in MediaPipe v0.9.3.0. Please upgrade the package and it should start working.

[Rameshv]

Thanks @ayushgdev
Will download and check in a bit

[Rameshv]

I just tried the MediaPipe v0.9.3.0 and the error is gone and it can produce the hair segmentation. But the issue now is, its not producing the correct segmentation.

heres the source image I tried

And heres the hair segmentation result

And the same image producing correct segmentation with selfie segmentaion

And with deeplabv3

As you can see both the selfie and deeplabv3 segmentaion produce correct results, but the hair segmentation is not.

any ideas what could be the issue @ayushgdev ??

[WillReynolds5]

Same problem as above, bad mask on every image i try.

[ayushgdev]

Hello @Rameshv hair segmentation will only produce a mask for the hair. While in contrast, the selfie segmentation will produce mask for human body. Hence both are different applications and incomparable.

[Rameshv]

Hi @ayushgdev I think you misunderstood me. What I said was the hair segmentation doesnt produce correct hair mask and not it should be same as self segmentation. if you can look at the hair mask produced, it is completely different from the source image.

This is the expected hair segment

but this is what it produced

Hope this helps.

[ayushgdev]

Hello @Rameshv
Thanks for bringing this issue. Answer to this will help others as well. As conveyed earlier, hair segmentation requires 4 channels while the output of mp.Image.create_from_file(image_file_name) creates an image representation with 3 channels. Hence, the last channel is treated as alpha channel and based on that, the computation happens.
To resolve this issue, we can add an empty alpha channel to the image.
So you can simply use opencv to read the image, add alpha channel, and then create a MediaPipe Image object using mp.Image() constructor.

#  image = mp.Image.create_from_file(image_file_name)      <----------- Comment this out

# read image using OpenCV
rgb_image = cv2.imread("person.jpeg")
rgba_image = cv2.cvtColor(rgb_image, cv2.COLOR_BGR2RGBA)
# set alpha channel to empty.
rgba_image[:,:,3] = 0
# create MP Image object from numpy array
image = _Image(image_format=_ImageFormat.SRGBA, data=rgba_image)

The output produced is below:

[mariomosko]

@ayushgdev
Is there any chance that you could show the entire implementation the _Image part is not very clear.
Thank you in advance

[ayushgdev]

_Image is essentially Image class. Check the imports below:

from mediapipe.python._framework_bindings import image as image_module
_Image = image_module.Image

Other code remains the same as this/image segmentation colab example.

[mariomosko]

Hello , thank you for your reply
,can you help with _ImageFormat as it is also not defined
thank you in advance!

[ayushgdev]

No problem. ImageFormat is imported as below

from mediapipe.python._framework_bindings import image_frame
_ImageFormat = image_frame.ImageFormat

[mariomosko]

@ayushgdev thank you

for me personally now it works!
thank you so much

[github-actions]

This issue has been marked stale because it has no recent activity since 7 days. It will be closed if no further activity occurs. Thank you.

[github-actions]

This issue was closed due to lack of activity after being marked stale for past 7 days.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[Gauravi1]

@Rameshv can you please provide me full code of hair segmentation