diff --git a/README.md b/README.md index 066bc32..34470f7 100644 --- a/README.md +++ b/README.md @@ -1 +1,69 @@ -# ComfyUI_SemanticSAM \ No newline at end of file +# ComfyUI_SemanticSAM +This is a ComfyUI node based-on [Semantic-SAM](https://github.com/UX-Decoder/Semantic-SAM) official implementation. We provide a workflow node for one-click segment. Compared with SAM, Semantic-SAM has better fine-grained capabilities and more candidate masks. Users can take this node as the pre-node for inpainting to obtain the mask region. + +This node have been valided on Ubuntu-20.04 & CUDA-11.8. Due to detectron2, this node is currently not supported on Windows (unless the compatibility of Detectron2 on Windows is addressed) +![image](./examples/SemanticSAM.png) + +## Get Start +### 1. Prepare repo +```sh +cd ComfyUI/custom +git clone https://github.com/eastoc/ComfyUI_SemanticSAM +cd ComfyUI_SemanticSAM +git clone https://github.com/facebookresearch/detectron2 +git clone https://github.com/facebookresearch/Mask2Former +``` +Because detectron2 has not been updated, if the CUDA=11.8, you can +```sh +git clone https://github.com/johnnynunez/detectron2 +``` +or you can try other 3rd detectron2 implementation. + +### 2. Download models +download Semantic-SAM model to ComfyUI_SemanticSAM/ckpt + + + + + + + + + + + + + + + + + + + + + + + + +
NameTraining DatasetBackbone1-IoU@Multi-Granularity1-IoU@COCO(Max|Oracle)download
Semantic-SAM | configSA-1BSwinT88.154.5|73.8model
Semantic-SAM | configSA-1BSwinL89.055.1|74.1model
+ +### 3. Install environment +Install [PyTorch & torchvision](https://pytorch.org/) through the official method. We have been valid that from pytorch 1.13.0 to 2.3.0. + +Install Semantic-SAM dependencies +```sh +pip install -r requirements.txt +``` + +Install detectron2 +```sh +cd detectron2 +pip install -e . +cd .. +``` +Install Mask2Former +```sh +cd Mask2Former/mask2former/modeling/pixel_decoder/ops +sh make.sh +``` + diff --git a/configs/semantic_sam_only_sa-1b_swinL.yaml b/configs/semantic_sam_only_sa-1b_swinL.yaml index c8c7399..294bfd6 100755 --- a/configs/semantic_sam_only_sa-1b_swinL.yaml +++ b/configs/semantic_sam_only_sa-1b_swinL.yaml @@ -38,7 +38,7 @@ MODEL: AUTOGRESSIVE: True BACKBONE: NAME: swin - PRETRAINED: 'custom_nodes/ComfyUI_Semantic_SAM/ckpt/swin_large_patch4_window12_384_22k.pth' + PRETRAINED: 'custom_nodes/ComfyUI_SemanticSAM/ckpt/swin_large_patch4_window12_384_22k.pth' LOAD_PRETRAINED: true SWIN: PRETRAIN_IMG_SIZE: 384 diff --git a/configs/semantic_sam_only_sa-1b_swinT.yaml b/configs/semantic_sam_only_sa-1b_swinT.yaml index e0bce43..53ff061 100755 --- a/configs/semantic_sam_only_sa-1b_swinT.yaml +++ b/configs/semantic_sam_only_sa-1b_swinT.yaml @@ -38,7 +38,7 @@ MODEL: AUTOGRESSIVE: True BACKBONE: NAME: swin - PRETRAINED: 'https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth' + PRETRAINED: 'custom_nodes/ComfyUI_SemanticSAM/ckpt/swin_tiny_patch4_window7_224.pth' LOAD_PRETRAINED: true SWIN: PRETRAIN_IMG_SIZE: 224 diff --git a/examples/SemanticSAM.png b/examples/SemanticSAM.png new file mode 100644 index 0000000..957f03e Binary files /dev/null and b/examples/SemanticSAM.png differ diff --git a/examples/workflow.png b/examples/workflow.png new file mode 100644 index 0000000..4928387 Binary files /dev/null and b/examples/workflow.png differ diff --git a/node.py b/node.py index 16f5d27..58c1bb1 100755 --- a/node.py +++ b/node.py @@ -10,8 +10,8 @@ import logging import cv2 sys.path.append(os.path.dirname(os.path.abspath(__file__))) -sam_model_list = {"L": {"model_url": "custom_nodes/ComfyUI_Semantic_SAM/ckpt/swinl_only_sam_many2many.pth"}, - "T": {"model_url": "custom_nodes/ComfyUI_Semantic_SAM/ckpt/swint_only_sam_many2many.pth"}} +sam_model_list = {"L": {"model_url": "custom_nodes/ComfyUI_SemanticSAM/ckpt/swinl_only_sam_many2many.pth"}, + "T": {"model_url": "custom_nodes/ComfyUI_SemanticSAM/ckpt/swint_only_sam_many2many.pth"}} def list_sam_model(): return list(sam_model_list.keys()) @@ -51,7 +51,7 @@ class PointPrompt(): FUNCTION = "main" def main(self, x_coord, y_coord): - return ([[x_coord,x_coord]],) + return ([[x_coord,y_coord]],) class SemanticSAMLoader(): @classmethod @@ -199,4 +199,4 @@ def mask2rgba(image, mask): #cv2.imwrite(f'mask_{i}.png', rgba_image*255) rgba_tensor = torch.from_numpy(rgba_image) rgba_tensor = torch.unsqueeze(rgba_tensor, dim=0) - return rgba_tensor \ No newline at end of file + return rgba_tensor diff --git a/semantic_sam/BaseModel.py b/semantic_sam/BaseModel.py index 5bcd05e..937a2fd 100755 --- a/semantic_sam/BaseModel.py +++ b/semantic_sam/BaseModel.py @@ -4,7 +4,7 @@ import logging import torch import torch.nn as nn -from ComfyUI_Semantic_SAM.utils.model import align_and_update_state_dicts +from ComfyUI_SemanticSAM.utils.model import align_and_update_state_dicts logger = logging.getLogger(__name__) diff --git a/semantic_sam/build_semantic_sam.py b/semantic_sam/build_semantic_sam.py index 6663ecd..b277e83 100755 --- a/semantic_sam/build_semantic_sam.py +++ b/semantic_sam/build_semantic_sam.py @@ -12,12 +12,12 @@ from torchvision import transforms import torch import os -from ComfyUI_Semantic_SAM.utils.arguments import load_opt_from_config_file -from ComfyUI_Semantic_SAM.semantic_sam.BaseModel import BaseModel -from ComfyUI_Semantic_SAM.semantic_sam import build_model -from ComfyUI_Semantic_SAM.tasks.automatic_mask_generator import SemanticSamAutomaticMaskGenerator -from ComfyUI_Semantic_SAM.tasks.interactive_idino_m2m_auto import show_anns -from ComfyUI_Semantic_SAM.tasks.interactive_predictor import SemanticSAMPredictor +from ComfyUI_SemanticSAM.utils.arguments import load_opt_from_config_file +from ComfyUI_SemanticSAM.semantic_sam.BaseModel import BaseModel +from ComfyUI_SemanticSAM.semantic_sam import build_model +from ComfyUI_SemanticSAM.tasks.automatic_mask_generator import SemanticSamAutomaticMaskGenerator +from ComfyUI_SemanticSAM.tasks.interactive_idino_m2m_auto import show_anns +from ComfyUI_SemanticSAM.tasks.interactive_predictor import SemanticSAMPredictor def prepare_image(image_pth): @@ -45,8 +45,8 @@ def build_semantic_sam(model_type, ckpt): current_path = os.getcwd() print("当前路径是:", current_path) - cfgs={'T':"custom_nodes/ComfyUI_Semantic_SAM/configs/semantic_sam_only_sa-1b_swinT.yaml", - 'L':"custom_nodes/ComfyUI_Semantic_SAM/configs/semantic_sam_only_sa-1b_swinL.yaml"} + cfgs={'T':"custom_nodes/ComfyUI_SemanticSAM/configs/semantic_sam_only_sa-1b_swinT.yaml", + 'L':"custom_nodes/ComfyUI_SemanticSAM/configs/semantic_sam_only_sa-1b_swinL.yaml"} sam_cfg=cfgs[model_type] opt = load_opt_from_config_file(sam_cfg) diff --git a/semantic_sam/language/LangEncoder/transformer.py b/semantic_sam/language/LangEncoder/transformer.py index 6d1c070..e8c6d73 100755 --- a/semantic_sam/language/LangEncoder/transformer.py +++ b/semantic_sam/language/LangEncoder/transformer.py @@ -12,7 +12,7 @@ from timm.models.layers import DropPath, trunc_normal_ from .registry import register_lang_encoder from detectron2.utils.comm import is_main_process -from ComfyUI_Semantic_SAM.utils.model import register_norm_module +from ComfyUI_SemanticSAM.utils.model import register_norm_module logger = logging.getLogger(__name__) diff --git a/semantic_sam/language/encoder.py b/semantic_sam/language/encoder.py index dbb3e4f..cd1dff9 100755 --- a/semantic_sam/language/encoder.py +++ b/semantic_sam/language/encoder.py @@ -7,7 +7,7 @@ from timm.models.layers import trunc_normal_ from .registry import register_model from ..utils import configurable from .LangEncoder import build_tokenizer, build_lang_encoder -from ComfyUI_Semantic_SAM.utils.prompt_engineering import prompt_engineering, get_prompt_templates +from ComfyUI_SemanticSAM.utils.prompt_engineering import prompt_engineering, get_prompt_templates class LanguageEncoder(nn.Module): diff --git a/semantic_sam/language/vlpencoder.py b/semantic_sam/language/vlpencoder.py index 021e934..f48ff3b 100755 --- a/semantic_sam/language/vlpencoder.py +++ b/semantic_sam/language/vlpencoder.py @@ -13,7 +13,7 @@ from timm.models.layers import trunc_normal_ from .registry import register_model from ..utils import configurable from .LangEncoder import build_tokenizer, build_lang_encoder -from ComfyUI_Semantic_SAM.utils.prompt_engineering import prompt_engineering, get_prompt_templates +from ComfyUI_SemanticSAM.utils.prompt_engineering import prompt_engineering, get_prompt_templates class LanguageEncoder(nn.Module): diff --git a/semantic_sam/utils/misc.py b/semantic_sam/utils/misc.py index 1caa0fa..0887649 100755 --- a/semantic_sam/utils/misc.py +++ b/semantic_sam/utils/misc.py @@ -20,7 +20,7 @@ import torch.distributed as dist import torchvision from torch import Tensor -from ComfyUI_Semantic_SAM.utils.constants import * +from ComfyUI_SemanticSAM.utils.constants import * def get_iou(gt_masks, pred_masks, ignore_label=-1): rev_ignore_mask = ~(gt_masks == ignore_label) diff --git a/tasks/automatic_mask_generator.py b/tasks/automatic_mask_generator.py index b44f569..13f0cb6 100755 --- a/tasks/automatic_mask_generator.py +++ b/tasks/automatic_mask_generator.py @@ -12,7 +12,7 @@ from typing import Any, Dict, List, Optional, Tuple # from # from .modeling import Sam # from .predictor import SamPredictor -from ComfyUI_Semantic_SAM.utils.sam_utils.amg import ( +from ComfyUI_SemanticSAM.utils.sam_utils.amg import ( MaskData, area_from_rle, batch_iterator, diff --git a/tasks/interactive_idino_m2m.py b/tasks/interactive_idino_m2m.py index 2b1e9dc..48f893c 100755 --- a/tasks/interactive_idino_m2m.py +++ b/tasks/interactive_idino_m2m.py @@ -8,7 +8,7 @@ import torch import numpy as np from torchvision import transforms -from ComfyUI_Semantic_SAM.utils.visualizer import Visualizer +from ComfyUI_SemanticSAM.utils.visualizer import Visualizer from typing import Tuple from PIL import Image from detectron2.data import MetadataCatalog diff --git a/tasks/interactive_idino_m2m_auto.py b/tasks/interactive_idino_m2m_auto.py index 4b65d0e..30282aa 100755 --- a/tasks/interactive_idino_m2m_auto.py +++ b/tasks/interactive_idino_m2m_auto.py @@ -8,7 +8,7 @@ import torch import numpy as np from torchvision import transforms -from ComfyUI_Semantic_SAM.utils.visualizer import Visualizer +from ComfyUI_SemanticSAM.utils.visualizer import Visualizer from typing import Tuple from PIL import Image from detectron2.data import MetadataCatalog diff --git a/tasks/interactive_predictor.py b/tasks/interactive_predictor.py index 6c516ba..3beb115 100755 --- a/tasks/interactive_predictor.py +++ b/tasks/interactive_predictor.py @@ -1,7 +1,7 @@ import torch import numpy as np from torchvision import transforms -from ComfyUI_Semantic_SAM.utils.visualizer import Visualizer +from ComfyUI_SemanticSAM.utils.visualizer import Visualizer from typing import Tuple from PIL import Image from detectron2.data import MetadataCatalog diff --git a/workflow/workflow.json b/workflow/workflow.json new file mode 100644 index 0000000..02d8796 --- /dev/null +++ b/workflow/workflow.json @@ -0,0 +1,239 @@ +{ + "last_node_id": 21, + "last_link_id": 22, + "nodes": [ + { + "id": 17, + "type": "LoadImage", + "pos": [ + 40, + 302 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 18 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "truck.jpg", + "image" + ] + }, + { + "id": 18, + "type": "PreviewImage", + "pos": [ + 1405, + 225 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 19 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 20, + "type": "SemanticSAMLoader", + "pos": [ + 435.5874909756752, + 49.649601025391036 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "SemanticSAM_Model", + "type": "SemanticSAM_Model", + "links": [ + 21 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "SemanticSAMLoader" + }, + "widgets_values": [ + "L" + ] + }, + { + "id": 16, + "type": "SemanticSAMSegment", + "pos": [ + 945, + 224 + ], + "size": { + "0": 315, + "1": 146 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "SemanticSAM_Model", + "link": 21, + "slot_index": 0 + }, + { + "name": "image", + "type": "IMAGE", + "link": 18 + }, + { + "name": "points", + "type": "POINTS", + "link": 22 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 19 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "SemanticSAMSegment" + }, + "widgets_values": [ + 0, + 0.5, + 6 + ] + }, + { + "id": 21, + "type": "PointPrompt", + "pos": [ + 422, + 400 + ], + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "POINTS", + "type": "POINTS", + "links": [ + 22 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PointPrompt" + }, + "widgets_values": [ + 375, + 500 + ] + } + ], + "links": [ + [ + 18, + 17, + 0, + 16, + 1, + "IMAGE" + ], + [ + 19, + 16, + 0, + 18, + 0, + "IMAGE" + ], + [ + 21, + 20, + 0, + 16, + 0, + "SemanticSAM_Model" + ], + [ + 22, + 21, + 0, + 16, + 2, + "POINTS" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9090909090909091, + "offset": { + "0": 75.92159993341575, + "1": 162.98676261097265 + } + } + }, + "version": 0.4 +} \ No newline at end of file