681 lines
25 KiB
Python
681 lines
25 KiB
Python
import torch
|
|
import numpy as np
|
|
import os
|
|
import json
|
|
from .inference.depth_service import DepthAnythingService
|
|
from .inference.detection_service import DetectionService
|
|
from .inference.segmentation_service import SAM2Service, SAM2VideoService
|
|
|
|
def filter_models(extension_list, subfolder):
|
|
base_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "models")
|
|
search_path = os.path.join(base_path, subfolder)
|
|
another_path = os.path.join(base_path, "another_utils")
|
|
|
|
found = []
|
|
for path in [search_path, another_path]:
|
|
if os.path.exists(path):
|
|
for root, dirs, files in os.walk(path):
|
|
for file in files:
|
|
if any(file.endswith(ext) for ext in extension_list):
|
|
found.append(file)
|
|
return sorted(list(set(found)))
|
|
|
|
class AnotherLoadYOLO:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
yolo_sizes = ["n", "s", "m", "l", "x"]
|
|
yolo_tasks = ["", "-pose", "-seg"]
|
|
yolo_versions = ["yolov8", "yolo11"]
|
|
presets = [f"{v}{s}{t}.pt" for v in yolo_versions for t in yolo_tasks for s in yolo_sizes]
|
|
local = filter_models([".pt", ".pth", ".safetensors"], "ultralytics")
|
|
models = sorted(list(set(presets + local)))
|
|
return {
|
|
"required": {
|
|
"model_name": (models, {"default": "yolov8m.pt"}),
|
|
"device": (["auto", "cuda", "cpu"],),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("ANOTHER_MODEL",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "AnotherUtils/inference"
|
|
def load(self, model_name, device="auto"):
|
|
return (AnotherLoadInferenceModel().load_model("YOLO", model_name, device),)
|
|
|
|
class AnotherLoadSAM2:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
presets = ["sam2_hiera_tiny.pt", "sam2_hiera_small.pt", "sam2_hiera_base_plus.pt", "sam2_hiera_large.pt"]
|
|
local = filter_models([".pt", ".pth", ".safetensors"], "sam2")
|
|
models = sorted(list(set(presets + local)))
|
|
return {
|
|
"required": {
|
|
"model_name": (models, {"default": "sam2_hiera_small.pt"}),
|
|
"mode": (["single_image", "video"], {"default": "single_image"}),
|
|
"device": (["auto", "cuda", "cpu"],),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("ANOTHER_MODEL",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "AnotherUtils/inference"
|
|
def load(self, model_name, mode="single_image", device="auto"):
|
|
model_type = "SAM2" if mode == "single_image" else "SAM2_VIDEO"
|
|
return (AnotherLoadInferenceModel().load_model(model_type, model_name, device),)
|
|
|
|
class AnotherLoadDepth:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
models = ["v3-tiny", "v3-small", "v3-medium", "v3-large"]
|
|
return {
|
|
"required": {
|
|
"model_name": (models, {"default": "v3-small"}),
|
|
"device": (["auto", "cuda", "cpu"],),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("ANOTHER_MODEL",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "AnotherUtils/inference"
|
|
def load(self, model_name, device="auto"):
|
|
return (AnotherLoadInferenceModel().load_model("DepthAnything", model_name, device),)
|
|
|
|
class AnotherLoadInferenceModel:
|
|
# Keeps this internal for backward compatibility of logic
|
|
# but we remove it from NODE_CLASS_MAPPINGS later
|
|
def load_model(self, model_type, model_name, device="auto"):
|
|
# Determine paths
|
|
base_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "models")
|
|
|
|
# Search for the model file
|
|
model_path = model_name
|
|
|
|
if "." in model_name:
|
|
found = False
|
|
# Check standard subfolders
|
|
for sub in ["sam2", "ultralytics", "another_utils", "ultralytics/bbox", "ultralytics/segm"]:
|
|
check_path = os.path.join(base_path, sub, model_name)
|
|
if os.path.exists(check_path):
|
|
model_path = check_path
|
|
found = True
|
|
break
|
|
|
|
if not found:
|
|
# If not found, we set the target path but the service will handle the download
|
|
subfolder = "sam2" if model_type == "SAM2" else "ultralytics"
|
|
model_path = os.path.join(base_path, subfolder, model_name)
|
|
|
|
model_data = {
|
|
"type": model_type,
|
|
"name": model_name,
|
|
"path": model_path,
|
|
"device": device,
|
|
"service": None
|
|
}
|
|
|
|
# Initialize the service (will trigger download if missing)
|
|
if model_type == "YOLO":
|
|
model_data["service"] = DetectionService(model_path, device=device)
|
|
elif model_type == "DepthAnything":
|
|
model_data["service"] = DepthAnythingService(model_name, device=device)
|
|
elif model_type == "SAM2":
|
|
model_data["service"] = SAM2Service(None, model_path, device=device)
|
|
elif model_type == "SAM2_VIDEO":
|
|
model_data["service"] = SAM2VideoService(None, model_path, device=device)
|
|
|
|
return model_data
|
|
|
|
class AnotherYOLOInference:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"model": ("ANOTHER_MODEL",),
|
|
"threshold": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BBOX", "KEYPOINTS", "MASK", "STRING", "IMAGE")
|
|
RETURN_NAMES = ("bboxes", "keypoints", "mask", "labels", "debug_image")
|
|
FUNCTION = "detect"
|
|
CATEGORY = "AnotherUtils/inference"
|
|
|
|
def detect(self, images, model, threshold):
|
|
if model["type"] != "YOLO":
|
|
raise ValueError("Model must be of type YOLO for AnotherYOLOInference")
|
|
|
|
service: DetectionService = model["service"]
|
|
service.threshold = threshold
|
|
|
|
all_bboxes = []
|
|
all_kpts = []
|
|
all_masks = []
|
|
all_labels = []
|
|
debug_images = []
|
|
|
|
# Process batch
|
|
for i in range(images.shape[0]):
|
|
img_np = (images[i].cpu().numpy() * 255).astype(np.uint8)
|
|
detections = service.infer(img_np, conf=threshold)
|
|
|
|
img_bboxes = [d["bbox"] for d in detections]
|
|
img_kpts = [d.get("pose", []) for d in detections]
|
|
img_labels = [d["label"] for d in detections]
|
|
|
|
# Merge masks for this image
|
|
current_img_masks = [d["mask"] for d in detections if "mask" in d]
|
|
if current_img_masks:
|
|
merged = np.any(current_img_masks, axis=0)
|
|
else:
|
|
merged = np.zeros((images.shape[1], images.shape[2]), dtype=bool)
|
|
|
|
# Draw debug image with bboxes, labels, and pose
|
|
debug_img = img_np.copy()
|
|
h, w = debug_img.shape[:2]
|
|
|
|
# COCO skeleton connections for drawing lines between keypoints
|
|
SKELETON = [
|
|
(0,1),(0,2),(1,3),(2,4), # head
|
|
(5,6),(5,7),(7,9),(6,8),(8,10), # arms
|
|
(5,11),(6,12),(11,12), # torso
|
|
(11,13),(13,15),(12,14),(14,16) # legs
|
|
]
|
|
# Colors per keypoint region (RGB)
|
|
KPT_COLOR = [
|
|
[255,0,0], # 0 nose - red
|
|
[255,85,0], # 1 l_eye
|
|
[255,170,0], # 2 r_eye
|
|
[255,255,0], # 3 l_ear
|
|
[170,255,0], # 4 r_ear
|
|
[0,255,0], # 5 l_shoulder - green
|
|
[0,255,85], # 6 r_shoulder
|
|
[0,255,170], # 7 l_elbow
|
|
[0,255,255], # 8 r_elbow
|
|
[0,170,255], # 9 l_wrist
|
|
[0,85,255], # 10 r_wrist
|
|
[0,0,255], # 11 l_hip - blue
|
|
[85,0,255], # 12 r_hip
|
|
[170,0,255], # 13 l_knee
|
|
[255,0,255], # 14 r_knee
|
|
[255,0,170], # 15 l_ankle
|
|
[255,0,85], # 16 r_ankle
|
|
]
|
|
|
|
for det in detections:
|
|
x1, y1, x2, y2 = [int(c) for c in det["bbox"]]
|
|
conf = det.get("confidence", 0)
|
|
label = f'{det["label"]} {conf:.2f}'
|
|
|
|
# Draw bbox rectangle (green, 2px)
|
|
t = 2 # thickness
|
|
debug_img[max(0,y1):min(h,y1+t), max(0,x1):min(w,x2)] = [0, 255, 0]
|
|
debug_img[max(0,y2-t):min(h,y2), max(0,x1):min(w,x2)] = [0, 255, 0]
|
|
debug_img[max(0,y1):min(h,y2), max(0,x1):min(w,x1+t)] = [0, 255, 0]
|
|
debug_img[max(0,y1):min(h,y2), max(0,x2-t):min(w,x2)] = [0, 255, 0]
|
|
|
|
# Draw label background
|
|
label_h = 16
|
|
label_w = min(len(label) * 8 + 4, w - x1)
|
|
ly1 = max(0, y1 - label_h)
|
|
debug_img[ly1:y1, x1:min(w, x1 + label_w)] = [0, 255, 0]
|
|
|
|
# Draw pose keypoints and skeleton if available
|
|
kpts = det.get("pose", [])
|
|
if kpts and len(kpts) >= 17:
|
|
# Draw skeleton lines first (so dots are on top)
|
|
for (a, b) in SKELETON:
|
|
ax, ay = int(kpts[a][0]), int(kpts[a][1])
|
|
bx, by = int(kpts[b][0]), int(kpts[b][1])
|
|
if ax > 0 and ay > 0 and bx > 0 and by > 0:
|
|
# Simple line drawing using numpy (Bresenham-lite)
|
|
n_steps = max(abs(bx - ax), abs(by - ay), 1)
|
|
for s in range(n_steps + 1):
|
|
px = int(ax + (bx - ax) * s / n_steps)
|
|
py = int(ay + (by - ay) * s / n_steps)
|
|
if 0 <= py < h and 0 <= px < w:
|
|
debug_img[max(0,py-1):min(h,py+1), max(0,px-1):min(w,px+1)] = [200, 200, 200]
|
|
|
|
# Draw keypoint dots
|
|
r = 3 # radius
|
|
for ki, (kx, ky) in enumerate(kpts[:17]):
|
|
kx, ky = int(kx), int(ky)
|
|
if kx > 0 and ky > 0 and 0 <= ky < h and 0 <= kx < w:
|
|
color = KPT_COLOR[ki] if ki < len(KPT_COLOR) else [255, 255, 255]
|
|
debug_img[max(0,ky-r):min(h,ky+r), max(0,kx-r):min(w,kx+r)] = color
|
|
|
|
all_bboxes.append(img_bboxes)
|
|
all_kpts.append(img_kpts)
|
|
all_masks.append(torch.from_numpy(merged.astype(np.float32)))
|
|
all_labels.append(img_labels)
|
|
debug_images.append(torch.from_numpy(debug_img.astype(np.float32) / 255.0))
|
|
|
|
return (all_bboxes, all_kpts, torch.stack(all_masks), json.dumps(all_labels), torch.stack(debug_images))
|
|
|
|
class AnotherSAM2Inference:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"model": ("ANOTHER_MODEL",),
|
|
},
|
|
"optional": {
|
|
"points_pos": ("STRING", {"default": "[]"}),
|
|
"points_neg": ("STRING", {"default": "[]"}),
|
|
"bboxes": ("BBOX",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "segment"
|
|
CATEGORY = "AnotherUtils/inference"
|
|
|
|
def segment(self, image, model, points_pos="[]", points_neg="[]", bboxes=None):
|
|
if model["type"] != "SAM2":
|
|
raise ValueError("Model must be of type SAM2 for AnotherSAM2Inference")
|
|
|
|
service: SAM2Service = model["service"]
|
|
|
|
pos = json.loads(points_pos) if isinstance(points_pos, str) else points_pos
|
|
neg = json.loads(points_neg) if isinstance(points_neg, str) else points_neg
|
|
|
|
# Convert JSON structure to simple list of coords
|
|
p_coords = []
|
|
p_labels = []
|
|
|
|
for p in pos:
|
|
p_coords.append([p["x"], p["y"]])
|
|
p_labels.append(1)
|
|
for p in neg:
|
|
p_coords.append([p["x"], p["y"]])
|
|
p_labels.append(0)
|
|
|
|
all_masks = []
|
|
for i in range(image.shape[0]):
|
|
img_np = (image[i].cpu().numpy() * 255).astype(np.uint8)
|
|
service.set_image(img_np)
|
|
|
|
# If bboxes is provided, use them. If it's a batch of bboxes, take the correct index.
|
|
current_bboxes = None
|
|
if bboxes is not None:
|
|
if isinstance(bboxes[0], list) and len(bboxes) == image.shape[0]:
|
|
current_bboxes = bboxes[i]
|
|
else:
|
|
current_bboxes = bboxes
|
|
|
|
mask = service.predict(
|
|
points=p_coords if p_coords else None,
|
|
labels=p_labels if p_labels else None,
|
|
bboxes=current_bboxes if current_bboxes else None
|
|
)
|
|
# mask is [N, H, W] bool. We merge all N detections into one mask.
|
|
if mask.shape[0] > 0:
|
|
merged_mask = np.any(mask, axis=0)
|
|
else:
|
|
merged_mask = np.zeros((image.shape[1], image.shape[2]), dtype=bool)
|
|
|
|
all_masks.append(torch.from_numpy(merged_mask.astype(np.float32)))
|
|
|
|
return (torch.stack(all_masks),)
|
|
|
|
|
|
class AnotherSAM2VideoAddPoints:
|
|
"""Add point prompts to a specific frame for video segmentation.
|
|
Copied from comfyui-segment-anything-2 Sam2VideoSegmentationAddPoints."""
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s):
|
|
return ""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("ANOTHER_MODEL",),
|
|
"coordinates_positive": ("STRING", {"forceInput": True}),
|
|
"frame_index": ("INT", {"default": 0}),
|
|
"object_index": ("INT", {"default": 0}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
"coordinates_negative": ("STRING", {"forceInput": True}),
|
|
"prev_state": ("ANOTHER_SAM2_STATE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("ANOTHER_MODEL", "ANOTHER_SAM2_STATE")
|
|
RETURN_NAMES = ("model", "state")
|
|
FUNCTION = "add_points"
|
|
CATEGORY = "AnotherUtils/inference"
|
|
|
|
def add_points(self, model, coordinates_positive, frame_index, object_index,
|
|
image=None, coordinates_negative=None, prev_state=None):
|
|
if model["type"] != "SAM2_VIDEO":
|
|
raise ValueError("Model must be loaded with mode='video' for video segmentation")
|
|
|
|
service: SAM2VideoService = model["service"]
|
|
|
|
# Parse coordinates
|
|
try:
|
|
pos = json.loads(coordinates_positive.replace("'", '"'))
|
|
pos = [(c['x'], c['y']) for c in pos]
|
|
except:
|
|
pos = []
|
|
|
|
neg = None
|
|
if coordinates_negative:
|
|
try:
|
|
neg_parsed = json.loads(coordinates_negative.replace("'", '"'))
|
|
neg = [(c['x'], c['y']) for c in neg_parsed]
|
|
except:
|
|
neg = None
|
|
|
|
# Initialize state from images if no previous state
|
|
if prev_state is None:
|
|
if image is None:
|
|
raise ValueError("Either 'image' or 'prev_state' is required")
|
|
B, H, W, C = image.shape
|
|
model_input_size = service.model.image_size
|
|
# Resize to model input size: [B,H,W,C] -> [B,C,H,W]
|
|
from comfy.utils import common_upscale
|
|
resized = common_upscale(
|
|
image.movedim(-1, 1),
|
|
model_input_size, model_input_size,
|
|
"bilinear", "disabled"
|
|
).movedim(1, -1)
|
|
# init_state expects [B, C, H, W]
|
|
service.init_state(resized.permute(0, 3, 1, 2).contiguous(), H, W)
|
|
num_frames = B
|
|
else:
|
|
service.inference_state = prev_state["inference_state"]
|
|
num_frames = prev_state["num_frames"]
|
|
|
|
# Add points
|
|
if pos:
|
|
service.add_points(frame_index, object_index, pos, neg)
|
|
|
|
state = {
|
|
"inference_state": service.inference_state,
|
|
"num_frames": num_frames,
|
|
}
|
|
return (model, state)
|
|
|
|
|
|
class AnotherSAM2VideoPropagate:
|
|
"""Propagate segmentation across all video frames.
|
|
Copied from comfyui-segment-anything-2 Sam2VideoSegmentation."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("ANOTHER_MODEL",),
|
|
"state": ("ANOTHER_SAM2_STATE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "propagate"
|
|
CATEGORY = "AnotherUtils/inference"
|
|
|
|
def propagate(self, model, state):
|
|
if model["type"] != "SAM2_VIDEO":
|
|
raise ValueError("Model must be loaded with mode='video'")
|
|
|
|
service: SAM2VideoService = model["service"]
|
|
service.inference_state = state["inference_state"]
|
|
|
|
mask_tensor = service.propagate()
|
|
return (mask_tensor,)
|
|
|
|
|
|
class AnotherDepthInference:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"model": ("ANOTHER_MODEL",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "infer"
|
|
CATEGORY = "AnotherUtils/inference"
|
|
|
|
def infer(self, images, model):
|
|
if model["type"] != "DepthAnything":
|
|
raise ValueError("Model must be of type DepthAnything for AnotherDepthInference")
|
|
|
|
service: DepthAnythingService = model["service"]
|
|
|
|
all_depths = []
|
|
for i in range(images.shape[0]):
|
|
img_np = (images[i].cpu().numpy() * 255).astype(np.uint8)
|
|
depth = service.infer(img_np)
|
|
all_depths.append(torch.from_numpy(depth))
|
|
|
|
depth_batch = torch.stack(all_depths)
|
|
# depth_batch is [B, H, W]. Return as image [B, H, W, 1] and mask [B, H, W]
|
|
return (depth_batch.unsqueeze(-1).repeat(1, 1, 1, 3), depth_batch)
|
|
|
|
class AnotherBBoxToPoints:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"bboxes": ("BBOX",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("points_json",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "AnotherUtils/utils"
|
|
|
|
def convert(self, bboxes):
|
|
# bboxes: [[x1, y1, x2, y2], ...]
|
|
points = []
|
|
# If it's a batch of bboxes, we might need to handle it.
|
|
# For simplicity, if it's nested [[...]], we take the first image's boxes.
|
|
target_boxes = bboxes
|
|
if bboxes and isinstance(bboxes[0], list) and len(bboxes[0]) > 0 and isinstance(bboxes[0][0], (int, float, list)):
|
|
target_boxes = bboxes[0]
|
|
|
|
for box in target_boxes:
|
|
x1, y1, x2, y2 = box
|
|
cx = x1 + (x2 - x1) / 2
|
|
cy = y1 + (y2 - y1) / 2
|
|
points.append({"x": cx, "y": cy})
|
|
|
|
return (json.dumps(points),)
|
|
|
|
class AnotherPoseToPoints:
|
|
BODY_PARTS = [
|
|
"all", "face", "upper_body",
|
|
"nose", "eyes", "ears", "shoulders", "elbows", "wrists", "hips", "knees", "ankles",
|
|
"left_eye", "right_eye", "left_ear", "right_ear",
|
|
"left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
|
|
"left_wrist", "right_wrist", "left_hip", "right_hip",
|
|
"left_knee", "right_knee", "left_ankle", "right_ankle",
|
|
]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"keypoints": ("KEYPOINTS",),
|
|
"body_part": (cls.BODY_PARTS, {"default": "all"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("points_json",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "AnotherUtils/utils"
|
|
|
|
def convert(self, keypoints, body_part):
|
|
# 0: nose, 1: l_eye, 2: r_eye, 5: l_sh, 6: r_sh...
|
|
COCO_MAP = {
|
|
"nose": 0, "left_eye": 1, "right_eye": 2, "left_ear": 3, "right_ear": 4,
|
|
"left_shoulder": 5, "right_shoulder": 6, "left_elbow": 7, "right_elbow": 8,
|
|
"left_wrist": 9, "right_wrist": 10, "left_hip": 11, "right_hip": 12,
|
|
"left_knee": 13, "right_knee": 14, "left_ankle": 15, "right_ankle": 16
|
|
}
|
|
# Group aliases for convenience
|
|
ALIASES = {
|
|
"eyes": ["left_eye", "right_eye"],
|
|
"ears": ["left_ear", "right_ear"],
|
|
"shoulders": ["left_shoulder", "right_shoulder"],
|
|
"elbows": ["left_elbow", "right_elbow"],
|
|
"wrists": ["left_wrist", "right_wrist"],
|
|
"hips": ["left_hip", "right_hip"],
|
|
"knees": ["left_knee", "right_knee"],
|
|
"ankles": ["left_ankle", "right_ankle"],
|
|
"face": ["nose", "left_eye", "right_eye"],
|
|
"upper_body": ["nose", "left_eye", "right_eye", "left_shoulder", "right_shoulder"],
|
|
"all": list(COCO_MAP.keys()),
|
|
}
|
|
|
|
points = []
|
|
|
|
# Handle batching: keypoints is List[List[kpts_per_detection]]
|
|
target_kpts = keypoints
|
|
if keypoints and isinstance(keypoints[0], list):
|
|
target_kpts = keypoints[0] # Take first image's detections
|
|
|
|
# Resolve body_part to keypoint indices
|
|
if body_part in ALIASES:
|
|
selected_indices = [COCO_MAP[name] for name in ALIASES[body_part]]
|
|
elif body_part in COCO_MAP:
|
|
selected_indices = [COCO_MAP[body_part]]
|
|
else:
|
|
selected_indices = list(COCO_MAP.values()) # fallback to all
|
|
|
|
for det_kpts in target_kpts:
|
|
# det_kpts is [[x,y], [x,y], ...] with 17 entries
|
|
for idx in selected_indices:
|
|
if idx < len(det_kpts):
|
|
p = det_kpts[idx]
|
|
# Skip keypoints with 0,0 (not detected)
|
|
if float(p[0]) > 0 and float(p[1]) > 0:
|
|
points.append({"x": float(p[0]), "y": float(p[1])})
|
|
|
|
return (json.dumps(points),)
|
|
|
|
class AnotherImageToMask:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"channel": (["red", "green", "blue", "alpha", "luminance"], {"default": "luminance"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "AnotherUtils/utils"
|
|
|
|
def convert(self, image, channel):
|
|
if channel == "luminance":
|
|
mask = image.mean(dim=-1)
|
|
elif channel == "red":
|
|
mask = image[:, :, :, 0]
|
|
elif channel == "green":
|
|
mask = image[:, :, :, 1]
|
|
elif channel == "blue":
|
|
mask = image[:, :, :, 2]
|
|
elif channel == "alpha":
|
|
if image.shape[-1] == 4:
|
|
mask = image[:, :, :, 3]
|
|
else:
|
|
mask = torch.ones((image.shape[0], image.shape[1], image.shape[2]), device=image.device)
|
|
return (mask,)
|
|
|
|
class AnotherMaskToImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "AnotherUtils/utils"
|
|
|
|
def convert(self, mask):
|
|
# mask is [B, H, W]
|
|
if mask.ndim == 2:
|
|
mask = mask.unsqueeze(0)
|
|
return (mask.unsqueeze(-1).repeat(1, 1, 1, 3),)
|
|
|
|
class AnotherMaskMath:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask_a": ("MASK",),
|
|
"operation": (["multiply", "add", "subtract", "and", "or", "xor"], {"default": "multiply"}),
|
|
},
|
|
"optional": {
|
|
"mask_b": ("MASK",),
|
|
"value": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "apply"
|
|
CATEGORY = "AnotherUtils/utils"
|
|
|
|
def apply(self, mask_a, operation, mask_b=None, value=1.0):
|
|
# Ensure mask_a is tensor
|
|
res = mask_a.clone()
|
|
|
|
# If mask_b is not provided, use the scalar value
|
|
other = mask_b if mask_b is not None else value
|
|
|
|
if operation == "multiply":
|
|
res = res * other
|
|
elif operation == "add":
|
|
res = res + other
|
|
elif operation == "subtract":
|
|
res = res - other
|
|
elif operation == "and":
|
|
res = torch.min(res, torch.tensor(other))
|
|
elif operation == "or":
|
|
res = torch.max(res, torch.tensor(other))
|
|
elif operation == "xor":
|
|
res = torch.abs(res - other)
|
|
|
|
return (torch.clamp(res, 0.0, 1.0),)
|
|
|
|
class AnotherMaskBlur:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",),
|
|
"blur_radius": ("INT", {"default": 5, "min": 0, "max": 100, "step": 1}),
|
|
"sigma": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "blur"
|
|
CATEGORY = "AnotherUtils/utils"
|
|
|
|
def blur(self, mask, blur_radius, sigma):
|
|
if blur_radius == 0:
|
|
return (mask,)
|
|
|
|
import scipy.ndimage
|
|
# Convert to numpy for scipy
|
|
mask_np = mask.cpu().numpy()
|
|
blurred_np = np.zeros_like(mask_np)
|
|
|
|
for i in range(mask_np.shape[0]):
|
|
blurred_np[i] = scipy.ndimage.gaussian_filter(mask_np[i], sigma=sigma, radius=blur_radius)
|
|
|
|
return (torch.from_numpy(blurred_np).to(mask.device),)
|