Spaces:
Sleeping
Sleeping
File size: 5,932 Bytes
4179885 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
import gradio as gr
import numpy as np
from PIL import Image
import tempfile
import os
# Function to create the binary mask from the ImageEditor's output
def create_binary_mask(im_dict):
"""
Generates a binary mask from the drawing layer of the gr.ImageEditor output.
Args:
im_dict (dict): The dictionary output from gr.ImageEditor, containing
'background', 'layers', and 'composite'.
Returns:
tuple: A tuple containing:
- np.ndarray: The binary mask image (H, W) as a NumPy array (0 or 255).
- str or None: The filepath to the saved PNG mask for download, or None if no mask generated.
"""
if im_dict is None or im_dict["background"] is None:
print("No background image found.")
# Return a small blank placeholder and None for the file path
blank_preview = np.zeros((768, 1024), dtype=np.uint8)
return blank_preview, None
background_img = im_dict["background"]
h, w, _ = background_img.shape # Get original dimensions (Height, Width, Channels)
print(f"Original image dimensions: H={h}, W={w}")
# Check if any drawing layer exists and is not None
if not im_dict["layers"] or im_dict["layers"][0] is None:
print("No drawing layer found. Generating blank mask.")
# Nothing drawn yet, return a black mask of the original size
mask = np.zeros((h, w), dtype=np.uint8)
filepath = None # No file to download as nothing was drawn
else:
# Use the first layer (index 0) which usually contains the drawing
layer = im_dict["layers"][0]
print(f"Drawing layer dimensions: H={layer.shape[0]}, W={layer.shape[1]}")
# Ensure layer dimensions match background (Gradio ImageEditor usually handles this)
if layer.shape[0] != h or layer.shape[1] != w:
print(f"Warning: Layer size ({layer.shape[0]}x{layer.shape[1]}) doesn't match background ({h}x{w}). This shouldn't happen.")
# Handle potential mismatch if necessary, though unlikely with default editor behavior
# For now, proceed assuming they match or the layer is the correct reference
# Layer is RGBA, extract the Alpha channel (index 3)
alpha_channel = layer[:, :, 3]
# Create binary mask: white (255) where alpha > 0 (drawn), black (0) otherwise
mask = np.where(alpha_channel > 0, 255, 0).astype(np.uint8)
print(f"Generated binary mask dimensions: H={mask.shape[0]}, W={mask.shape[1]}")
# Save the mask to a temporary PNG file for download
try:
# Create a temporary file path
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
filepath = tmpfile.name
# Save the NumPy array as a PNG image using PIL
pil_image = Image.fromarray(mask)
pil_image.save(filepath, format="PNG")
print(f"Mask saved temporarily to: {filepath}")
except Exception as e:
print(f"Error saving mask to temporary file: {e}")
filepath = None # Indicate failure to save
# Return a blank mask in case of saving error
mask = np.zeros((h, w), dtype=np.uint8)
# Return the mask NumPy array for preview and the filepath for download
# The DownloadButton component will become active/functional if filepath is not None
return mask, filepath
# --- Gradio App Layout ---
with gr.Blocks() as demo:
gr.Markdown("## Binary Mask Generator")
gr.Markdown(
"Upload or paste an image. Use the brush tool (select it!) to draw the area "
"you want to mask. Click 'Generate Mask' to see the result and download it."
)
with gr.Row():
# --- Left Column ---
with gr.Column(scale=1): # Adjust scale as needed
image_editor = gr.ImageEditor(
label="Draw on Image",
# type="numpy" is essential for processing layers
type="numpy",
# DON'T set crop_size, height, or width to keep original dimensions
# sources allow upload, paste, webcam etc.
sources=["upload"],
# Set a default brush for clarity (optional, but helpful)
brush=gr.Brush(colors=["#FF0000"], color_mode="fixed"), # Red fixed brush
interactive=True,
canvas_size=(768, 1024)
)
generate_button = gr.Button("Generate Mask", variant="primary")
# --- Right Column ---
with gr.Column(scale=1): # Adjust scale as needed
mask_preview = gr.Image(
label="Binary Mask Preview",
# Use numpy for consistency, PIL would also work
type="numpy",
interactive=False, # Preview is not interactive
)
# Download button - its value (the file path) is set by the function's output
download_button = gr.DownloadButton(
label="Download Mask (PNG)",
interactive=True, # Button starts interactive
)
# --- Event Handling ---
generate_button.click(
fn=create_binary_mask,
inputs=[image_editor],
# Output 1 goes to mask_preview (image data)
# Output 2 goes to download_button (file path for the 'value' argument)
outputs=[mask_preview, download_button]
)
# --- Launch the App ---
if __name__ == "__main__":
# Cleaning up old temp files on startup (optional but good practice)
temp_dir = tempfile.gettempdir()
for item in os.listdir(temp_dir):
if item.endswith(".png") and item.startswith("tmp"): # Be specific to avoid deleting wrong files
try:
os.remove(os.path.join(temp_dir, item))
except Exception:
pass # Ignore if file is locked etc.
demo.launch()
|