diff --git a/__init__.py b/__init__.py index 37557e9..f0b7a7f 100644 --- a/__init__.py +++ b/__init__.py @@ -2,12 +2,12 @@ from .node import * from .install import * NODE_CLASS_MAPPINGS = { - '🧩 SAMModelLoader': SAMModelLoader, - '🧩 DinoModelLoader': GroundingDinoModelLoader, - '🧩 DinoSAMSegment': GroundingDinoSAMSegment, - '🧩 InvertMask': InvertMask, - "🧩 IsMaskEmpty": IsMaskEmptyNode, - "🧩 SegmentAnything": SegmentAnything, + '🧩 SAMModel Loader': SAMModelLoader, + '🧩 DinoModel Loader': GroundingDinoModelLoader, + '🧩 Dino SAMSegment': GroundingDinoSAMSegment, + '🧩 Invert Mask': InvertMask, + "🧩 Is Mask Empty": IsMaskEmptyNode, + "🧩 Segment Anything": SegmentAnything, } __all__ = ['NODE_CLASS_MAPPINGS'] diff --git a/__pycache__/__init__.cpython-311.pyc b/__pycache__/__init__.cpython-311.pyc index 6524ca5..294b960 100644 Binary files a/__pycache__/__init__.cpython-311.pyc and b/__pycache__/__init__.cpython-311.pyc differ diff --git a/__pycache__/node.cpython-311.pyc b/__pycache__/node.cpython-311.pyc index 893e8c5..617de81 100644 Binary files a/__pycache__/node.cpython-311.pyc and b/__pycache__/node.cpython-311.pyc differ diff --git a/node.py b/node.py index cbf6334..bf336c3 100644 --- a/node.py +++ b/node.py @@ -3,7 +3,7 @@ import sys sys.path.append( os.path.dirname(os.path.abspath(__file__)) ) - +from nodes import NODE_CLASS_MAPPINGS as ALL_NODE import copy import torch import numpy as np @@ -383,7 +383,11 @@ class SegmentAnything: sam_load_model = SAMModelLoader().main(sam_model)[0] dino_load_model = GroundingDinoModelLoader().main(dino_model)[0] output = GroundingDinoSAMSegment().main(dino_load_model, sam_load_model, image, prompt, threshold, threshold) - return (output) + mask = output[1] + invert_mask = (1.0 - mask).to(image.device) + alpha_image = ALL_NODE["JoinImageWithAlpha"]().execute(image, invert_mask)[0] + ui = ALL_NODE["PreviewImage"]().save_images(alpha_image)["ui"] + return {"ui":ui, "result": output} class InvertMask: @classmethod