change to selfie_multiclass

app.py

@@ -9,15 +9,17 @@ from scipy.ndimage import binary_dilation
 BG_COLOR = np.array([0, 0, 0], dtype=np.uint8) # black
 MASK_COLOR = np.array([255, 255, 255], dtype=np.uint8) # white
 
-MODEL_PATH = "checkpoints/hair_segmenter.tflite"
+MODEL_PATH = "checkpoints/selfie_multiclass_256x256.tflite"
 base_options = python.BaseOptions(model_asset_path=MODEL_PATH)
 options = vision.ImageSegmenterOptions(base_options=base_options,output_category_mask=True)
 segmenter = vision.ImageSegmenter.create_from_options(options)
+MASK_CATEGORY = segmenter.labels
 
-def segment(input_image):
+def segment(input_image, category):
     image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(input_image))
     segmentation_result = segmenter.segment(image)
     category_mask = segmentation_result.category_mask
+    category_mask = category_mask == MASK_CATEGORY.index(category)
 
     # Generate solid color images for showing the output segmentation mask.
     image_data = image.numpy_view()
@@ -37,6 +39,7 @@ with gr.Blocks() as app:
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type='pil', label='Upload image')
+            category = gr.Dropdown(label='Category', choices=MASK_CATEGORY, value=MASK_CATEGORY[1])
             submit_btn = gr.Button(value='Submit', variant='primary')
         with gr.Column():
             output_image = gr.Image(type='pil', label='Image Output')
@@ -45,6 +48,7 @@ with gr.Blocks() as app:
         fn=segment,
         inputs=[
             input_image,
+            category,
         ],
         outputs=[output_image]
     )

checkpoints/selfie_multiclass_256x256.tflite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6748b1253a99067ef71f7e26ca71096cd449baefa8f101900ea23016507e0e0
+size 16371837