mask_metadata and requirement.txt

grounded_sam_demo.py

@@ -3,6 +3,7 @@
 import copy
 
 import numpy as np
+import json
 import torch
 from PIL import Image, ImageDraw, ImageFont
 
@@ -98,6 +99,35 @@ def show_box(box, ax, label):
     ax.text(x0, y0, label)
 
 
+def save_mask_data(output_dir, mask_list, box_list, label_list):
+    value = 0  # 0 for background
+
+    mask_img = torch.zeros(mask_list.shape[-2:])
+    for idx, mask in enumerate(mask_list):
+        mask_img[mask.cpu().numpy()[0] == True] = value + idx + 1
+    plt.figure(figsize=(10, 10))
+    plt.imshow(mask_img.numpy())
+    plt.axis('off')
+    plt.savefig(os.path.join(output_dir, 'mask.jpg'), bbox_inches="tight", dpi=300, pad_inches=0.0)
+
+    json_data = [{
+        'value': value,
+        'label': 'background'
+    }]
+    for label, box in zip(label_list, box_list):
+        value += 1
+        name, logit = label.split('(')
+        logit = logit[:-1] # the last is ')'
+        json_data.append({
+            'value': value,
+            'label': name,
+            'logit': float(logit),
+            'box': box.numpy().tolist(),
+        })
+    with open(os.path.join(output_dir, 'mask.json'), 'w') as f:
+        json.dump(json_data, f)
+
+
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser("Grounded-Segment-Anything Demo", add_help=True)
@@ -176,6 +206,12 @@ def show_box(box, ax, label):
         show_mask(mask.cpu().numpy(), plt.gca(), random_color=True)
     for box, label in zip(boxes_filt, pred_phrases):
         show_box(box.numpy(), plt.gca(), label)
+
     plt.axis('off')
-    plt.savefig(os.path.join(output_dir, "grounded_sam_output.jpg"), bbox_inches="tight")
+    plt.savefig(
+        os.path.join(output_dir, "grounded_sam_output.jpg"), 
+        bbox_inches="tight", dpi=300, pad_inches=0.0
+    )
+
+    save_mask_data(output_dir, masks, boxes_filt, pred_phrases)
 

requirements.txt

@@ -0,0 +1,20 @@
+addict
+diffusers
+gradio
+huggingface_hub
+matplotlib
+numpy
+onnxruntime
+opencv_python
+Pillow
+pycocotools
+PyYAML
+requests
+setuptools
+supervision
+termcolor
+timm
+torch
+torchvision
+transformers
+yapf