ai-mirow-projects
/
swarms
mirror of https://github.com/kyegomez/swarms
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

no dino

Browse Source
main
Kye 2 years ago
parent 6f0fb1e3e2
commit 4b23afb4e7
1 changed files with 146 additions and 146 deletions
Show Stats Download Patch File Download Diff File

296
swarms/agents/workers/multi_modal.py
Unescape View File

@ -32,11 +32,11 @@ from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain.llms.openai import OpenAI from langchain.llms.openai import OpenAI
# Grounding DINO # Grounding DINO
import groundingdino.datasets.transforms as T # import groundingdino.datasets.transforms as T
from groundingdino.models import build_model # from groundingdino.models import build_model
from groundingdino.util import box_ops # from groundingdino.util import box_ops
from groundingdino.util.slconfig import SLConfig # from groundingdino.util.slconfig import SLConfig
from groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap # from groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap
# segment anything # segment anything
from segment_anything import build_sam, SamPredictor, SamAutomaticMaskGenerator from segment_anything import build_sam, SamPredictor, SamAutomaticMaskGenerator
@ -1023,149 +1023,149 @@ class Segmenting:
) )
return updated_image_path return updated_image_path
class Text2Box: # class Text2Box:
def __init__(self, device): # def __init__(self, device):
print(f"Initializing ObjectDetection to {device}") # print(f"Initializing ObjectDetection to {device}")
self.device = device # self.device = device
self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 # self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32
self.model_checkpoint_path = os.path.join("checkpoints","groundingdino") # self.model_checkpoint_path = os.path.join("checkpoints","groundingdino")
self.model_config_path = os.path.join("checkpoints","grounding_config.py") # self.model_config_path = os.path.join("checkpoints","grounding_config.py")
self.download_parameters() # self.download_parameters()
self.box_threshold = 0.3 # self.box_threshold = 0.3
self.text_threshold = 0.25 # self.text_threshold = 0.25
self.grounding = (self.load_model()).to(self.device) # self.grounding = (self.load_model()).to(self.device)
def download_parameters(self): # def download_parameters(self):
url = "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth" # 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): # if not os.path.exists(self.model_checkpoint_path):
wget.download(url,out=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" # 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): # if not os.path.exists(self.model_config_path):
wget.download(config_url,out=self.model_config_path) # wget.download(config_url,out=self.model_config_path)
def load_image(self,image_path): # def load_image(self,image_path):
# load image # # load image
image_pil = Image.open(image_path).convert("RGB") # load image # image_pil = Image.open(image_path).convert("RGB") # load image
transform = T.Compose( # transform = T.Compose(
[ # [
T.RandomResize([512], max_size=1333), # T.RandomResize([512], max_size=1333),
T.ToTensor(), # T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), # T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
] # ]
) # )
image, _ = transform(image_pil, None) # 3, h, w # image, _ = transform(image_pil, None) # 3, h, w
return image_pil, image # return image_pil, image
def load_model(self): # def load_model(self):
args = SLConfig.fromfile(self.model_config_path) # args = SLConfig.fromfile(self.model_config_path)
args.device = self.device # args.device = self.device
model = build_model(args) # model = build_model(args)
checkpoint = torch.load(self.model_checkpoint_path, map_location="cpu") # checkpoint = torch.load(self.model_checkpoint_path, map_location="cpu")
load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) # load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False)
print(load_res) # print(load_res)
_ = model.eval() # _ = model.eval()
return model # return model
def get_grounding_boxes(self, image, caption, with_logits=True): # def get_grounding_boxes(self, image, caption, with_logits=True):
caption = caption.lower() # caption = caption.lower()
caption = caption.strip() # caption = caption.strip()
if not caption.endswith("."): # if not caption.endswith("."):
caption = caption + "." # caption = caption + "."
image = image.to(self.device) # image = image.to(self.device)
with torch.no_grad(): # with torch.no_grad():
outputs = self.grounding(image[None], captions=[caption]) # outputs = self.grounding(image[None], captions=[caption])
logits = outputs["pred_logits"].cpu().sigmoid()[0] # (nq, 256) # logits = outputs["pred_logits"].cpu().sigmoid()[0] # (nq, 256)
boxes = outputs["pred_boxes"].cpu()[0] # (nq, 4) # boxes = outputs["pred_boxes"].cpu()[0] # (nq, 4)
logits.shape[0] # logits.shape[0]
# filter output # # filter output
logits_filt = logits.clone() # logits_filt = logits.clone()
boxes_filt = boxes.clone() # boxes_filt = boxes.clone()
filt_mask = logits_filt.max(dim=1)[0] > self.box_threshold # filt_mask = logits_filt.max(dim=1)[0] > self.box_threshold
logits_filt = logits_filt[filt_mask] # num_filt, 256 # logits_filt = logits_filt[filt_mask] # num_filt, 256
boxes_filt = boxes_filt[filt_mask] # num_filt, 4 # boxes_filt = boxes_filt[filt_mask] # num_filt, 4
logits_filt.shape[0] # logits_filt.shape[0]
# get phrase # # get phrase
tokenlizer = self.grounding.tokenizer # tokenlizer = self.grounding.tokenizer
tokenized = tokenlizer(caption) # tokenized = tokenlizer(caption)
# build pred # # build pred
pred_phrases = [] # pred_phrases = []
for logit, box in zip(logits_filt, boxes_filt): # for logit, box in zip(logits_filt, boxes_filt):
pred_phrase = get_phrases_from_posmap(logit > self.text_threshold, tokenized, tokenlizer) # pred_phrase = get_phrases_from_posmap(logit > self.text_threshold, tokenized, tokenlizer)
if with_logits: # if with_logits:
pred_phrases.append(pred_phrase + f"({str(logit.max().item())[:4]})") # pred_phrases.append(pred_phrase + f"({str(logit.max().item())[:4]})")
else: # else:
pred_phrases.append(pred_phrase) # pred_phrases.append(pred_phrase)
return boxes_filt, pred_phrases # return boxes_filt, pred_phrases
def plot_boxes_to_image(self, image_pil, tgt): # def plot_boxes_to_image(self, image_pil, tgt):
H, W = tgt["size"] # H, W = tgt["size"]
boxes = tgt["boxes"] # boxes = tgt["boxes"]
labels = tgt["labels"] # labels = tgt["labels"]
assert len(boxes) == len(labels), "boxes and labels must have same length" # assert len(boxes) == len(labels), "boxes and labels must have same length"
draw = ImageDraw.Draw(image_pil) # draw = ImageDraw.Draw(image_pil)
mask = Image.new("L", image_pil.size, 0) # mask = Image.new("L", image_pil.size, 0)
mask_draw = ImageDraw.Draw(mask) # mask_draw = ImageDraw.Draw(mask)
# draw boxes and masks # # draw boxes and masks
for box, label in zip(boxes, labels): # for box, label in zip(boxes, labels):
# from 0..1 to 0..W, 0..H # # from 0..1 to 0..W, 0..H
box = box * torch.Tensor([W, H, W, H]) # box = box * torch.Tensor([W, H, W, H])
# from xywh to xyxy # # from xywh to xyxy
box[:2] -= box[2:] / 2 # box[:2] -= box[2:] / 2
box[2:] += box[:2] # box[2:] += box[:2]
# random color # # random color
color = tuple(np.random.randint(0, 255, size=3).tolist()) # color = tuple(np.random.randint(0, 255, size=3).tolist())
# draw # # draw
x0, y0, x1, y1 = box # x0, y0, x1, y1 = box
x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1) # x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1)
draw.rectangle([x0, y0, x1, y1], outline=color, width=6) # draw.rectangle([x0, y0, x1, y1], outline=color, width=6)
# draw.text((x0, y0), str(label), fill=color) # # draw.text((x0, y0), str(label), fill=color)
font = ImageFont.load_default() # font = ImageFont.load_default()
if hasattr(font, "getbbox"): # if hasattr(font, "getbbox"):
bbox = draw.textbbox((x0, y0), str(label), font) # bbox = draw.textbbox((x0, y0), str(label), font)
else: # else:
w, h = draw.textsize(str(label), font) # w, h = draw.textsize(str(label), font)
bbox = (x0, y0, w + x0, y0 + h) # bbox = (x0, y0, w + x0, y0 + h)
# bbox = draw.textbbox((x0, y0), str(label)) # # bbox = draw.textbbox((x0, y0), str(label))
draw.rectangle(bbox, fill=color) # draw.rectangle(bbox, fill=color)
draw.text((x0, y0), str(label), fill="white") # draw.text((x0, y0), str(label), fill="white")
mask_draw.rectangle([x0, y0, x1, y1], fill=255, width=2) # mask_draw.rectangle([x0, y0, x1, y1], fill=255, width=2)
return image_pil, mask # return image_pil, mask
@prompts(name="Detect the Give Object", # @prompts(name="Detect the Give Object",
description="useful when you only want to detect or find out given objects in the picture" # 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, " # "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") # "representing the image_path, the text description of the object to be found")
def inference(self, inputs): # def inference(self, inputs):
image_path, det_prompt = inputs.split(",") # image_path, det_prompt = inputs.split(",")
print(f"image_path={image_path}, text_prompt={det_prompt}") # print(f"image_path={image_path}, text_prompt={det_prompt}")
image_pil, image = self.load_image(image_path) # image_pil, image = self.load_image(image_path)
boxes_filt, pred_phrases = self.get_grounding_boxes(image, det_prompt) # boxes_filt, pred_phrases = self.get_grounding_boxes(image, det_prompt)
size = image_pil.size # size = image_pil.size
pred_dict = { # pred_dict = {
"boxes": boxes_filt, # "boxes": boxes_filt,
"size": [size[1], size[0]], # H,W # "size": [size[1], size[0]], # H,W
"labels": pred_phrases,} # "labels": pred_phrases,}
image_with_box = self.plot_boxes_to_image(image_pil, pred_dict)[0] # 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_path = get_new_image_name(image_path, func_name="detect-something")
updated_image = image_with_box.resize(size) # updated_image = image_with_box.resize(size)
updated_image.save(updated_image_path) # updated_image.save(updated_image_path)
print( # print(
f"\nProcessed ObejectDetecting, Input Image: {image_path}, Object to be Detect {det_prompt}, " # f"\nProcessed ObejectDetecting, Input Image: {image_path}, Object to be Detect {det_prompt}, "
f"Output Image: {updated_image_path}") # f"Output Image: {updated_image_path}")
return updated_image_path # return updated_image_path
class Inpainting: class Inpainting:

Write Preview
Loading…
Cancel
Save