|
import gradio as gr |
|
import numpy as np |
|
from PIL import Image, ImageDraw |
|
|
|
|
|
class ImageMask(gr.components.Image): |
|
""" |
|
Sets: source="canvas", tool="sketch" |
|
""" |
|
|
|
is_template = True |
|
|
|
def __init__(self, **kwargs): |
|
super().__init__(source="upload", |
|
tool="sketch", |
|
interactive=False, |
|
**kwargs) |
|
|
|
def preprocess(self, x): |
|
if x is None: |
|
return x |
|
if self.tool == "sketch" and self.source in ["upload", "webcam" |
|
] and type(x) != dict: |
|
decode_image = gr.processing_utils.decode_base64_to_image(x) |
|
width, height = decode_image.size |
|
mask = np.ones((height, width, 4), dtype=np.uint8) |
|
mask[..., -1] = 255 |
|
mask = self.postprocess(mask) |
|
x = {'image': x, 'mask': mask} |
|
return super().preprocess(x) |
|
|
|
|
|
def get_valid_mask(mask: np.ndarray): |
|
"""Convert mask from gr.Image(0 to 255, RGBA) to binary mask. |
|
""" |
|
if mask.ndim == 3: |
|
mask_pil = Image.fromarray(mask).convert('L') |
|
mask = np.array(mask_pil) |
|
if mask.max() == 255: |
|
mask = mask / 255 |
|
return mask |
|
|
|
|
|
def draw_points_on_image(image, |
|
points, |
|
curr_point=None, |
|
highlight_all=True, |
|
radius_scale=0.01): |
|
overlay_rgba = Image.new("RGBA", image.size, 0) |
|
overlay_draw = ImageDraw.Draw(overlay_rgba) |
|
for point_key, point in points.items(): |
|
if ((curr_point is not None and curr_point == point_key) |
|
or highlight_all): |
|
p_color = (255, 0, 0) |
|
t_color = (0, 0, 255) |
|
|
|
else: |
|
p_color = (255, 0, 0, 35) |
|
t_color = (0, 0, 255, 35) |
|
|
|
rad_draw = int(image.size[0] * radius_scale) |
|
|
|
p_start = point.get("start_temp", point["start"]) |
|
p_target = point["target"] |
|
|
|
if p_start is not None and p_target is not None: |
|
p_draw = int(p_start[0]), int(p_start[1]) |
|
t_draw = int(p_target[0]), int(p_target[1]) |
|
|
|
overlay_draw.line( |
|
(p_draw[0], p_draw[1], t_draw[0], t_draw[1]), |
|
fill=(255, 255, 0), |
|
width=2, |
|
) |
|
|
|
if p_start is not None: |
|
p_draw = int(p_start[0]), int(p_start[1]) |
|
overlay_draw.ellipse( |
|
( |
|
p_draw[0] - rad_draw, |
|
p_draw[1] - rad_draw, |
|
p_draw[0] + rad_draw, |
|
p_draw[1] + rad_draw, |
|
), |
|
fill=p_color, |
|
) |
|
|
|
if curr_point is not None and curr_point == point_key: |
|
|
|
overlay_draw.text(p_draw, "p", align="center", fill=(0, 0, 0)) |
|
|
|
if p_target is not None: |
|
t_draw = int(p_target[0]), int(p_target[1]) |
|
overlay_draw.ellipse( |
|
( |
|
t_draw[0] - rad_draw, |
|
t_draw[1] - rad_draw, |
|
t_draw[0] + rad_draw, |
|
t_draw[1] + rad_draw, |
|
), |
|
fill=t_color, |
|
) |
|
|
|
if curr_point is not None and curr_point == point_key: |
|
|
|
overlay_draw.text(t_draw, "t", align="center", fill=(0, 0, 0)) |
|
|
|
return Image.alpha_composite(image.convert("RGBA"), |
|
overlay_rgba).convert("RGB") |
|
|
|
|
|
def draw_mask_on_image(image, mask): |
|
im_mask = np.uint8(mask * 255) |
|
im_mask_rgba = np.concatenate( |
|
( |
|
np.tile(im_mask[..., None], [1, 1, 3]), |
|
45 * np.ones( |
|
(im_mask.shape[0], im_mask.shape[1], 1), dtype=np.uint8), |
|
), |
|
axis=-1, |
|
) |
|
im_mask_rgba = Image.fromarray(im_mask_rgba).convert("RGBA") |
|
|
|
return Image.alpha_composite(image.convert("RGBA"), |
|
im_mask_rgba).convert("RGB") |
|
|
|
|
|
def on_change_single_global_state(keys, |
|
value, |
|
global_state, |
|
map_transform=None): |
|
if map_transform is not None: |
|
value = map_transform(value) |
|
|
|
curr_state = global_state |
|
if isinstance(keys, str): |
|
last_key = keys |
|
|
|
else: |
|
for k in keys[:-1]: |
|
curr_state = curr_state[k] |
|
|
|
last_key = keys[-1] |
|
|
|
curr_state[last_key] = value |
|
return global_state |
|
|
|
|
|
def get_latest_points_pair(points_dict): |
|
if not points_dict: |
|
return None |
|
point_idx = list(points_dict.keys()) |
|
latest_point_idx = max(point_idx) |
|
return latest_point_idx |
|
|
