support GroundingDINO and segment-anything

app.py

@@ -87,6 +87,11 @@ VISUAL_CHATGPT_SUFFIX_CN = """你对文件名的正确性非常严格，而且
 Thought: Do I need to use a tool? {agent_scratchpad}
 """
 
+import os
+
+os.system('pip install  git+https://github.com/IDEA-Research/GroundingDINO.git')
+os.system('pip install  git+https://github.com/facebookresearch/segment-anything.git')
+
 from visual_foundation_models import *
 from langchain.agents.initialize import initialize_agent
 from langchain.agents.tools import Tool

requirements.txt

@@ -1,32 +1,32 @@
 --extra-index-url https://download.pytorch.org/whl/cu113
-torch==1.12.1
-torchvision==0.13.1
-numpy==1.23.1
-transformers==4.26.1
-albumentations==1.3.0
-opencv-contrib-python==4.3.0.36
-imageio==2.9.0
-imageio-ffmpeg==0.4.2
-pytorch-lightning==1.5.0
-omegaconf==2.1.1
-test-tube>=0.7.5
-streamlit==1.12.1
-einops==0.3.0
-webdataset==0.2.5
-kornia==0.6
-open_clip_torch==2.0.2
-invisible-watermark>=0.1.5
-streamlit-drawable-canvas==0.8.0
-torchmetrics==0.6.0
-timm==0.6.12
-addict==2.4.0
-yapf==0.32.0
-prettytable==3.6.0
-safetensors==0.2.7
-basicsr==1.4.2
 langchain==0.0.101
-diffusers==0.14.0
+torch==1.13.1
+torchvision==0.14.1
+wget==3.2
+accelerate
+addict
+albumentations
+basicsr
+controlnet-aux
+diffusers
+einops
 gradio
+imageio
+imageio-ffmpeg
+invisible-watermark
+kornia
+numpy
+omegaconf
+open_clip_torch
 openai
-accelerate
-controlnet-aux==0.0.1
+opencv-python
+prettytable
+safetensors
+streamlit
+test-tube
+timm
+torchmetrics
+transformers
+webdataset
+yapf
+

visual_foundation_models.py

@@ -19,6 +19,18 @@ import math
 
 from langchain.llms.openai import OpenAI
 
+# Grounding DINO
+import groundingdino.datasets.transforms as T
+from groundingdino.models import build_model
+from groundingdino.util import box_ops
+from groundingdino.util.slconfig import SLConfig
+from groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap
+
+# segment anything
+from segment_anything import build_sam, SamPredictor, SamAutomaticMaskGenerator
+import matplotlib.pyplot as plt
+import wget
+
 def prompts(name, description):
     def decorator(func):
         func.name = name
@@ -101,76 +113,6 @@ def get_new_image_name(org_img_name, func_name="update"):
     return os.path.join(head, new_file_name)
 
 
-class MaskFormer:
-    def __init__(self, device):
-        print(f"Initializing MaskFormer to {device}")
-        self.device = device
-        self.processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
-        self.model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined").to(device)
-
-    def inference(self, image_path, text):
-        threshold = 0.5
-        min_area = 0.02
-        padding = 20
-        original_image = Image.open(image_path)
-        image = original_image.resize((512, 512))
-        inputs = self.processor(text=text, images=image, padding="max_length", return_tensors="pt").to(self.device)
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-        mask = torch.sigmoid(outputs[0]).squeeze().cpu().numpy() > threshold
-        area_ratio = len(np.argwhere(mask)) / (mask.shape[0] * mask.shape[1])
-        if area_ratio < min_area:
-            return None
-        true_indices = np.argwhere(mask)
-        mask_array = np.zeros_like(mask, dtype=bool)
-        for idx in true_indices:
-            padded_slice = tuple(slice(max(0, i - padding), i + padding + 1) for i in idx)
-            mask_array[padded_slice] = True
-        visual_mask = (mask_array * 255).astype(np.uint8)
-        image_mask = Image.fromarray(visual_mask)
-        return image_mask.resize(original_image.size)
-
-
-class ImageEditing:
-    def __init__(self, device):
-        print(f"Initializing ImageEditing to {device}")
-        self.device = device
-        self.mask_former = MaskFormer(device=self.device)
-        self.revision = 'fp16' if 'cuda' in device else None
-        self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32
-        self.inpaint = StableDiffusionInpaintPipeline.from_pretrained(
-            "runwayml/stable-diffusion-inpainting", revision=self.revision, torch_dtype=self.torch_dtype).to(device)
-
-    @prompts(name="Remove Something From The Photo",
-             description="useful when you want to remove and object or something from the photo "
-                         "from its description or location. "
-                         "The input to this tool should be a comma separated string of two, "
-                         "representing the image_path and the object need to be removed. ")
-    def inference_remove(self, inputs):
-        image_path, to_be_removed_txt = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
-        return self.inference_replace(f"{image_path},{to_be_removed_txt},background")
-
-    @prompts(name="Replace Something From The Photo",
-             description="useful when you want to replace an object from the object description or "
-                         "location with another object from its description. "
-                         "The input to this tool should be a comma separated string of three, "
-                         "representing the image_path, the object to be replaced, the object to be replaced with ")
-    def inference_replace(self, inputs):
-        image_path, to_be_replaced_txt, replace_with_txt = inputs.split(",")
-        original_image = Image.open(image_path)
-        original_size = original_image.size
-        mask_image = self.mask_former.inference(image_path, to_be_replaced_txt)
-        updated_image = self.inpaint(prompt=replace_with_txt, image=original_image.resize((512, 512)),
-                                     mask_image=mask_image.resize((512, 512))).images[0]
-        updated_image_path = get_new_image_name(image_path, func_name="replace-something")
-        updated_image = updated_image.resize(original_size)
-        updated_image.save(updated_image_path)
-        print(
-            f"\nProcessed ImageEditing, Input Image: {image_path}, Replace {to_be_replaced_txt} to {replace_with_txt}, "
-            f"Output Image: {updated_image_path}")
-        return updated_image_path
-
-
 class InstructPix2Pix:
     def __init__(self, device):
         print(f"Initializing InstructPix2Pix to {device}")
@@ -283,7 +225,7 @@ class CannyText2Image:
         self.seed = -1
         self.a_prompt = 'best quality, extremely detailed'
         self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, ' \
-                            'fewer digits, cropped, worst quality, low quality'
+                        'fewer digits, cropped, worst quality, low quality'
 
     @prompts(name="Generate Image Condition On Canny Image",
              description="useful when you want to generate a new real image from both the user description and a canny image."
@@ -340,7 +282,7 @@ class LineText2Image:
         self.seed = -1
         self.a_prompt = 'best quality, extremely detailed'
         self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, ' \
-                            'fewer digits, cropped, worst quality, low quality'
+                        'fewer digits, cropped, worst quality, low quality'
 
     @prompts(name="Generate Image Condition On Line Image",
              description="useful when you want to generate a new real image from both the user description "
@@ -398,7 +340,7 @@ class HedText2Image:
         self.seed = -1
         self.a_prompt = 'best quality, extremely detailed'
         self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, ' \
-                            'fewer digits, cropped, worst quality, low quality'
+                        'fewer digits, cropped, worst quality, low quality'
 
     @prompts(name="Generate Image Condition On Soft Hed Boundary Image",
              description="useful when you want to generate a new real image from both the user description "
@@ -456,7 +398,7 @@ class ScribbleText2Image:
         self.seed = -1
         self.a_prompt = 'best quality, extremely detailed'
         self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, ' \
-                            'fewer digits, cropped, worst quality, low quality'
+                        'fewer digits, cropped, worst quality, low quality'
 
     @prompts(name="Generate Image Condition On Sketch Image",
              description="useful when you want to generate a new real image from both the user description and "
@@ -512,7 +454,7 @@ class PoseText2Image:
         self.unconditional_guidance_scale = 9.0
         self.a_prompt = 'best quality, extremely detailed'
         self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit,' \
-                            ' fewer digits, cropped, worst quality, low quality'
+                        ' fewer digits, cropped, worst quality, low quality'
 
     @prompts(name="Generate Image Condition On Pose Image",
              description="useful when you want to generate a new real image from both the user description "
@@ -536,73 +478,6 @@ class PoseText2Image:
         return updated_image_path
 
 
-class Image2Seg:
-    def __init__(self, device):
-        print("Initializing Image2Seg")
-        self.image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
-        self.image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
-        self.ade_palette = [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
-                            [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
-                            [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
-                            [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
-                            [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
-                            [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
-                            [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
-                            [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
-                            [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
-                            [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
-                            [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
-                            [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
-                            [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
-                            [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
-                            [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
-                            [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
-                            [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
-                            [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
-                            [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
-                            [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
-                            [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
-                            [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
-                            [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
-                            [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
-                            [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
-                            [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
-                            [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
-                            [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
-                            [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
-                            [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
-                            [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
-                            [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
-                            [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
-                            [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
-                            [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
-                            [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
-                            [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
-                            [102, 255, 0], [92, 0, 255]]
-
-    @prompts(name="Segmentation On Image",
-             description="useful when you want to detect segmentations of the image. "
-                         "like: segment this image, or generate segmentations on this image, "
-                         "or perform segmentation on this image. "
-                         "The input to this tool should be a string, representing the image_path")
-    def inference(self, inputs):
-        image = Image.open(inputs)
-        pixel_values = self.image_processor(image, return_tensors="pt").pixel_values
-        with torch.no_grad():
-            outputs = self.image_segmentor(pixel_values)
-        seg = self.image_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
-        color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)  # height, width, 3
-        palette = np.array(self.ade_palette)
-        for label, color in enumerate(palette):
-            color_seg[seg == label, :] = color
-        color_seg = color_seg.astype(np.uint8)
-        segmentation = Image.fromarray(color_seg)
-        updated_image_path = get_new_image_name(inputs, func_name="segmentation")
-        segmentation.save(updated_image_path)
-        print(f"\nProcessed Image2Seg, Input Image: {inputs}, Output Pose: {updated_image_path}")
-        return updated_image_path
-
-
 class SegText2Image:
     def __init__(self, device):
         print(f"Initializing SegText2Image to {device}")
@@ -617,7 +492,7 @@ class SegText2Image:
         self.seed = -1
         self.a_prompt = 'best quality, extremely detailed'
         self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit,' \
-                            ' fewer digits, cropped, worst quality, low quality'
+                        ' fewer digits, cropped, worst quality, low quality'
 
     @prompts(name="Generate Image Condition On Segmentations",
              description="useful when you want to generate a new real image from both the user description and segmentations. "
@@ -676,7 +551,7 @@ class DepthText2Image:
         self.seed = -1
         self.a_prompt = 'best quality, extremely detailed'
         self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit,' \
-                            ' fewer digits, cropped, worst quality, low quality'
+                        ' fewer digits, cropped, worst quality, low quality'
 
     @prompts(name="Generate Image Condition On Depth",
              description="useful when you want to generate a new real image from both the user description and depth image. "
@@ -747,7 +622,7 @@ class NormalText2Image:
         self.seed = -1
         self.a_prompt = 'best quality, extremely detailed'
         self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit,' \
-                            ' fewer digits, cropped, worst quality, low quality'
+                        ' fewer digits, cropped, worst quality, low quality'
 
     @prompts(name="Generate Image Condition On Normal Map",
              description="useful when you want to generate a new real image from both the user description and normal map. "
@@ -793,10 +668,284 @@ class VisualQuestionAnswering:
               f"Output Answer: {answer}")
         return answer
 
+
+class Segmenting:
+    def __init__(self, device):
+        print(f"Inintializing Segmentation to {device}")
+        self.device = device
+        self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32
+        self.model_checkpoint_path = os.path.join("checkpoints", "sam")
+
+        self.download_parameters()
+        self.sam = build_sam(checkpoint=self.model_checkpoint_path).to(device)
+        self.sam_predictor = SamPredictor(self.sam)
+        self.mask_generator = SamAutomaticMaskGenerator(self.sam)
+
+    def download_parameters(self):
+        url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
+        if not os.path.exists(self.model_checkpoint_path):
+            wget.download(url, out=self.model_checkpoint_path)
+
+    def show_mask(self, mask, ax, random_color=False):
+        if random_color:
+            color = np.concatenate([np.random.random(3), np.array([1])], axis=0)
+        else:
+            color = np.array([30 / 255, 144 / 255, 255 / 255, 1])
+        h, w = mask.shape[-2:]
+        mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+        ax.imshow(mask_image)
+
+    def show_box(self, box, ax, label):
+        x0, y0 = box[0], box[1]
+        w, h = box[2] - box[0], box[3] - box[1]
+        ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))
+        ax.text(x0, y0, label)
+
+    def get_mask_with_boxes(self, image_pil, image, boxes_filt):
+
+        size = image_pil.size
+        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()
+        transformed_boxes = self.sam_predictor.transform.apply_boxes_torch(boxes_filt, image.shape[:2]).to(self.device)
+
+        masks, _, _ = self.sam_predictor.predict_torch(
+            point_coords=None,
+            point_labels=None,
+            boxes=transformed_boxes.to(self.device),
+            multimask_output=False,
+        )
+        return masks
+
+    def segment_image_with_boxes(self, image_pil, image_path, boxes_filt, pred_phrases):
+
+        image = cv2.imread(image_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        self.sam_predictor.set_image(image)
+
+        masks = self.get_mask_with_boxes(image_pil, image, boxes_filt)
+
+        # draw output image
+        plt.figure(figsize=(10, 10))
+        plt.imshow(image)
+        for mask in masks:
+            self.show_mask(mask.cpu().numpy(), plt.gca(), random_color=True)
+
+        updated_image_path = get_new_image_name(image_path, func_name="segmentation")
+        plt.axis('off')
+        plt.savefig(
+            updated_image_path,
+            bbox_inches="tight", dpi=300, pad_inches=0.0
+        )
+        return updated_image_path
+
+    @prompts(name="Segment the Image",
+             description="useful when you want to segment all the part of the image, but not segment a certain object."
+                         "like: segment all the object in this image, or generate segmentations on this image, "
+                         "or segment the image,"
+                         "or perform segmentation on this image, "
+                         "or segment all the object in this image."
+                         "The input to this tool should be a string, representing the image_path")
+    def inference_all(self, image_path):
+        image = cv2.imread(image_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        masks = self.mask_generator.generate(image)
+        plt.figure(figsize=(20, 20))
+        plt.imshow(image)
+        if len(masks) == 0:
+            return
+        sorted_anns = sorted(masks, key=(lambda x: x['area']), reverse=True)
+        ax = plt.gca()
+        ax.set_autoscale_on(False)
+        polygons = []
+        color = []
+        for ann in sorted_anns:
+            m = ann['segmentation']
+            img = np.ones((m.shape[0], m.shape[1], 3))
+            color_mask = np.random.random((1, 3)).tolist()[0]
+            for i in range(3):
+                img[:, :, i] = color_mask[i]
+            ax.imshow(np.dstack((img, m)))
+
+        updated_image_path = get_new_image_name(image_path, func_name="segment-image")
+        plt.axis('off')
+        plt.savefig(
+            updated_image_path,
+            bbox_inches="tight", dpi=300, pad_inches=0.0
+        )
+        return updated_image_path
+
+
+class Text2Box:
+    def __init__(self, device):
+        print(f"Initializing ObjectDetection to {device}")
+        self.device = device
+        self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32
+        self.model_checkpoint_path = os.path.join("checkpoints", "groundingdino")
+        self.model_config_path = os.path.join("checkpoints", "grounding_config.py")
+        self.download_parameters()
+        self.box_threshold = 0.3
+        self.text_threshold = 0.25
+        self.grounding = (self.load_model()).to(self.device)
+
+    def download_parameters(self):
+        url = "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth"
+        if not os.path.exists(self.model_checkpoint_path):
+            wget.download(url, out=self.model_checkpoint_path)
+        config_url = "https://raw.githubusercontent.com/IDEA-Research/GroundingDINO/main/groundingdino/config/GroundingDINO_SwinT_OGC.py"
+        if not os.path.exists(self.model_config_path):
+            wget.download(config_url, out=self.model_config_path)
+
+    def load_image(self, image_path):
+        # load image
+        image_pil = Image.open(image_path).convert("RGB")  # load image
+
+        transform = T.Compose(
+            [
+                T.RandomResize([512], max_size=1333),
+                T.ToTensor(),
+                T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
+        image, _ = transform(image_pil, None)  # 3, h, w
+        return image_pil, image
+
+    def load_model(self):
+        args = SLConfig.fromfile(self.model_config_path)
+        args.device = self.device
+        model = build_model(args)
+        checkpoint = torch.load(self.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_boxes(self, image, caption, with_logits=True):
+        caption = caption.lower()
+        caption = caption.strip()
+        if not caption.endswith("."):
+            caption = caption + "."
+        image = image.to(self.device)
+        with torch.no_grad():
+            outputs = self.grounding(image[None], captions=[caption])
+        logits = outputs["pred_logits"].cpu().sigmoid()[0]  # (nq, 256)
+        boxes = outputs["pred_boxes"].cpu()[0]  # (nq, 4)
+        logits.shape[0]
+
+        # filter output
+        logits_filt = logits.clone()
+        boxes_filt = boxes.clone()
+        filt_mask = logits_filt.max(dim=1)[0] > self.box_threshold
+        logits_filt = logits_filt[filt_mask]  # num_filt, 256
+        boxes_filt = boxes_filt[filt_mask]  # num_filt, 4
+        logits_filt.shape[0]
+
+        # get phrase
+        tokenlizer = self.grounding.tokenizer
+        tokenized = tokenlizer(caption)
+        # build pred
+        pred_phrases = []
+        for logit, box in zip(logits_filt, boxes_filt):
+            pred_phrase = get_phrases_from_posmap(logit > self.text_threshold, tokenized, tokenlizer)
+            if with_logits:
+                pred_phrases.append(pred_phrase + f"({str(logit.max().item())[:4]})")
+            else:
+                pred_phrases.append(pred_phrase)
+
+        return boxes_filt, pred_phrases
+
+    def plot_boxes_to_image(self, image_pil, tgt):
+        H, W = tgt["size"]
+        boxes = tgt["boxes"]
+        labels = tgt["labels"]
+        assert len(boxes) == len(labels), "boxes and labels must have same length"
+
+        draw = ImageDraw.Draw(image_pil)
+        mask = Image.new("L", image_pil.size, 0)
+        mask_draw = ImageDraw.Draw(mask)
+
+        # draw boxes and masks
+        for box, label in zip(boxes, labels):
+            # from 0..1 to 0..W, 0..H
+            box = box * torch.Tensor([W, H, W, H])
+            # from xywh to xyxy
+            box[:2] -= box[2:] / 2
+            box[2:] += box[:2]
+            # random color
+            color = tuple(np.random.randint(0, 255, size=3).tolist())
+            # draw
+            x0, y0, x1, y1 = box
+            x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1)
+
+            draw.rectangle([x0, y0, x1, y1], outline=color, width=6)
+            # draw.text((x0, y0), str(label), fill=color)
+
+            font = ImageFont.load_default()
+            if hasattr(font, "getbbox"):
+                bbox = draw.textbbox((x0, y0), str(label), font)
+            else:
+                w, h = draw.textsize(str(label), font)
+                bbox = (x0, y0, w + x0, y0 + h)
+            # bbox = draw.textbbox((x0, y0), str(label))
+            draw.rectangle(bbox, fill=color)
+            draw.text((x0, y0), str(label), fill="white")
+
+            mask_draw.rectangle([x0, y0, x1, y1], fill=255, width=2)
+
+        return image_pil, mask
+
+    @prompts(name="Detect the Give Object",
+             description="useful when you only want to detect or find out given objects in the picture"
+                         "The input to this tool should be a comma separated string of two, "
+                         "representing the image_path, the text description of the object to be found")
+    def inference(self, inputs):
+        image_path, det_prompt = inputs.split(",")
+        print(f"image_path={image_path}, text_prompt={det_prompt}")
+        image_pil, image = self.load_image(image_path)
+
+        boxes_filt, pred_phrases = self.get_grounding_boxes(image, det_prompt)
+
+        size = image_pil.size
+        pred_dict = {
+            "boxes": boxes_filt,
+            "size": [size[1], size[0]],  # H,W
+            "labels": pred_phrases, }
+
+        image_with_box = self.plot_boxes_to_image(image_pil, pred_dict)[0]
+
+        updated_image_path = get_new_image_name(image_path, func_name="detect-something")
+        updated_image = image_with_box.resize(size)
+        updated_image.save(updated_image_path)
+        print(
+            f"\nProcessed ObejectDetecting, Input Image: {image_path}, Object to be Detect {det_prompt}, "
+            f"Output Image: {updated_image_path}")
+        return updated_image_path
+
+
+class Inpainting:
+    def __init__(self, device):
+        self.device = device
+        self.revision = 'fp16' if 'cuda' in self.device else None
+        self.torch_dtype = torch.float16 if 'cuda' in self.device else torch.float32
+
+        self.inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+            "runwayml/stable-diffusion-inpainting", revision=self.revision, torch_dtype=self.torch_dtype).to(device)
+
+    def __call__(self, prompt, original_image, mask_image):
+        update_image = self.inpaint(prompt=prompt, image=original_image.resize((512, 512)),
+                                    mask_image=mask_image.resize((512, 512))).images[0]
+        return update_image
+
+
 class InfinityOutPainting:
-    template_model = True # Add this line to show this is a template model.
+    template_model = True  # Add this line to show this is a template model.
+
     def __init__(self, ImageCaptioning, ImageEditing, VisualQuestionAnswering):
-        # self.llm = OpenAI(temperature=0)
+        self.llm = OpenAI(temperature=0)
         self.ImageCaption = ImageCaptioning
         self.ImageEditing = ImageEditing
         self.ImageVQA = VisualQuestionAnswering
@@ -814,16 +963,16 @@ class InfinityOutPainting:
 
     def get_BLIP_caption(self, image):
         inputs = self.ImageCaption.processor(image, return_tensors="pt").to(self.ImageCaption.device,
-                                                                                self.ImageCaption.torch_dtype)
+                                                                            self.ImageCaption.torch_dtype)
         out = self.ImageCaption.model.generate(**inputs)
         BLIP_caption = self.ImageCaption.processor.decode(out[0], skip_special_tokens=True)
         return BLIP_caption
 
-    # def check_prompt(self, prompt):
-    #     check = f"Here is a paragraph with adjectives. " \
-    #             f"{prompt} " \
-    #             f"Please change all plural forms in the adjectives to singular forms. "
-    #     return self.llm(check)
+    def check_prompt(self, prompt):
+        check = f"Here is a paragraph with adjectives. " \
+                f"{prompt} " \
+                f"Please change all plural forms in the adjectives to singular forms. "
+        return self.llm(check)
 
     def get_imagine_caption(self, image, imagine):
         BLIP_caption = self.get_BLIP_caption(image)
@@ -836,14 +985,13 @@ class InfinityOutPainting:
                          f"You should make the painting as vivid and realistic as possible" \
                          f"You can not use words like painting or picture" \
                          f"and you should use no more than 50 words to describe it"
-        # caption = self.llm(imagine_prompt) if imagine else BLIP_caption
-        caption = BLIP_caption
-        # caption = self.check_prompt(caption)
+        caption = self.llm(imagine_prompt) if imagine else BLIP_caption
+        caption = self.check_prompt(caption)
         print(f'BLIP observation: {BLIP_caption}, ChatGPT imagine to {caption}') if imagine else print(
             f'Prompt: {caption}')
         return caption
 
-    def resize_image(self, image, max_size=100000, multiple=8):
+    def resize_image(self, image, max_size=1000000, multiple=8):
         aspect_ratio = image.size[0] / image.size[1]
         new_width = int(math.sqrt(max_size * aspect_ratio))
         new_height = int(new_width / aspect_ratio)
@@ -889,4 +1037,94 @@ class InfinityOutPainting:
         out_painted_image.save(updated_image_path)
         print(f"\nProcessed InfinityOutPainting, Input Image: {image_path}, Input Resolution: {resolution}, "
               f"Output Image: {updated_image_path}")
+        return updated_image_path
+
+
+class ObjectSegmenting:
+    template_model = True  # Add this line to show this is a template model.
+
+    def __init__(self, Text2Box: Text2Box, Segmenting: Segmenting):
+        # self.llm = OpenAI(temperature=0)
+        self.grounding = Text2Box
+        self.sam = Segmenting
+
+    @prompts(name="Segment the given object",
+             description="useful when you only want to segment the certain objects in the picture"
+                         "according to the given text"
+                         "like: segment the cat,"
+                         "or can you segment an obeject for me"
+                         "The input to this tool should be a comma separated string of two, "
+                         "representing the image_path, the text description of the object to be found")
+    def inference(self, inputs):
+        image_path, det_prompt = inputs.split(",")
+        print(f"image_path={image_path}, text_prompt={det_prompt}")
+        image_pil, image = self.grounding.load_image(image_path)
+        boxes_filt, pred_phrases = self.grounding.get_grounding_boxes(image, det_prompt)
+        updated_image_path = self.sam.segment_image_with_boxes(image_pil, image_path, boxes_filt, pred_phrases)
+        print(
+            f"\nProcessed ObejectSegmenting, Input Image: {image_path}, Object to be Segment {det_prompt}, "
+            f"Output Image: {updated_image_path}")
+        return updated_image_path
+
+
+class ImageEditing:
+    template_model = True
+
+    def __init__(self, Text2Box: Text2Box, Segmenting: Segmenting, Inpainting: Inpainting):
+        print(f"Initializing ImageEditing")
+        self.sam = Segmenting
+        self.grounding = Text2Box
+        self.inpaint = Inpainting
+
+    def pad_edge(self, mask, padding):
+        # mask Tensor [H,W]
+        mask = mask.numpy()
+        true_indices = np.argwhere(mask)
+        mask_array = np.zeros_like(mask, dtype=bool)
+        for idx in true_indices:
+            padded_slice = tuple(slice(max(0, i - padding), i + padding + 1) for i in idx)
+            mask_array[padded_slice] = True
+        new_mask = (mask_array * 255).astype(np.uint8)
+        # new_mask
+        return new_mask
+
+    @prompts(name="Remove Something From The Photo",
+             description="useful when you want to remove and object or something from the photo "
+                         "from its description or location. "
+                         "The input to this tool should be a comma separated string of two, "
+                         "representing the image_path and the object need to be removed. ")
+    def inference_remove(self, inputs):
+        image_path, to_be_removed_txt = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
+        return self.inference_replace_sam(f"{image_path},{to_be_removed_txt},background")
+
+    @prompts(name="Replace Something From The Photo",
+             description="useful when you want to replace an object from the object description or "
+                         "location with another object from its description. "
+                         "The input to this tool should be a comma separated string of three, "
+                         "representing the image_path, the object to be replaced, the object to be replaced with ")
+    def inference_replace_sam(self, inputs):
+        image_path, to_be_replaced_txt, replace_with_txt = inputs.split(",")
+
+        print(f"image_path={image_path}, to_be_replaced_txt={to_be_replaced_txt}")
+        image_pil, image = self.grounding.load_image(image_path)
+        boxes_filt, pred_phrases = self.grounding.get_grounding_boxes(image, to_be_replaced_txt)
+        image = cv2.imread(image_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        self.sam.sam_predictor.set_image(image)
+        masks = self.sam.get_mask_with_boxes(image_pil, image, boxes_filt)
+        mask = torch.sum(masks, dim=0).unsqueeze(0)
+        mask = torch.where(mask > 0, True, False)
+        mask = mask.squeeze(0).squeeze(0).cpu()  # tensor
+
+        mask = self.pad_edge(mask, padding=20)  # numpy
+        mask_image = Image.fromarray(mask)
+
+        updated_image = self.inpaint(prompt=replace_with_txt, original_image=image_pil,
+                                     mask_image=mask_image)
+        updated_image_path = get_new_image_name(image_path, func_name="replace-something")
+        updated_image = updated_image.resize(image_pil.size)
+        updated_image.save(updated_image_path)
+        print(
+            f"\nProcessed ImageEditing, Input Image: {image_path}, Replace {to_be_replaced_txt} to {replace_with_txt}, "
+            f"Output Image: {updated_image_path}")
         return updated_image_path