493 lines
16 KiB
Python
493 lines
16 KiB
Python
"""
|
|
ComfyUI nodes for DINO-X API integration
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import supervision as sv
|
|
import torch
|
|
from dds_cloudapi_sdk import Client, Config
|
|
from dds_cloudapi_sdk.tasks.detection import DetectionTask
|
|
from dds_cloudapi_sdk.tasks.dinox import DinoxTask
|
|
from dds_cloudapi_sdk.tasks.types import DetectionTarget
|
|
|
|
|
|
class PromptType(str, Enum):
|
|
TEXT = "text"
|
|
UNIVERSAL = "universal"
|
|
|
|
|
|
class DinoxBaseNode:
|
|
"""Base class for DINO-X API nodes with common functionality"""
|
|
|
|
def __init__(self):
|
|
self.output_dir = Path("./outputs/dinox")
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _prepare_image(self, image: torch.Tensor) -> np.ndarray:
|
|
"""Convert input tensor to OpenCV format"""
|
|
try:
|
|
if not isinstance(image, torch.Tensor):
|
|
raise ValueError("Expected torch.Tensor input")
|
|
|
|
# Ensure BHWC format
|
|
if len(image.shape) != 4:
|
|
raise ValueError(
|
|
f"Expected 4D tensor [B,H,W,C], got shape {image.shape}"
|
|
)
|
|
|
|
# Convert to numpy and scale to [0,255]
|
|
image_np = (image[0].cpu().numpy() * 255).astype(
|
|
np.uint8
|
|
) # Take first batch
|
|
|
|
# Convert to BGR for OpenCV
|
|
return image_np[..., ::-1].copy() # RGB to BGR
|
|
except Exception as e:
|
|
raise ValueError(f"Error processing input image: {str(e)}")
|
|
|
|
def _save_temp_image(self, image: np.ndarray) -> str:
|
|
"""Save image to temporary file and return path"""
|
|
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
|
|
temp_path = temp_file.name
|
|
try:
|
|
success = cv2.imwrite(temp_path, image)
|
|
if not success:
|
|
raise RuntimeError("Failed to save temporary image")
|
|
temp_file.flush()
|
|
os.fsync(temp_file.fileno())
|
|
return temp_path
|
|
except Exception as e:
|
|
try:
|
|
Path(temp_path).unlink(missing_ok=True)
|
|
except:
|
|
pass
|
|
raise RuntimeError(f"Failed to save temporary image: {str(e)}")
|
|
|
|
|
|
class DinoxDetectorNode(DinoxBaseNode):
|
|
"""Node for DINO-X object detection"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",), # [B,H,W,C] torch.Tensor
|
|
"text_prompt": (
|
|
"STRING",
|
|
{
|
|
"default": "wheel . eye . helmet . mouse . mouth . vehicle . steering wheel . ear . nose",
|
|
"multiline": True,
|
|
},
|
|
),
|
|
"api_token": ("STRING", {"default": ""}),
|
|
"bbox_threshold": (
|
|
"FLOAT",
|
|
{"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.05},
|
|
),
|
|
"iou_threshold": (
|
|
"FLOAT",
|
|
{"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.05},
|
|
),
|
|
"prompt_type": (["text", "universal"],),
|
|
"return_pose": ("BOOLEAN", {"default": False}),
|
|
"return_hand": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING")
|
|
RETURN_NAMES = ("box_annotated", "binary_mask", "pose_data", "hand_data")
|
|
FUNCTION = "detect_and_annotate"
|
|
CATEGORY = "detection"
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(
|
|
cls, api_token, text_prompt, bbox_threshold, iou_threshold, prompt_type
|
|
):
|
|
if not api_token:
|
|
return "API token cannot be empty"
|
|
|
|
if prompt_type == "text" and not text_prompt.strip():
|
|
return "Text prompt cannot be empty when prompt type is 'text'"
|
|
|
|
if bbox_threshold < 0.0 or bbox_threshold > 1.0:
|
|
return "bbox_threshold must be between 0.0 and 1.0"
|
|
|
|
if iou_threshold < 0.0 or iou_threshold > 1.0:
|
|
return "iou_threshold must be between 0.0 and 1.0"
|
|
|
|
return True
|
|
|
|
def detect_and_annotate(
|
|
self,
|
|
image: torch.Tensor,
|
|
text_prompt: str,
|
|
api_token: str,
|
|
bbox_threshold: float,
|
|
iou_threshold: float,
|
|
prompt_type: str = "text",
|
|
return_pose: bool = False,
|
|
return_hand: bool = False,
|
|
) -> Tuple[torch.Tensor, torch.Tensor, str, str]:
|
|
"""
|
|
Process an image with DINO-X API to detect and annotate objects
|
|
|
|
Args:
|
|
image: Input image tensor [B,H,W,C]
|
|
text_prompt: Text description of objects to detect
|
|
api_token: DINO-X API token
|
|
bbox_threshold: Detection confidence threshold
|
|
iou_threshold: IoU threshold for NMS
|
|
prompt_type: Type of prompt ("text" or "universal")
|
|
return_pose: Whether to return pose keypoints
|
|
return_hand: Whether to return hand keypoints
|
|
|
|
Returns:
|
|
tuple: (box_annotated, binary_mask, pose_data, hand_data)
|
|
"""
|
|
# Convert input tensor to OpenCV format
|
|
image_bgr = self._prepare_image(image)
|
|
|
|
# Save to temp file
|
|
temp_path = self._save_temp_image(image_bgr)
|
|
|
|
try:
|
|
# Initialize API client
|
|
config = Config(api_token)
|
|
client = Client(config)
|
|
|
|
# Prepare targets
|
|
targets = [DetectionTarget.BBox, DetectionTarget.Mask]
|
|
if return_pose:
|
|
targets.append(DetectionTarget.PoseKeypoints)
|
|
if return_hand:
|
|
targets.append(DetectionTarget.HandKeypoints)
|
|
|
|
# Prepare prompt
|
|
if prompt_type == "text":
|
|
prompt = {"type": "text", "text": text_prompt}
|
|
else:
|
|
prompt = {"type": "universal", "universal": 1}
|
|
|
|
# Upload image and run detection
|
|
image_url = client.upload_file(temp_path)
|
|
task = DinoxTask(
|
|
model="DINO-X-1.0", # Set model version
|
|
image_url=image_url,
|
|
prompt=prompt,
|
|
bbox_threshold=bbox_threshold,
|
|
iou_threshold=iou_threshold,
|
|
targets=targets,
|
|
)
|
|
client.run_task(task)
|
|
|
|
# Handle case where no detections are found
|
|
if not task.result or not hasattr(task.result, "objects"):
|
|
predictions = []
|
|
else:
|
|
predictions = task.result.objects
|
|
|
|
# Initialize empty outputs
|
|
h, w = image_bgr.shape[:2]
|
|
empty_mask = torch.zeros((1, h, w), dtype=torch.float32)
|
|
rgb_image = (
|
|
torch.from_numpy(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)).float()
|
|
/ 255.0
|
|
)
|
|
rgb_image = rgb_image.unsqueeze(0) # Add batch dimension
|
|
|
|
if not predictions:
|
|
return (
|
|
rgb_image,
|
|
empty_mask,
|
|
"[]", # Empty JSON array for pose data
|
|
"[]", # Empty JSON array for hand data
|
|
)
|
|
|
|
# Process predictions
|
|
boxes = []
|
|
masks = []
|
|
pose_points = []
|
|
hand_points = []
|
|
confidences = []
|
|
class_names = []
|
|
|
|
for obj in predictions:
|
|
if hasattr(obj, "bbox"):
|
|
boxes.append(obj.bbox)
|
|
confidences.append(obj.score)
|
|
class_names.append(obj.category)
|
|
|
|
if hasattr(obj, "mask") and obj.mask:
|
|
if hasattr(obj.mask, "counts") and hasattr(obj.mask, "size"):
|
|
masks.append(
|
|
DetectionTask.rle2mask(
|
|
DetectionTask.string2rle(obj.mask.counts), obj.mask.size
|
|
)
|
|
)
|
|
|
|
if return_pose and hasattr(obj, "pose"):
|
|
pose_points.append(obj.pose)
|
|
|
|
if return_hand and hasattr(obj, "hand"):
|
|
hand_points.append(obj.hand)
|
|
|
|
if not boxes:
|
|
return (
|
|
rgb_image,
|
|
empty_mask,
|
|
"[]", # Empty JSON array for pose data
|
|
"[]", # Empty JSON array for hand data
|
|
)
|
|
|
|
# Create annotations
|
|
boxes = np.array(boxes)
|
|
masks = (
|
|
np.array(masks) if masks else np.zeros((len(boxes), h, w), dtype=bool)
|
|
)
|
|
labels = [
|
|
f"{name} {conf:.2f}" for name, conf in zip(class_names, confidences)
|
|
]
|
|
|
|
# Create detections object
|
|
detections = sv.Detections(
|
|
xyxy=boxes,
|
|
mask=masks.astype(bool),
|
|
class_id=np.arange(len(boxes)),
|
|
)
|
|
|
|
# Create box annotations
|
|
box_annotated = image_bgr.copy()
|
|
box_annotator = sv.BoxAnnotator()
|
|
box_annotated = box_annotator.annotate(
|
|
scene=box_annotated, detections=detections
|
|
)
|
|
label_annotator = sv.LabelAnnotator()
|
|
box_annotated = label_annotator.annotate(
|
|
scene=box_annotated, detections=detections, labels=labels
|
|
)
|
|
|
|
# Create binary mask
|
|
binary_mask = np.zeros((h, w), dtype=np.uint8)
|
|
for mask in masks:
|
|
binary_mask |= mask.astype(np.uint8)
|
|
|
|
# Convert outputs to tensors
|
|
box_annotated = cv2.cvtColor(box_annotated, cv2.COLOR_BGR2RGB)
|
|
box_tensor = torch.from_numpy(box_annotated).float() / 255.0
|
|
box_tensor = box_tensor.unsqueeze(0) # Add batch dimension
|
|
|
|
mask_tensor = torch.from_numpy(binary_mask).float()
|
|
mask_tensor = mask_tensor.unsqueeze(0) # Add batch dimension
|
|
|
|
# Convert keypoints to JSON strings
|
|
pose_json = json.dumps(pose_points) if pose_points else "[]"
|
|
hand_json = json.dumps(hand_points) if hand_points else "[]"
|
|
|
|
return (
|
|
box_tensor,
|
|
mask_tensor,
|
|
pose_json,
|
|
hand_json,
|
|
)
|
|
|
|
finally:
|
|
# Clean up temp file
|
|
try:
|
|
Path(temp_path).unlink(missing_ok=True)
|
|
except Exception as e:
|
|
print(f"Warning: Failed to clean up temp file: {str(e)}")
|
|
|
|
|
|
class DinoxRegionVLNode(DinoxBaseNode):
|
|
"""Node for DINO-X region-based vision-language tasks"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",), # [B,H,W,C] torch.Tensor
|
|
"regions": (
|
|
"STRING",
|
|
{"multiline": True},
|
|
), # JSON array of [x1,y1,x2,y2] coordinates
|
|
"api_token": ("STRING", {"default": ""}),
|
|
"prompt_type": (["text", "universal"],),
|
|
"text_prompt": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
},
|
|
),
|
|
"return_caption": ("BOOLEAN", {"default": True}),
|
|
"return_roc": ("BOOLEAN", {"default": False}),
|
|
"return_ocr": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("captions", "roc_results", "ocr_results")
|
|
FUNCTION = "analyze_regions"
|
|
CATEGORY = "detection"
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, api_token, regions, prompt_type, text_prompt):
|
|
if not api_token:
|
|
return "API token cannot be empty"
|
|
|
|
if not regions.strip():
|
|
return "Regions cannot be empty"
|
|
|
|
if prompt_type == "text" and not text_prompt.strip():
|
|
return "Text prompt cannot be empty when prompt type is 'text'"
|
|
|
|
try:
|
|
import json
|
|
|
|
regions_list = json.loads(regions)
|
|
if not isinstance(regions_list, list):
|
|
return "Regions must be a JSON array"
|
|
if len(regions_list) > 1000:
|
|
return "Maximum 1000 regions allowed"
|
|
for region in regions_list:
|
|
if not isinstance(region, list) or len(region) != 4:
|
|
return "Each region must be an array of 4 coordinates [x1,y1,x2,y2]"
|
|
except Exception as e:
|
|
return f"Invalid regions format: {str(e)}"
|
|
|
|
return True
|
|
|
|
def analyze_regions(
|
|
self,
|
|
image: torch.Tensor,
|
|
regions: str,
|
|
api_token: str,
|
|
prompt_type: str = "text",
|
|
text_prompt: str = "",
|
|
return_caption: bool = True,
|
|
return_roc: bool = False,
|
|
return_ocr: bool = False,
|
|
) -> Tuple[str, str, str]:
|
|
"""
|
|
Analyze image regions using DINO-X API
|
|
|
|
Args:
|
|
image: Input image tensor [B,H,W,C]
|
|
regions: JSON array of region coordinates [[x1,y1,x2,y2],...]
|
|
api_token: DINO-X API token
|
|
prompt_type: Type of prompt ("text" or "universal")
|
|
text_prompt: Text prompt for detection (required for "text" type)
|
|
return_caption: Whether to return region captions
|
|
return_roc: Whether to return ROC results
|
|
return_ocr: Whether to return OCR results
|
|
|
|
Returns:
|
|
tuple: (captions, roc_results, ocr_results) as JSON strings
|
|
"""
|
|
# Convert input tensor to OpenCV format
|
|
image_bgr = self._prepare_image(image)
|
|
|
|
# Save to temp file
|
|
temp_path = self._save_temp_image(image_bgr)
|
|
|
|
try:
|
|
# Initialize API client
|
|
config = Config(api_token)
|
|
client = Client(config)
|
|
|
|
# Parse regions
|
|
import json
|
|
|
|
regions_list = json.loads(regions)
|
|
|
|
# Prepare targets
|
|
targets = []
|
|
if return_caption:
|
|
targets.append("caption")
|
|
if return_roc:
|
|
targets.append("roc")
|
|
if return_ocr:
|
|
targets.append("ocr")
|
|
|
|
# Prepare prompt
|
|
if prompt_type == "text":
|
|
prompt = {"type": "text", "text": text_prompt}
|
|
else:
|
|
prompt = {"type": "universal", "universal": 1}
|
|
|
|
# Upload image and run analysis
|
|
image_url = client.upload_file(temp_path)
|
|
task = DinoxTask(
|
|
model="DINO-X-1.0", # Set model version
|
|
image_url=image_url,
|
|
prompt=prompt,
|
|
regions=regions_list,
|
|
targets=targets,
|
|
)
|
|
client.run_task(task)
|
|
|
|
# Process results
|
|
if not task.result or not hasattr(task.result, "objects"):
|
|
return "{}", "{}", "{}"
|
|
|
|
# Extract results by type
|
|
captions = {}
|
|
roc_results = {}
|
|
ocr_results = {}
|
|
|
|
for idx, obj in enumerate(task.result.objects):
|
|
region_key = f"region_{idx}"
|
|
region = obj.get("region", [])
|
|
|
|
if return_caption:
|
|
captions[region_key] = {
|
|
"region": region,
|
|
"caption": obj.get("caption", ""),
|
|
}
|
|
|
|
if return_roc:
|
|
roc_results[region_key] = {
|
|
"region": region,
|
|
"roc": obj.get("roc", ""),
|
|
}
|
|
|
|
if return_ocr:
|
|
ocr_results[region_key] = {
|
|
"region": region,
|
|
"ocr": obj.get("ocr", ""),
|
|
}
|
|
|
|
return (
|
|
json.dumps(captions, indent=2),
|
|
json.dumps(roc_results, indent=2),
|
|
json.dumps(ocr_results, indent=2),
|
|
)
|
|
|
|
finally:
|
|
# Clean up temp file
|
|
try:
|
|
Path(temp_path).unlink(missing_ok=True)
|
|
except Exception as e:
|
|
print(f"Warning: Failed to clean up temp file: {str(e)}")
|
|
|
|
|
|
# Register nodes
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DinoxDetector": DinoxDetectorNode,
|
|
"DinoxRegionVL": DinoxRegionVLNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DinoxDetector": "DINO-X Object Detector",
|
|
"DinoxRegionVL": "DINO-X Region Analysis",
|
|
}
|