| |
|
|
| import os |
| import sys |
| sys.path.append( |
| os.path.dirname(os.path.abspath(__file__)) |
| ) |
|
|
| import copy |
| import torch |
| import numpy as np |
| from PIL import Image |
| import logging |
| from torch.hub import download_url_to_file |
| from urllib.parse import urlparse |
| import folder_paths |
| import comfy.model_management |
| from sam_hq.predictor import SamPredictorHQ |
| from sam_hq.build_sam_hq import sam_model_registry |
|
|
| import glob |
| import folder_paths |
|
|
| logger = logging.getLogger('comfyui_segment_anything') |
|
|
| sam_model_dir_name = "sams" |
| sam_model_list = { |
| "sam_vit_h (2.56GB)": { |
| "model_url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth" |
| }, |
| "sam_vit_l (1.25GB)": { |
| "model_url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth" |
| }, |
| "sam_vit_b (375MB)": { |
| "model_url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth" |
| }, |
| "sam_hq_vit_h (2.57GB)": { |
| "model_url": "https://huggingface.co/lkeab/hq-sam/resolve/main/sam_hq_vit_h.pth" |
| }, |
| "sam_hq_vit_l (1.25GB)": { |
| "model_url": "https://huggingface.co/lkeab/hq-sam/resolve/main/sam_hq_vit_l.pth" |
| }, |
| "sam_hq_vit_b (379MB)": { |
| "model_url": "https://huggingface.co/lkeab/hq-sam/resolve/main/sam_hq_vit_b.pth" |
| }, |
| "mobile_sam(39MB)": { |
| "model_url": "https://github.com/ChaoningZhang/MobileSAM/blob/master/weights/mobile_sam.pt" |
| } |
| } |
|
|
| groundingdino_model_dir_name = "grounding-dino" |
| groundingdino_model_list = { |
| "GroundingDINO_SwinT_OGC (694MB)": { |
| "config_url": "https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/GroundingDINO_SwinT_OGC.cfg.py", |
| "model_url": "https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/groundingdino_swint_ogc.pth", |
| }, |
| "GroundingDINO_SwinB (938MB)": { |
| "config_url": "https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/GroundingDINO_SwinB.cfg.py", |
| "model_url": "https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/groundingdino_swinb_cogcoor.pth" |
| }, |
| } |
|
|
| def get_bert_base_uncased_model_path(): |
| comfy_bert_model_base = os.path.join(folder_paths.models_dir, 'bert-base-uncased') |
| if glob.glob(os.path.join(comfy_bert_model_base, '**/model.safetensors'), recursive=True): |
| print('grounding-dino is using models/bert-base-uncased') |
| return comfy_bert_model_base |
| return 'bert-base-uncased' |
|
|
| def list_sam_model(): |
| return list(sam_model_list.keys()) |
|
|
| def load_sam_model(model_name): |
| sam_checkpoint_path = get_local_filepath( |
| sam_model_list[model_name]["model_url"], sam_model_dir_name) |
| model_file_name = os.path.basename(sam_checkpoint_path) |
| model_type = model_file_name.split('.')[0] |
| if 'hq' not in model_type and 'mobile' not in model_type: |
| model_type = '_'.join(model_type.split('_')[:-1]) |
| sam = sam_model_registry[model_type](checkpoint=sam_checkpoint_path) |
| sam_device = comfy.model_management.get_torch_device() |
| sam.to(device=sam_device) |
| sam.eval() |
| sam.model_name = model_file_name |
| return sam |
|
|
| def get_local_filepath(url, dirname, local_file_name=None): |
| if not local_file_name: |
| parsed_url = urlparse(url) |
| local_file_name = os.path.basename(parsed_url.path) |
|
|
| destination = folder_paths.get_full_path(dirname, local_file_name) |
| if destination: |
| logger.warn(f'using extra model: {destination}') |
| return destination |
|
|
| folder = os.path.join(folder_paths.models_dir, dirname) |
| if not os.path.exists(folder): |
| os.makedirs(folder) |
|
|
| destination = os.path.join(folder, local_file_name) |
| if not os.path.exists(destination): |
| logger.warn(f'downloading {url} to {destination}') |
| download_url_to_file(url, destination) |
| return destination |
|
|
| def load_groundingdino_model(model_name): |
| from local_groundingdino.util.utils import clean_state_dict as local_groundingdino_clean_state_dict |
| from local_groundingdino.util.slconfig import SLConfig as local_groundingdino_SLConfig |
| from local_groundingdino.models import build_model as local_groundingdino_build_model |
| dino_model_args = local_groundingdino_SLConfig.fromfile( |
| get_local_filepath( |
| groundingdino_model_list[model_name]["config_url"], |
| groundingdino_model_dir_name |
| ), |
| ) |
|
|
| if dino_model_args.text_encoder_type == 'bert-base-uncased': |
| dino_model_args.text_encoder_type = get_bert_base_uncased_model_path() |
| |
| dino = local_groundingdino_build_model(dino_model_args) |
| checkpoint = torch.load( |
| get_local_filepath( |
| groundingdino_model_list[model_name]["model_url"], |
| groundingdino_model_dir_name, |
| ), |
| ) |
| dino.load_state_dict(local_groundingdino_clean_state_dict( |
| checkpoint['model']), strict=False) |
| device = comfy.model_management.get_torch_device() |
| dino.to(device=device) |
| dino.eval() |
| return dino |
|
|
| def list_groundingdino_model(): |
| return list(groundingdino_model_list.keys()) |
|
|
| def groundingdino_predict( |
| dino_model, |
| image, |
| prompt, |
| threshold |
| ): |
| from local_groundingdino.datasets import transforms as T |
| def load_dino_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 get_grounding_output(model, image, caption, box_threshold): |
| caption = caption.lower() |
| caption = caption.strip() |
| if not caption.endswith("."): |
| caption = caption + "." |
| device = comfy.model_management.get_torch_device() |
| image = image.to(device) |
| with torch.no_grad(): |
| outputs = model(image[None], captions=[caption]) |
| logits = outputs["pred_logits"].sigmoid()[0] |
| boxes = outputs["pred_boxes"][0] |
| |
| 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] |
| return boxes_filt.cpu() |
|
|
| dino_image = load_dino_image(image.convert("RGB")) |
| boxes_filt = get_grounding_output( |
| dino_model, dino_image, prompt, threshold |
| ) |
| H, W = image.size[1], image.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] |
| return boxes_filt |
|
|
| def create_tensor_output(image_np, masks, boxes_filt): |
| output_masks, output_images = [], [] |
| boxes_filt = boxes_filt.numpy().astype(int) if boxes_filt is not None else None |
| for mask in masks: |
| image_np_copy = copy.deepcopy(image_np) |
| image_np_copy[~np.any(mask, axis=0)] = np.array([0, 0, 0, 0]) |
| output_image, output_mask = split_image_mask( |
| Image.fromarray(image_np_copy)) |
| output_masks.append(output_mask) |
| output_images.append(output_image) |
| return (output_images, output_masks) |
|
|
| def split_image_mask(image): |
| image_rgb = image.convert("RGB") |
| image_rgb = np.array(image_rgb).astype(np.float32) / 255.0 |
| image_rgb = torch.from_numpy(image_rgb)[None,] |
| if 'A' in image.getbands(): |
| mask = np.array(image.getchannel('A')).astype(np.float32) / 255.0 |
| mask = torch.from_numpy(mask)[None,] |
| else: |
| mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") |
| return (image_rgb, mask) |
|
|
| def sam_segment( |
| sam_model, |
| image, |
| boxes |
| ): |
| if boxes.shape[0] == 0: |
| return None |
| sam_is_hq = False |
| |
| if hasattr(sam_model, 'model_name') and 'hq' in sam_model.model_name: |
| sam_is_hq = True |
| predictor = SamPredictorHQ(sam_model, sam_is_hq) |
| image_np = np.array(image) |
| image_np_rgb = image_np[..., :3] |
| predictor.set_image(image_np_rgb) |
| transformed_boxes = predictor.transform.apply_boxes_torch( |
| boxes, image_np.shape[:2]) |
| sam_device = comfy.model_management.get_torch_device() |
| masks, _, _ = predictor.predict_torch( |
| point_coords=None, |
| point_labels=None, |
| boxes=transformed_boxes.to(sam_device), |
| multimask_output=False) |
| masks = masks.permute(1, 0, 2, 3).cpu().numpy() |
| return create_tensor_output(image_np, masks, boxes) |
|
|
