nodes/eastoc-ComfyUI_SemanticSAM
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update readme

Browse Source
This commit is contained in:
eastoc
2024-08-14 03:19:26 +08:00
parent 1c3cda4e92
commit c28fa92e40
17 changed files with 331 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+69 -1
View File
@@ -1 +1,69 @@
# ComfyUI_SemanticSAM
# 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
<table><tbody>
<!-- START TABLE -->
<!-- TABLE HEADER -->
<th valign="bottom">Name</th>
<th valign="bottom">Training Dataset</th>
<th valign="bottom">Backbone</th>
<th valign="bottom">1-IoU@Multi-Granularity</th>
<th valign="bottom">1-IoU@COCO(Max|Oracle)</th>
<th valign="bottom">download</th>
<tr><td align="left">Semantic-SAM | <a href="configs/semantic_sam_only_sa-1b_swinT.yaml">config</a></td>
<td align="center">SA-1B</td>
<td align="center">SwinT</td>
<td align="center">88.1</td>
<td align="center">54.5|73.8</td>
<td align="center"><a href="https://github.com/UX-Decoder/Semantic-SAM/releases/download/checkpoint/swint_only_sam_many2many.pth">model</a></td>
<tr><td align="left">Semantic-SAM | <a href="configs/semantic_sam_only_sa-1b_swinL.yaml">config</a></td>
<td align="center">SA-1B</td>
<td align="center">SwinL</td>
<td align="center">89.0</td>
<td align="center">55.1|74.1</td>
<td align="center"><a href="https://github.com/UX-Decoder/Semantic-SAM/releases/download/checkpoint/swinl_only_sam_many2many.pth">model</a></td>
</tbody></table>
### 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
```
configs/semantic_sam_only_sa-1b_swinL.yaml
+1 -1
View File
@@ -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
configs/semantic_sam_only_sa-1b_swinT.yaml
+1 -1
View File
@@ -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
examples/SemanticSAM.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 143 KiB

examples/workflow.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 193 KiB

node.py
+4 -4
View File
@@ -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
return rgba_tensor
semantic_sam/BaseModel.py
+1 -1
View File
@@ -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__)
semantic_sam/build_semantic_sam.py
+8 -8
View File
@@ -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)
semantic_sam/language/LangEncoder/transformer.py
+1 -1
View File
@@ -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__)
semantic_sam/language/encoder.py
+1 -1
View File
@@ -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):
semantic_sam/language/vlpencoder.py
+1 -1
View File
@@ -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):
semantic_sam/utils/misc.py
+1 -1
View File
@@ -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)
tasks/automatic_mask_generator.py
+1 -1
View File
@@ -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,
tasks/interactive_idino_m2m.py
+1 -1
View File
@@ -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
tasks/interactive_idino_m2m_auto.py
+1 -1
View File
@@ -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
tasks/interactive_predictor.py
+1 -1
View File
@@ -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
workflow/workflow.json
+239
View File
@@ -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
}