add segment mask

app.py

@@ -1,7 +1,22 @@
 import gradio as gr
-from PIL import Image
+import numpy as np
+import mediapipe as mp
 import torch
+
+from PIL import Image
 from diffusers import AutoPipelineForImage2Image, DPMSolverMultistepScheduler
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+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
+
+segment_model = "checkpoints/selfie_multiclass_256x256.tflite"
+base_options = python.BaseOptions(model_asset_path=segment_model)
+options = vision.ImageSegmenterOptions(base_options=base_options,output_category_mask=True)
+segmenter = vision.ImageSegmenter.create_from_options(options)
+MASK_CATEGORY = segmenter.labels
 
 base_model = "SG161222/RealVisXL_V4.0"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -13,29 +28,64 @@ pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.
 pipeline.to(device)
 generator = torch.Generator(device).manual_seed(0)
 
-def image_to_image(input_image, prompt, guidance_scale, num_inference_steps):
-    # resize image to 512x512
-    input_image = input_image.resize((512, 512))
+def image_to_image(input_image, mask_image, prompt, negative_prompt, category, guidance_scale, num_inference_steps):
     # Generate the output image
     output_image = pipeline(
         generator=generator,
-        prompt=prompt, image=input_image, 
-        guidance_scale=guidance_scale, num_inference_steps = num_inference_steps
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=input_image, 
+        mask_image=mask_image,
+        guidance_scale=guidance_scale, 
+        num_inference_steps = num_inference_steps,
     ).images[0]
     
     return output_image
 
+def segment_image(input_image, prompt, negative_prompt, category, guidance_scale, num_inference_steps):
+    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_np = category_mask.numpy_view()
+    target_mask = category_mask_np == MASK_CATEGORY.index(category)
+
+    # 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
+
+    dilated_mask = binary_dilation(target_mask, iterations=4)
+    condition = np.stack((dilated_mask,) * 3, axis=-1) > 0.2
+
+    output_image = np.where(condition, fg_image, bg_image)
+    output_image = Image.fromarray(output_image)
+    return output_image
+
 with gr.Blocks() as grApp:
-    input_image = gr.Image(label="Input Image", type="pil")
-    prompt = gr.Textbox(lines=3, label="Prompt")
-    guidance_scale = gr.Slider(minimum=0, maximum=1, value=0.75, label="Guidance Scale")
-    num_inference_steps = gr.Slider(minimum=10, maximum=100, value=25, label="Number of Inference Steps")
-    output_image = gr.Image(label="Output Image", type="pil")
-    generate_btn = gr.Button("Generate Image")
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(lines=1, label="Prompt")
+            negative_prompt = gr.Textbox(lines=2, label="Negative Prompt")
+            category = gr.Dropdown(label='Mask Category', choices=MASK_CATEGORY, value=MASK_CATEGORY[1])
+            guidance_scale = gr.Slider(minimum=0, maximum=1, value=0.75, label="Guidance Scale")
+            num_inference_steps = gr.Slider(minimum=10, maximum=100, value=25, label="Number of Inference Steps")
+            input_image = gr.Image(label="Input Image", type="pil")
+            generate_btn = gr.Button("Generate Image")
+        with gr.Column():
+            mask_image = gr.Image(label="Mask Image", type="pil")
+        with gr.Column():
+            output_image = gr.Image(label="Output Image", type="pil")
+    
     generate_btn.click(
+        fn=segment_image,
+        inputs=[input_image, prompt, category, guidance_scale, num_inference_steps],
+        outputs=[mask_image],
+    ).then(
         fn=image_to_image,
-        inputs=[input_image, prompt, guidance_scale, num_inference_steps],
-        outputs=output_image,
+        inputs=[input_image, mask_image, prompt, negative_prompt, category, guidance_scale, num_inference_steps],
+        outputs=[output_image],
     )
 
 grApp.launch()

checkpoints/selfie_multiclass_256x256.tflite

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

requirements.txt

@@ -2,4 +2,5 @@ gradio
 torch
 diffusers
 transformers
-accelerate
+accelerate
+mediapipe