|
import os, sys |
|
import random |
|
import time |
|
import warnings |
|
from scipy.ndimage import binary_dilation |
|
from PIL import Image, ImageDraw, ImageFont |
|
|
|
def downloadStuff(): |
|
|
|
|
|
os.system("pip install --upgrade diffusers[torch]") |
|
|
|
|
|
|
|
|
|
|
|
sys.path.append(os.path.join(os.getcwd(), "GroundingDINO")) |
|
sys.path.append(os.path.join(os.getcwd(), "segment_anything")) |
|
|
|
|
|
|
|
warnings.filterwarnings("ignore") |
|
|
|
import gradio as gr |
|
import argparse |
|
|
|
import numpy as np |
|
import torch |
|
import torchvision |
|
from PIL import Image, ImageDraw, ImageFont |
|
|
|
|
|
from .GroundingDINO.groundingdino.datasets import transforms as T |
|
from .GroundingDINO.groundingdino.models import build_model |
|
from .GroundingDINO.groundingdino.util.slconfig import SLConfig |
|
from .GroundingDINO.groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap |
|
|
|
|
|
|
|
from .segment_anything.segment_anything.build_sam import build_sam |
|
from .segment_anything.segment_anything.predictor import SamPredictor |
|
import numpy as np |
|
|
|
|
|
import torch |
|
from diffusers import StableDiffusionInpaintPipeline |
|
|
|
|
|
from transformers import BlipProcessor, BlipForConditionalGeneration |
|
|
|
|
|
def generate_caption(processor, blip_model, raw_image): |
|
|
|
inputs = processor(raw_image, return_tensors="pt").to( |
|
"cuda", torch.float16) |
|
out = blip_model.generate(**inputs) |
|
caption = processor.decode(out[0], skip_special_tokens=True) |
|
return caption |
|
|
|
|
|
def transform_image(image_pil): |
|
|
|
transform = T.Compose( |
|
[ |
|
T.RandomResize([800], max_size=1333), |
|
T.ToTensor(), |
|
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
|
] |
|
) |
|
image, _ = transform(image_pil, None) |
|
return image |
|
|
|
|
|
def load_model(model_config_path, model_checkpoint_path, device): |
|
args = SLConfig.fromfile(model_config_path) |
|
args.device = device |
|
model = build_model(args) |
|
checkpoint = torch.load(model_checkpoint_path, map_location="cpu") |
|
load_res = model.load_state_dict( |
|
clean_state_dict(checkpoint["model"]), strict=False) |
|
print(load_res) |
|
_ = model.eval() |
|
return model |
|
|
|
|
|
def get_grounding_output(model, image, caption, box_threshold, text_threshold, with_logits=True): |
|
image = image.to("cuda") |
|
|
|
caption = caption.lower() |
|
caption = caption.strip() |
|
if not caption.endswith("."): |
|
caption = caption + "." |
|
|
|
start_time = time.time() |
|
model.to("cuda") |
|
with torch.no_grad(): |
|
outputs = model(image[None], captions=[caption]) |
|
print(f"Model forward time: {time.time() - start_time}") |
|
|
|
start_time = time.time() |
|
logits = outputs["pred_logits"].cpu().sigmoid()[0] |
|
boxes = outputs["pred_boxes"].cpu()[0] |
|
logits.shape[0] |
|
print(f"Process output time: {time.time() - start_time}") |
|
|
|
|
|
start_time = time.time() |
|
logits_filt = logits.clone() |
|
boxes_filt = boxes.clone() |
|
filt_mask = logits_filt.max(dim=1)[0] > box_threshold |
|
logits_filt = logits_filt[filt_mask] |
|
boxes_filt = boxes_filt[filt_mask] |
|
logits_filt.shape[0] |
|
print(f"Filter output time: {time.time() - start_time}") |
|
|
|
|
|
start_time = time.time() |
|
tokenlizer = model.tokenizer |
|
tokenized = tokenlizer(caption) |
|
print(f"Tokenize time: {time.time() - start_time}") |
|
|
|
pred_phrases = [] |
|
scores = [] |
|
start_time = time.time() |
|
for logit, box in zip(logits_filt, boxes_filt): |
|
pred_phrase = get_phrases_from_posmap( |
|
logit > text_threshold, tokenized, tokenlizer) |
|
if with_logits: |
|
pred_phrases.append( |
|
pred_phrase + f"({str(logit.max().item())[:4]})") |
|
else: |
|
pred_phrases.append(pred_phrase) |
|
scores.append(logit.max().item()) |
|
|
|
print(f"Build pred time: {time.time() - start_time}") |
|
|
|
return boxes_filt, torch.Tensor(scores), pred_phrases |
|
|
|
|
|
def draw_mask(mask, draw, random_color=False): |
|
if random_color: |
|
color = (random.randint(0, 255), random.randint( |
|
0, 255), random.randint(0, 255), 153) |
|
else: |
|
color = (30, 144, 255, 153) |
|
|
|
nonzero_coords = np.transpose(np.nonzero(mask)) |
|
|
|
for coord in nonzero_coords: |
|
draw.point(coord[::-1], fill=color) |
|
|
|
|
|
def draw_box(box, draw, label): |
|
|
|
color = tuple(np.random.randint(0, 255, size=3).tolist()) |
|
|
|
draw.rectangle(((box[0], box[1]), (box[2], box[3])), |
|
outline=color, width=2) |
|
|
|
if label: |
|
font = ImageFont.load_default() |
|
if hasattr(font, "getbbox"): |
|
bbox = draw.textbbox((box[0], box[1]), str(label), font) |
|
else: |
|
w, h = draw.textsize(str(label), font) |
|
bbox = (box[0], box[1], w + box[0], box[1] + h) |
|
draw.rectangle(bbox, fill=color) |
|
draw.text((box[0], box[1]), str(label), fill="white") |
|
|
|
draw.text((box[0], box[1]), label) |
|
|
|
|
|
config_file = f'{os.environ.get("path", ".")}/grounded_sam/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py' |
|
ckpt_repo_id = "ShilongLiu/GroundingDINO" |
|
ckpt_filenmae = f'{os.environ.get("path", ".")}/checkpoints/groundingdino_swint_ogc.pth' |
|
sam_checkpoint = f'{os.environ.get("path", ".")}/checkpoints/sam_hq_vit_h.pth' |
|
output_dir = "outputs" |
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
|
|
blip_processor = None |
|
blip_model = None |
|
groundingdino_model = load_model( |
|
config_file, ckpt_filenmae, device=device) |
|
sam_predictor = None |
|
inpaint_pipeline = None |
|
|
|
|
|
def run_grounded_sam(input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, iou_threshold, inpaint_mode): |
|
|
|
global blip_processor, blip_model, groundingdino_model, sam_predictor, inpaint_pipeline |
|
|
|
|
|
os.makedirs(output_dir, exist_ok=True) |
|
|
|
start_time = time.time() |
|
image_pil = input_image.convert("RGB") |
|
transformed_image = transform_image(image_pil) |
|
print(f"Transform image time: {time.time() - start_time}") |
|
|
|
start_time = time.time() |
|
print(f"Load model time: {time.time() - start_time}") |
|
|
|
|
|
start_time = time.time() |
|
boxes_filt, scores, pred_phrases = get_grounding_output( |
|
groundingdino_model, transformed_image, text_prompt, box_threshold, text_threshold |
|
) |
|
print(f"Run model time: {time.time() - start_time}") |
|
|
|
size = image_pil.size |
|
|
|
|
|
start_time = time.time() |
|
H, W = size[1], size[0] |
|
for i in range(boxes_filt.size(0)): |
|
boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H]) |
|
boxes_filt[i][:2] -= boxes_filt[i][2:] / 2 |
|
boxes_filt[i][2:] += boxes_filt[i][:2] |
|
|
|
boxes_filt = boxes_filt.cpu() |
|
print(f"Process boxes time: {time.time() - start_time}") |
|
|
|
|
|
start_time = time.time() |
|
print(f"Before NMS: {boxes_filt.shape[0]} boxes") |
|
nms_idx = torchvision.ops.nms( |
|
boxes_filt, scores, iou_threshold).numpy().tolist() |
|
boxes_filt = boxes_filt[nms_idx] |
|
pred_phrases = [pred_phrases[idx] for idx in nms_idx] |
|
print(f"After NMS: {boxes_filt.shape[0]} boxes") |
|
print(f"NMS time: {time.time() - start_time}") |
|
|
|
start_time = time.time() |
|
if sam_predictor is None: |
|
|
|
assert sam_checkpoint, 'sam_checkpoint is not found!' |
|
sam = build_sam(checkpoint=sam_checkpoint) |
|
sam.to(device=device) |
|
sam_predictor = SamPredictor(sam) |
|
print(f"Initialize SAM time: {time.time() - start_time}") |
|
|
|
image = np.array(image_pil) |
|
sam_predictor.set_image(image) |
|
|
|
start_time = time.time() |
|
transformed_boxes = sam_predictor.transform.apply_boxes_torch( |
|
boxes_filt, image.shape[:2]).to(device) |
|
print(f"Transform boxes time: {time.time() - start_time}") |
|
|
|
start_time = time.time() |
|
masks, _, _ = sam_predictor.predict_torch( |
|
point_coords=None, |
|
point_labels=None, |
|
boxes=transformed_boxes, |
|
multimask_output=False, |
|
) |
|
print(f"Predict time: {time.time() - start_time}") |
|
|
|
|
|
|
|
|
|
mask_image = Image.new('RGBA', size, color=(0, 0, 0, 0)) |
|
|
|
mask_draw = ImageDraw.Draw(mask_image) |
|
for mask in masks: |
|
draw_mask(mask[0].cpu().numpy(), mask_draw, random_color=True) |
|
|
|
image_draw = ImageDraw.Draw(image_pil) |
|
|
|
for box, label in zip(boxes_filt, pred_phrases): |
|
draw_box(box, image_draw, label) |
|
|
|
image_pil = image_pil.convert('RGBA') |
|
image_pil.alpha_composite(mask_image) |
|
|
|
return [image_pil, mask_image] |
|
|
|
def toBlackWhiteMask(input_pil): |
|
if input_pil.mode != 'RGBA': |
|
input_pil = input_pil.convert('RGBA') |
|
|
|
new_image = Image.new('RGBA', input_pil.size, (255, 255, 255, 255)) |
|
|
|
input_data = input_pil.getdata() |
|
new_data = [] |
|
|
|
for item in input_data: |
|
if item[3] == 0: |
|
new_data.append((0, 0, 0, 255)) |
|
else: |
|
new_data.append((255, 255, 255, 255)) |
|
|
|
new_image.putdata(new_data) |
|
|
|
return new_image |
|
|
|
def expand_white_pixels(input_pil, expand_by=1): |
|
|
|
grayscale = input_pil.convert('L') |
|
|
|
|
|
binary_mask = np.array(grayscale) > 245 |
|
|
|
|
|
dilated_mask = binary_dilation(binary_mask, iterations=expand_by) |
|
|
|
|
|
expanded_image = Image.fromarray(np.uint8(dilated_mask * 255)) |
|
|
|
return expanded_image |
|
|
|
def get_mask(input_pil, positive_prompt, expand_by=0): |
|
result = run_grounded_sam(input_image=input_pil, text_prompt=positive_prompt, task_type="seg", inpaint_prompt=None, box_threshold=0.3, |
|
text_threshold=0.25, iou_threshold=0.8, inpaint_mode="merge") |
|
|
|
result = result[1] |
|
|
|
black_white_result = toBlackWhiteMask(result) |
|
|
|
if expand_by > 0: |
|
black_white_result = expand_white_pixels(black_white_result, expand_by=expand_by) |
|
|
|
return black_white_result |
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser("Grounded SAM demo", add_help=True) |
|
parser.add_argument("--debug", action="store_true", |
|
help="using debug mode") |
|
parser.add_argument("--share", action="store_true", help="share the app") |
|
parser.add_argument('--no-gradio-queue', action="store_true", |
|
help='path to the SAM checkpoint') |
|
args = parser.parse_args() |
|
|
|
print(args) |
|
|
|
block = gr.Blocks() |
|
if not args.no_gradio_queue: |
|
block = block.queue() |
|
|
|
with block: |
|
with gr.Row(): |
|
with gr.Column(): |
|
input_image = gr.Image( |
|
source='upload', type="pil", value="demo1.jpg") |
|
task_type = gr.Dropdown( |
|
["det", "seg", "inpainting", "automatic"], value="seg", label="task_type") |
|
text_prompt = gr.Textbox(label="Text Prompt", placeholder="bear . beach .") |
|
inpaint_prompt = gr.Textbox(label="Inpaint Prompt", placeholder="A dinosaur, detailed, 4K.") |
|
run_button = gr.Button(label="Run") |
|
with gr.Accordion("Advanced options", open=False): |
|
box_threshold = gr.Slider( |
|
label="Box Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.001 |
|
) |
|
text_threshold = gr.Slider( |
|
label="Text Threshold", minimum=0.0, maximum=1.0, value=0.25, step=0.001 |
|
) |
|
iou_threshold = gr.Slider( |
|
label="IOU Threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.001 |
|
) |
|
inpaint_mode = gr.Dropdown( |
|
["merge", "first"], value="merge", label="inpaint_mode") |
|
|
|
with gr.Column(): |
|
gallery = gr.Gallery( |
|
label="Generated images", show_label=False, elem_id="gallery" |
|
).style(preview=True, grid=2, object_fit="scale-down") |
|
|
|
run_button.click(fn=run_grounded_sam, inputs=[ |
|
input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, iou_threshold, inpaint_mode], outputs=gallery) |
|
|
|
block.launch(debug=args.debug, share=args.share, show_error=True) |
|
|