File size: 5,132 Bytes
8fa1f84 |
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 137 138 139 140 141 142 143 144 145 |
import cv2
from PIL import Image
import numpy as np
import torch
import gradio as gr
class ImageSketcher(gr.Image):
"""
Code is from https://github.com/jshilong/GPT4RoI/blob/7c157b5f33914f21cfbc804fb301d3ce06324193/gpt4roi/app.py#L365
Fix the bug of gradio.Image that cannot upload with tool == 'sketch'.
"""
is_template = True # Magic to make this work with gradio.Block, don't remove unless you know what you're doing.
def __init__(self, **kwargs):
super().__init__(tool='boxes', **kwargs)
def preprocess(self, x):
if x is None:
return x
if self.tool == 'boxes' and self.source in ['upload', 'webcam']:
if isinstance(x, str):
x = {'image': x, 'boxes': []}
else:
assert isinstance(x, dict)
assert isinstance(x['image'], str)
assert isinstance(x['boxes'], list)
x = super().preprocess(x)
return x
def process_mask_to_show(mask):
'''
Process the mask to show on the gradio.Image
'''
mask = np.array(mask > 0.1, dtype=np.uint8) * 255
mask_stacked = np.stack([mask] * 3, axis=-1)
return mask_stacked
def img_add_masks(img_, colored_mask, mask, linewidth=2):
if type(img_) is np.ndarray:
img = Image.fromarray(img_, mode='RGB').convert('RGBA')
else:
img = img_.copy()
h, w = img.height, img.width
# contour
temp = np.zeros((h, w, 1))
contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(temp, contours, -1, (255, 255, 255), linewidth)
color = np.array([1, 1, 1, 1])
contour_mask = temp * color.reshape(1, 1, -1)
overlay_inner = Image.fromarray(colored_mask.astype(np.uint8), 'RGBA')
img.paste(overlay_inner, (0, 0), overlay_inner)
overlay_contour = Image.fromarray(contour_mask.astype(np.uint8), 'RGBA')
img.paste(overlay_contour, (0, 0), overlay_contour)
return img
def gen_colored_masks(
annotation,
random_color=False,
):
"""
Code is largely based on https://github.com/CASIA-IVA-Lab/FastSAM/blob/4d153e909f0ad9c8ecd7632566e5a24e21cf0071/utils/tools_gradio.py#L130
"""
device = annotation.device
mask_sum = annotation.shape[0]
height = annotation.shape[1]
weight = annotation.shape[2]
areas = torch.sum(annotation, dim=(1, 2))
sorted_indices = torch.argsort(areas, descending=False)
annotation = annotation[sorted_indices]
index = (annotation != 0).to(torch.long).argmax(dim=0)
if random_color:
color = torch.rand((mask_sum, 1, 1, 3)).to(device)
else:
color = torch.ones((mask_sum, 1, 1, 3)).to(device) * torch.tensor(
[30 / 255, 144 / 255, 255 / 255]
).to(device)
transparency = torch.ones((mask_sum, 1, 1, 1)).to(device) * 0.6
visual = torch.cat([color, transparency], dim=-1)
mask_image = torch.unsqueeze(annotation, -1) * visual
mask = torch.zeros((height, weight, 4)).to(device)
h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight))
indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
mask[h_indices, w_indices, :] = mask_image[indices]
mask_cpu = mask.cpu().numpy()
return mask_cpu, sorted_indices
def mask_foreground(mask, trans=0.6, random_color=True):
if random_color:
color = np.concatenate([np.random.random(3) * 255, np.array([trans * 255])], axis=0)
else:
color = np.array([30, 144, 255, trans * 255])
h, w = mask.shape[-2:]
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
return mask_image
def mask_background(mask, trans=0.5):
h, w = mask.shape[-2:]
mask_image = (1 - mask.reshape(h, w, 1)) * np.array([0, 0, 0, trans * 255])
return mask_image
def mask_select_point(all_masks, output_mask_2_raw, mask_order, evt: gr.SelectData):
h, w = output_mask_2_raw.height, output_mask_2_raw.width
pointed_mask = None
for i in range(len(mask_order)):
idx = mask_order[i]
msk = all_masks[idx]
if msk[evt.index[1], evt.index[0]] == 1:
pointed_mask = msk.copy()
break
if pointed_mask is not None:
contours, hierarchy = cv2.findContours(pointed_mask.astype("uint8"), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
ret = output_mask_2_raw.copy()
temp = np.zeros((h, w, 1))
contours, _ = cv2.findContours(msk.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(temp, contours, -1, (255, 255, 255), 3)
color = np.array([1, 1, 1, 1])
contour_mask = temp * color.reshape(1, 1, -1)
colored_mask = mask_background(pointed_mask)
overlay_inner = Image.fromarray(colored_mask.astype(np.uint8), 'RGBA')
ret.paste(overlay_inner, (0, 0), overlay_inner)
overlay_contour = Image.fromarray(contour_mask.astype(np.uint8), 'RGBA')
ret.paste(overlay_contour, (0, 0), overlay_contour)
return ret, pointed_mask
else:
return output_mask_2_raw, None
|