Add SAM support

extras/GroundingDINO/config/GroundingDINO_SwinT_OGC.py

@@ -0,0 +1,43 @@
+batch_size = 1
+modelname = "groundingdino"
+backbone = "swin_T_224_1k"
+position_embedding = "sine"
+pe_temperatureH = 20
+pe_temperatureW = 20
+return_interm_indices = [1, 2, 3]
+backbone_freeze_keywords = None
+enc_layers = 6
+dec_layers = 6
+pre_norm = False
+dim_feedforward = 2048
+hidden_dim = 256
+dropout = 0.0
+nheads = 8
+num_queries = 900
+query_dim = 4
+num_patterns = 0
+num_feature_levels = 4
+enc_n_points = 4
+dec_n_points = 4
+two_stage_type = "standard"
+two_stage_bbox_embed_share = False
+two_stage_class_embed_share = False
+transformer_activation = "relu"
+dec_pred_bbox_embed_share = True
+dn_box_noise_scale = 1.0
+dn_label_noise_ratio = 0.5
+dn_label_coef = 1.0
+dn_bbox_coef = 1.0
+embed_init_tgt = True
+dn_labelbook_size = 2000
+max_text_len = 256
+text_encoder_type = "bert-base-uncased"
+use_text_enhancer = True
+use_fusion_layer = True
+use_checkpoint = True
+use_transformer_ckpt = True
+use_text_cross_attention = True
+text_dropout = 0.0
+fusion_dropout = 0.0
+fusion_droppath = 0.1
+sub_sentence_present = True

extras/inpaint_mask.py

@@ -1,4 +1,39 @@
+from PIL import Image
+import numpy as np
+import torch
 from rembg import remove, new_session
+from groundingdino.util.inference import Model as GroundingDinoModel
+
+from modules.model_loader import load_file_from_url
+from modules.config import path_inpaint
+
+config_file = 'extras/GroundingDINO/config/GroundingDINO_SwinT_OGC.py'
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+groundingdino_model = None
+
+
+def run_grounded_sam(input_image, text_prompt, box_threshold, text_threshold):
+
+    global groundingdino_model
+
+    if groundingdino_model is None:
+        filename = load_file_from_url(
+            url="https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth",
+            model_dir=path_inpaint)
+        groundingdino_model = GroundingDinoModel(model_config_path=config_file, model_checkpoint_path=filename, device=device)
+
+
+    # run grounding dino model
+    boxes, _ = groundingdino_model.predict_with_caption(
+        image=np.array(input_image),
+        caption=text_prompt,
+        box_threshold=box_threshold,
+        text_threshold=text_threshold
+    )
+
+    return boxes.xyxy
 
 
 def generate_mask_from_image(image, mask_model, extras):
@@ -8,9 +43,15 @@ def generate_mask_from_image(image, mask_model, extras):
     if 'image' in image:
         image = image['image']
 
+    if mask_model == 'sam':
+        boxes = run_grounded_sam(Image.fromarray(image), extras['sam_prompt_text'], box_threshold=extras['box_threshold'], text_threshold=extras['text_threshold'])
+        extras['sam_prompt'] = []
+        for idx, box in enumerate(boxes):
+            extras['sam_prompt'] += [{"type": "rectangle", "data": box.tolist()}]
+
     return remove(
         image,
-        session=new_session(mask_model),
+        session=new_session(mask_model, **extras),
         only_mask=True,
         **extras
     )