61 lines
2.0 KiB
Python
61 lines
2.0 KiB
Python
import torch
|
|
import numpy as np
|
|
import json
|
|
import io
|
|
import base64
|
|
from PIL import Image
|
|
|
|
class PointCollectorSAM2:
|
|
"""
|
|
Interactive Point Collector for SAM 2.
|
|
Outputs pixel coordinates in JSON format: [{"x": 1, "y": 2}]
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
# Hidden widgets populated by JS
|
|
"coordinates": ("STRING", {"multiline": False, "default": "[]"}),
|
|
"neg_coordinates": ("STRING", {"multiline": False, "default": "[]"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("pos_points_json", "neg_points_json")
|
|
FUNCTION = "collect"
|
|
CATEGORY = "AnotherUtils/sam2"
|
|
OUTPUT_NODE = True
|
|
|
|
def collect(self, image, coordinates, neg_coordinates):
|
|
# Coordinates from JS are already in pixel units relative to the image
|
|
# We just need to ensure they are valid JSON strings for our SAM2 nodes
|
|
|
|
pos_json = coordinates if coordinates and coordinates.strip() else "[]"
|
|
neg_json = neg_coordinates if neg_coordinates and neg_coordinates.strip() else "[]"
|
|
|
|
# Send image to the JS widget via a UI message
|
|
img_base64 = self.tensor_to_base64(image)
|
|
|
|
return {
|
|
"ui": {"bg_image": [img_base64]},
|
|
"result": (pos_json, neg_json)
|
|
}
|
|
|
|
def tensor_to_base64(self, tensor):
|
|
# Convert from [B, H, W, C] to PIL Image (first frame)
|
|
img_array = tensor[0].cpu().numpy()
|
|
img_array = (img_array * 255).astype(np.uint8)
|
|
pil_img = Image.fromarray(img_array)
|
|
|
|
buffered = io.BytesIO()
|
|
pil_img.save(buffered, format="JPEG", quality=75)
|
|
img_bytes = buffered.getvalue()
|
|
img_base64 = base64.b64encode(img_bytes).decode('utf-8')
|
|
|
|
return img_base64
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, image, coordinates, neg_coordinates):
|
|
return float("nan") # Always update UI when points change
|