From ff5dbd90258f4235e4ef3db93736367b2fd97095 Mon Sep 17 00:00:00 2001 From: TTPlanetPig <152850462+TTPlanetPig@users.noreply.github.com> Date: Fri, 13 Jun 2025 11:16:26 +0800 Subject: [PATCH] Add bounding box selection option --- README.md | 8 ++++++-- nodes.py | 41 +++++++++++++++++++++++++++++++++++------ 2 files changed, 41 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index ece78dc..b9ff84d 100644 --- a/README.md +++ b/README.md @@ -8,12 +8,16 @@ This repository provides a custom [ComfyUI](https://github.com/comfyanonymous/Co Downloads a chosen Qwen 2.5-VL model into `models/Qwen` and returns the loaded model and processor. You can also choose which device to load the model onto (e.g. `cuda:1` if you have multiple GPUs). ### `QwenVLDetection` -Runs a detection prompt on an input image using the loaded model. The node outputs the raw text response and the bounding box coordinates. +Runs a detection prompt on an input image using the loaded model. The node outputs the raw text response and the bounding box coordinates. You can specify which boxes to return using the **bbox_selection** parameter: + +- `all` – return all boxes (default) +- `merge` – merge all boxes into a single bounding box +- Comma-separated indices such as `0`, `1,2` or `0,2` – return only the selected boxes, sorted by detection confidence The bounding boxes are converted to absolute pixel coordinates so that they can be directly fed into the SAM2 nodes. ## Usage 1. Place this repository inside your `ComfyUI/custom_nodes` directory. 2. From the **Download and Load Qwen2.5-VL Model** node, select the model you want to use and, if necessary, choose the device (such as `cuda:1`) where it should be loaded. The snapshot download will resume automatically if a previous attempt was interrupted. -3. Connect the output model to **Qwen2.5-VL Object Detection**, provide an image and the object you want to locate (e.g. `cat`). The node will automatically build the detection prompt and return the bounding boxes. +3. Connect the output model to **Qwen2.5-VL Object Detection**, provide an image and the object you want to locate (e.g. `cat`). Optionally set **bbox_selection** to `merge` or specific indices like `0,2` if you only want certain boxes. The node will automatically build the detection prompt and return the selected bounding boxes. 4. Feed the resulting bounding boxes into the SAM2 workflow for further processing. diff --git a/nodes.py b/nodes.py index c5012b0..4287c3e 100644 --- a/nodes.py +++ b/nodes.py @@ -2,7 +2,7 @@ import os import ast import json from dataclasses import dataclass -from typing import List, Dict, Any +from typing import List, Dict, Any, Tuple import torch from PIL import Image @@ -22,7 +22,13 @@ def parse_json(json_output: str) -> str: return json_output -def parse_boxes(text: str, img_width: int, img_height: int, input_w: int, input_h: int) -> List[List[int]]: +def parse_boxes( + text: str, + img_width: int, + img_height: int, + input_w: int, + input_h: int, +) -> List[List[int]]: text = parse_json(text) try: data = ast.literal_eval(text) @@ -30,9 +36,10 @@ def parse_boxes(text: str, img_width: int, img_height: int, input_w: int, input_ end_idx = text.rfind('"}') + len('"}') truncated = text[:end_idx] + "]" data = ast.literal_eval(truncated) - boxes = [] + items: List[Tuple[float, List[int]]] = [] for item in data: box = item.get("bbox_2d") or item.get("bbox") or item + score = float(item.get("score", 0)) y1, x1, y2, x2 = box[1], box[0], box[3], box[2] abs_y1 = int(y1 / input_h * img_height) abs_x1 = int(x1 / input_w * img_width) @@ -42,8 +49,9 @@ def parse_boxes(text: str, img_width: int, img_height: int, input_w: int, input_ abs_x1, abs_x2 = abs_x2, abs_x1 if abs_y1 > abs_y2: abs_y1, abs_y2 = abs_y2, abs_y1 - boxes.append([abs_x1, abs_y1, abs_x2, abs_y2]) - return boxes + items.append((score, [abs_x1, abs_y1, abs_x2, abs_y2])) + items.sort(key=lambda x: x[0], reverse=True) + return [box for _, box in items] @dataclass @@ -130,6 +138,7 @@ class QwenVLDetection: "qwen_model": ("QWEN_MODEL",), "image": ("IMAGE",), "target": ("STRING", {"default": "object"}), + "bbox_selection": ("STRING", {"default": "all"}), }, } @@ -138,7 +147,7 @@ class QwenVLDetection: FUNCTION = "detect" CATEGORY = "Qwen2.5-VL" - def detect(self, qwen_model: QwenModel, image, target: str): + def detect(self, qwen_model: QwenModel, image, target: str, bbox_selection: str = "all"): model = qwen_model.model processor = qwen_model.processor device = next(model.parameters()).device @@ -168,6 +177,26 @@ class QwenVLDetection: input_h = inputs['image_grid_thw'][0][1] * 14 input_w = inputs['image_grid_thw'][0][2] * 14 boxes = parse_boxes(output_text, image.width, image.height, input_w, input_h) + + selection = bbox_selection.strip().lower() + if selection == "merge": + if boxes: + x1 = min(b[0] for b in boxes) + y1 = min(b[1] for b in boxes) + x2 = max(b[2] for b in boxes) + y2 = max(b[3] for b in boxes) + boxes = [[x1, y1, x2, y2]] + else: + boxes = [] + elif selection != "all": + idxs = [] + for part in selection.replace(",", " ").split(): + try: + idxs.append(int(part)) + except Exception: + continue + boxes = [boxes[i] for i in idxs if 0 <= i < len(boxes)] + return (output_text, boxes)