indigoman/text_image_editing

1# Gradio and other necessary imports2import gradio as gr3import subprocess4 5subprocess.run(["bash", "setup.sh"])6from segment_anything import SamPredictor, sam_model_registry7from diffusers import StableDiffusionInpaintPipeline8from GroundingDINO.groundingdino.util.inference import load_model, load_image, predict, annotate9from GroundingDINO.groundingdino.util import box_ops10from PIL import Image11import torch12import numpy as np13 14import os15device = torch.device("cpu")16# ----SAM17 18print("path", os.getcwd())19 20model_type = "vit_h"21predictor = SamPredictor(sam_model_registry[model_type](checkpoint="./GroundingDINO/weights/sam_vit_h_4b8939.pth").to(device))22# ------Stable Diffusion23pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting", torch_dtype=torch.float32).to(device)24# ----Grounding DINO25groundingdino_model = load_model("./GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py", "./GroundingDINO/weights/groundingdino_swint_ogc.pth")26 27BOX_TRESHOLD = 0.328TEXT_TRESHOLD = 0.2529 30def show_mask(mask, image, random_color=True):31    if random_color:32        color = np.concatenate([np.random.random(3), np.array([0.8])], axis=0)33    else:34        color = np.array([30/255, 144/255, 255/255, 0.6])35    h, w = mask.shape[-2:]36    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)37 38    annotated_frame_pil = Image.fromarray(image).convert("RGBA")39    mask_image_pil = Image.fromarray((mask_image.cpu().numpy() * 255).astype(np.uint8)).convert("RGBA")40 41    return np.array(Image.alpha_composite(annotated_frame_pil, mask_image_pil))42 43def process_boxes(boxes, src):44    H, W, _ = src.shape45    boxes_xyxy = box_ops.box_cxcywh_to_xyxy(boxes) * torch.Tensor([W, H, W, H])46    return predictor.transform.apply_boxes_torch(boxes_xyxy, src.shape[:2]).to(device)47 48def edit_image(path: str, item: str, prompt=str, box_threshold=BOX_TRESHOLD, text_threshold=TEXT_TRESHOLD):49    src, img = load_image(path)50    boxes, logits, phrases = predict(51        model=groundingdino_model,52        image=img,53        caption=item,54        box_threshold=box_threshold,55        text_threshold=text_threshold56    )57    predictor.set_image(src)58    new_boxes = process_boxes(boxes, src)59    masks, _, _ = predictor.predict_torch(60        point_coords=None,61        point_labels=None,62        boxes=new_boxes,63        multimask_output=False,64    )65    img_annotated_mask = show_mask(masks[0][0].cpu(),66        annotate(image_source=src, boxes=boxes, logits=logits, phrases=phrases)[...,::-1]67    )68    return pipe(prompt=prompt,69        image=Image.fromarray(src).resize((512, 512)),70        mask_image=Image.fromarray(masks[0][0].cpu().numpy()).resize((512, 512))71    ).images[0]72 73# Define the Gradio interface74iface = gr.Interface(75    fn=edit_image, 76    inputs=[77        gr.inputs.Textbox(label="Image Path"),78        gr.inputs.Textbox(label="Caption"),79    ], 80    outputs=gr.outputs.Image(type="numpy"),81)82 83iface = gr.Interface(84    fn=edit_image, 85    inputs=[86        gr.inputs.Image(type="filepath", label="Upload Image"),87        gr.inputs.Textbox(label="Item"),88        gr.inputs.Textbox(label="Prompt"),89        gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.3, label="Box Threshold"),90        gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.2, label="Text Threshold")91    ], 92    outputs=gr.outputs.Image(type="numpy"),93)94 95iface.launch(inbrowser=True)96 97 98# path = './fire3.jpg'99# edit_image(path, "fire hydrant", "phone booth", 0.5, 0.2)