53 lines
1.2 KiB
Python
53 lines
1.2 KiB
Python
import folder_paths
|
|
import torch
|
|
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
|
|
from einops import rearrange, repeat
|
|
|
|
|
|
class SAM_Load_Embedding:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"filename": ("STRING", {
|
|
"multiline": False,
|
|
"default": "embeddings"
|
|
}),
|
|
"embedding_id": ("STRING", {
|
|
"multiline": False,
|
|
"default": "embedding"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("EMBEDDINGS",)
|
|
RETURN_NAMES = ("EMBEDDINGS",)
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
FUNCTION = "process"
|
|
|
|
CATEGORY = "image"
|
|
|
|
def process(self, filename, embedding_id):
|
|
import json
|
|
import numpy as np
|
|
|
|
data = {}
|
|
with open(filename, 'r') as f:
|
|
data = json.load(f)
|
|
|
|
# Convert list to numpy ndarray
|
|
data['image_embedding'] = np.array(data['image_embedding'])
|
|
|
|
return (data, )
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SAM_Load_Embedding": SAM_Load_Embedding
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SAM_Load_Embedding": "SAM_Load_Embedding "
|
|
}
|