259 lines
8.7 KiB
Python
259 lines
8.7 KiB
Python
"""
|
|
SAM3 Server - API endpoints for interactive point selection
|
|
|
|
Provides endpoints for the SAM3 interactive editor:
|
|
- /sam3/prepare: Load model and image
|
|
- /sam3/detect: Run SAM3 with point prompts
|
|
"""
|
|
|
|
import io
|
|
import json
|
|
import os
|
|
import threading
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
import folder_paths
|
|
|
|
# Global state for SAM3 predictor
|
|
sam3_predictor = None
|
|
sam3_lock = threading.Condition()
|
|
|
|
|
|
def load_image_from_path(filename, file_type, subfolder=""):
|
|
"""Load image from ComfyUI storage"""
|
|
# Get the appropriate directory based on type
|
|
if file_type == "input":
|
|
input_dir = folder_paths.get_input_directory()
|
|
elif file_type == "temp":
|
|
input_dir = folder_paths.get_temp_directory()
|
|
elif file_type == "output":
|
|
input_dir = folder_paths.get_output_directory()
|
|
else:
|
|
input_dir = folder_paths.get_input_directory()
|
|
|
|
# Build full path
|
|
if subfolder:
|
|
full_path = os.path.join(input_dir, subfolder, filename)
|
|
else:
|
|
full_path = os.path.join(input_dir, filename)
|
|
|
|
print(f"[SAM3 Server] Loading image from: {full_path}")
|
|
|
|
# Load image
|
|
image = Image.open(full_path)
|
|
if image.mode != 'RGB':
|
|
image = image.convert('RGB')
|
|
|
|
return image
|
|
|
|
|
|
@PromptServer.instance.routes.post("/sam3/prepare")
|
|
async def sam3_prepare(request):
|
|
"""
|
|
Load SAM3 model and prepare image for interactive editing
|
|
|
|
Request JSON:
|
|
{
|
|
"sam_model_name": "auto",
|
|
"filename": "image.png",
|
|
"type": "temp",
|
|
"subfolder": ""
|
|
}
|
|
"""
|
|
global sam3_predictor
|
|
|
|
try:
|
|
data = await request.json()
|
|
filename = data.get("filename", "")
|
|
file_type = data.get("type", "temp")
|
|
subfolder = data.get("subfolder", "")
|
|
|
|
print(f"[SAM3 Server] Preparing SAM3 for image: {filename}")
|
|
|
|
# Load image
|
|
image = load_image_from_path(filename, file_type, subfolder)
|
|
|
|
# Load SAM3 model if not already loaded
|
|
# For now, we'll load the model on-demand when /sam3/detect is called
|
|
# This is because the model is managed by LoadSAM3Model node
|
|
|
|
with sam3_lock:
|
|
sam3_predictor = {
|
|
"image": image,
|
|
"filename": filename,
|
|
"type": file_type,
|
|
"subfolder": subfolder,
|
|
}
|
|
print(f"[SAM3 Server] Image prepared: {image.size}")
|
|
|
|
return web.Response(status=200)
|
|
|
|
except Exception as e:
|
|
print(f"[SAM3 Server] Error in /sam3/prepare: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return web.Response(status=500, text=str(e))
|
|
|
|
|
|
@PromptServer.instance.routes.post("/sam3/detect")
|
|
async def sam3_detect(request):
|
|
"""
|
|
Run SAM3 detection with point prompts
|
|
|
|
Request JSON:
|
|
{
|
|
"positive_points": [[x1, y1], [x2, y2], ...],
|
|
"negative_points": [[x3, y3], ...],
|
|
"threshold": 0.7
|
|
}
|
|
|
|
Returns: PNG image of the mask
|
|
"""
|
|
global sam3_predictor
|
|
|
|
try:
|
|
data = await request.json()
|
|
positive_points = data.get("positive_points", [])
|
|
negative_points = data.get("negative_points", [])
|
|
threshold = data.get("threshold", 0.7)
|
|
|
|
print(f"[SAM3 Server] Detecting with {len(positive_points)} positive and {len(negative_points)} negative points")
|
|
|
|
if sam3_predictor is None:
|
|
return web.Response(status=400, text="SAM3 not prepared. Call /sam3/prepare first.")
|
|
|
|
# Get the stored image
|
|
image = sam3_predictor.get("image")
|
|
if image is None:
|
|
return web.Response(status=400, text="No image loaded")
|
|
|
|
# Load SAM3 model on-demand
|
|
# We need to import here to avoid circular imports
|
|
from .sam3_lib.model_builder import build_sam3_image_model
|
|
from .sam3_lib.model.sam3_image_processor import SAM3ImageProcessor
|
|
|
|
# Build model (cached globally if possible)
|
|
if "model" not in sam3_predictor:
|
|
print(f"[SAM3 Server] Loading SAM3 model...")
|
|
model = build_sam3_image_model()
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model = model.to(device).eval()
|
|
|
|
processor = SAM3ImageProcessor(model, device=device)
|
|
|
|
sam3_predictor["model"] = model
|
|
sam3_predictor["processor"] = processor
|
|
sam3_predictor["device"] = device
|
|
print(f"[SAM3 Server] SAM3 model loaded on {device}")
|
|
|
|
processor = sam3_predictor["processor"]
|
|
device = sam3_predictor["device"]
|
|
|
|
# Set image and extract features
|
|
print(f"[SAM3 Server] Setting image...")
|
|
state = processor.set_image(image)
|
|
|
|
# Prepare points and labels
|
|
all_points = []
|
|
all_labels = []
|
|
|
|
# Add positive points (label=1)
|
|
for point in positive_points:
|
|
all_points.append(point)
|
|
all_labels.append(1)
|
|
|
|
# Add negative points (label=0)
|
|
for point in negative_points:
|
|
all_points.append(point)
|
|
all_labels.append(0)
|
|
|
|
print(f"[SAM3 Server] Total points: {len(all_points)}")
|
|
|
|
# Run detection with points if we have any
|
|
if len(all_points) > 0:
|
|
# Convert to tensors
|
|
points_tensor = torch.tensor(all_points, dtype=torch.float32).unsqueeze(0) # [1, N, 2]
|
|
labels_tensor = torch.tensor(all_labels, dtype=torch.int32).unsqueeze(0) # [1, N]
|
|
|
|
print(f"[SAM3 Server] Running point prompt detection...")
|
|
state = processor.add_point_prompt(points_tensor, labels_tensor, state)
|
|
|
|
# Get masks from state
|
|
masks = state.get("masks", None)
|
|
scores = state.get("scores", None)
|
|
|
|
if masks is None or len(masks) == 0:
|
|
print(f"[SAM3 Server] No masks generated")
|
|
# Return empty mask
|
|
mask_array = np.zeros((image.height, image.width), dtype=np.uint8)
|
|
else:
|
|
print(f"[SAM3 Server] Got {len(masks)} masks, scores: {scores}")
|
|
|
|
# Filter by threshold if scores available
|
|
if scores is not None:
|
|
valid_indices = scores > threshold
|
|
if valid_indices.sum() > 0:
|
|
masks = masks[valid_indices]
|
|
scores = scores[valid_indices]
|
|
print(f"[SAM3 Server] Filtered to {len(masks)} masks above threshold {threshold}")
|
|
else:
|
|
# Take best mask if none above threshold
|
|
best_idx = torch.argmax(scores)
|
|
masks = masks[best_idx:best_idx+1]
|
|
scores = scores[best_idx:best_idx+1]
|
|
print(f"[SAM3 Server] No masks above threshold, using best mask with score {scores[0]:.3f}")
|
|
|
|
# Combine multiple masks using bitwise OR
|
|
mask_array = np.zeros((image.height, image.width), dtype=np.uint8)
|
|
for mask in masks:
|
|
# Convert tensor to numpy
|
|
if isinstance(mask, torch.Tensor):
|
|
mask_np = mask.cpu().numpy()
|
|
else:
|
|
mask_np = mask
|
|
|
|
# Ensure correct shape
|
|
if mask_np.ndim == 3:
|
|
mask_np = mask_np[0] # Remove batch dimension
|
|
|
|
# Resize if needed
|
|
if mask_np.shape != (image.height, image.width):
|
|
from PIL import Image as PILImage
|
|
mask_pil = PILImage.fromarray((mask_np > 0).astype(np.uint8) * 255)
|
|
mask_pil = mask_pil.resize((image.width, image.height), PILImage.NEAREST)
|
|
mask_np = np.array(mask_pil) > 0
|
|
|
|
# Combine with OR
|
|
mask_array = np.logical_or(mask_array, mask_np > 0).astype(np.uint8) * 255
|
|
|
|
else:
|
|
print(f"[SAM3 Server] No points provided, returning empty mask")
|
|
mask_array = np.zeros((image.height, image.width), dtype=np.uint8)
|
|
|
|
# Convert mask to PIL Image
|
|
mask_image = Image.fromarray(mask_array, mode='L')
|
|
|
|
# Convert to PNG bytes
|
|
img_buffer = io.BytesIO()
|
|
mask_image.save(img_buffer, format='PNG')
|
|
img_buffer.seek(0)
|
|
|
|
print(f"[SAM3 Server] Returning mask PNG")
|
|
|
|
return web.Response(
|
|
body=img_buffer.getvalue(),
|
|
headers={'Content-Type': 'image/png'}
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"[SAM3 Server] Error in /sam3/detect: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return web.Response(status=500, text=str(e))
|
|
|
|
|
|
print("[SAM3 Server] Registered routes: /sam3/prepare, /sam3/detect")
|