fix

.idea/.gitignore

@@ -0,0 +1,8 @@
+# 默认忽略的文件
+/shelf/
+/workspace.xml
+# 基于编辑器的 HTTP 客户端请求
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/comfyui-segment-anything-marko.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="NUMPY" />
+    <option name="myDocStringFormat" value="NumPy" />
+  </component>
+</module>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="Black">
+    <option name="sdkName" value="Python 3.12" />
+  </component>
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.12" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/comfyui-segment-anything-marko.iml" filepath="$PROJECT_DIR$/.idea/comfyui-segment-anything-marko.iml" />
+    </modules>
+  </component>
+</project>

__init__.py

@@ -0,0 +1,12 @@
+from .node import *
+from .install import *
+
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "AutomaticMask(segment anything)": AutomaticMask
+}
+
+__all__ = ['NODE_CLASS_MAPPINGS']
+
+

install.py

@@ -0,0 +1,17 @@
+import sys
+import os.path
+import subprocess
+
+custom_nodes_path = os.path.dirname(os.path.abspath(__file__))
+
+def build_pip_install_cmds(args):
+    if "python_embeded" in sys.executable or "python_embedded" in sys.executable:
+        return [sys.executable, '-s', '-m', 'pip', 'install'] + args
+    else:
+        return [sys.executable, '-m', 'pip', 'install'] + args
+
+def ensure_package():
+    cmds = build_pip_install_cmds(['-r', 'requirements.txt'])
+    subprocess.run(cmds, cwd=custom_nodes_path)
+
+ensure_package()

node.py

@@ -0,0 +1,76 @@
+import sys
+import os
+import numpy as np
+from PIL import Image
+import torch
+import matplotlib.pyplot as plt
+import cv2
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
+import folder_paths
+
+sys.path.append(
+    os.path.dirname(os.path.abspath(__file__))
+)
+
+def show_anns(anns, image_shape):
+    if len(anns) == 0:
+        return
+    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
+
+    img = np.ones((image_shape[0], image_shape[1], 4))
+    img[:,:,3] = 0
+    for ann in sorted_anns:
+        m = ann['segmentation']
+        color_mask = np.concatenate([np.random.random(3), [0.35]])
+        img[m] = color_mask
+    # 将带有标注的numpy图像转化为torch张量
+    annotated_img_tensor = torch.from_numpy(img)
+
+    return annotated_img_tensor
+
+class AutomaticMask:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+            },
+        }
+
+    # RETURN_NAMES = ("IMAGE",)
+    FUNCTION = "main"
+    CATEGORY = "segment_anything"
+    RETURN_TYPES = ("IMAGE",)
+
+    def main(self, image):
+        sam_checkpoint = folder_paths.get_full_path('sams', 'sam_vit_h_4b8939.pth')
+        model_type = "vit_h"
+
+        device = "cuda"
+
+        sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+
+        sam.to(device=device)
+
+        mask_generator = SamAutomaticMaskGenerator(sam)
+
+        image_res = []
+        for item in image:
+            image_shape = (item.shape[0], item.shape[1])
+            print(image_shape)
+            item = Image.fromarray(
+                np.clip(255. * item.cpu().numpy(), 0, 255).astype(np.uint8)).convert('RGBA')
+            image_np = np.array(item)
+            image_np_rgb = image_np[..., :3]
+
+
+            # 生成蒙版
+            masks = mask_generator.generate(image_np_rgb)
+            annotated_image_tensor = show_anns(masks, image_shape)
+
+            image_res.append(annotated_image_tensor)
+
+        return (image_res,)

requirements.txt

@@ -0,0 +1,5 @@
+segment_anything
+cv2import
+matplotlib
+torch
+numpy