88 lines
2.5 KiB
Python
88 lines
2.5 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from sam2.build_sam import load_model
|
|
from sam2.sam2_image_predictor import SAM2ImagePredictor
|
|
|
|
#waiting for checkpoints
|
|
#https://huggingface.co/spaces/yunyangx/EfficientTAM/tree/main
|
|
class EfficientTAMLoader:
|
|
"""Loads an EfficientTAM model for use in ComfyUI"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (["efficienttam_ti", "efficienttam_s"], {"default": "efficienttam_ti"}),
|
|
"device": (["cuda", "cpu"], {"default": "cuda"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SAM_MODEL",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "EfficientTAM"
|
|
|
|
def load_model(self, model_name, device):
|
|
model = load_model(
|
|
variant=model_name,
|
|
device=device,
|
|
mode="eval"
|
|
)
|
|
return (model,)
|
|
|
|
class EfficientTAMPredictor:
|
|
"""Predicts masks using EfficientTAM"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("SAM_MODEL",),
|
|
"image": ("IMAGE",), # ComfyUI image tensor (B,H,W,C)
|
|
"points": ("COORDS",), # List of [x,y] coordinates
|
|
"point_labels": ("LABELS",), # List of 1s (foreground) or 0s (background)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",) # Single channel mask (B,H,W)
|
|
FUNCTION = "predict"
|
|
CATEGORY = "EfficientTAM"
|
|
|
|
def predict(self, model, image, points, point_labels):
|
|
# Convert from BHWC to BCHW
|
|
image = image.permute(0,3,1,2)
|
|
|
|
# Initialize predictor
|
|
predictor = SAM2ImagePredictor(model)
|
|
|
|
# Set image
|
|
predictor.set_image(image)
|
|
|
|
# Convert points to torch tensor
|
|
points = torch.as_tensor(points, device=predictor.device)
|
|
point_labels = torch.as_tensor(point_labels, device=predictor.device)
|
|
|
|
# Get prediction
|
|
masks, iou_preds, _ = predictor.predict(
|
|
point_coords=points,
|
|
point_labels=point_labels,
|
|
multimask_output=True,
|
|
return_logits=True
|
|
)
|
|
|
|
# Return mask with highest IoU
|
|
best_idx = torch.argmax(iou_preds)
|
|
mask = masks[best_idx:best_idx+1]
|
|
|
|
return (mask,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"EfficientTAMLoader": EfficientTAMLoader,
|
|
"EfficientTAMPredictor": EfficientTAMPredictor
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"EfficientTAMLoader": "Load EfficientTAM Model",
|
|
"EfficientTAMPredictor": "EfficientTAM Predict"
|
|
} |