|
|
|
|
|
import gradio as gr |
|
|
import numpy as np |
|
|
from PIL import Image |
|
|
import os |
|
|
|
|
|
|
|
|
def generate_mask(input_dict): |
|
|
|
|
|
|
|
|
original_image = input_dict['image'] |
|
|
sketch_mask = input_dict['mask'] |
|
|
if sketch_mask.mode == 'RGBA': |
|
|
alpha_channel = sketch_mask.split()[3] |
|
|
binary_mask_np = np.array(alpha_channel) |
|
|
binary_mask_np = np.where(binary_mask_np > 0, 255, 0).astype(np.uint8) |
|
|
binary_mask_pil = Image.fromarray(binary_mask_np, mode='L') |
|
|
return binary_mask_pil |
|
|
else: |
|
|
return sketch_mask |
|
|
|
|
|
|
|
|
iface = gr.Interface( |
|
|
fn=generate_mask, |
|
|
inputs=gr.Image(type="pil", tool="sketch", label="Draw Mask on Image", source="upload"), |
|
|
outputs=gr.Image(type="pil", label="Generated Mask"), |
|
|
title="Image Mask Generator", |
|
|
description="Upload an image and draw on it to create a mask." |
|
|
) |
|
|
|
|
|
|
|
|
iface.launch() |