374 lines
14 KiB
Python
374 lines
14 KiB
Python
|
|
from .imagefunc import *
|
|
|
|
select_list = ["all", "first", "by_index"]
|
|
sort_method_list = ["left_to_right", "top_to_bottom", "big_to_small"]
|
|
|
|
def sort_bboxes(bboxes:list, method:str) -> list:
|
|
sorted_bboxes = []
|
|
if method == "left_to_right":
|
|
sorted_bboxes = sorted(bboxes, key=lambda bbox: bbox[0])
|
|
elif method == "top_to_bottom":
|
|
sorted_bboxes = sorted(bboxes, key=lambda bbox: bbox[1])
|
|
else:# bit_to_small
|
|
sorted_bboxes = sorted(bboxes, key=lambda bbox: (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), reverse=True)
|
|
return sorted_bboxes
|
|
|
|
def select_bboxes(bboxes:list, bbox_select:str, select_index:str) -> list:
|
|
indexs = extract_numbers(select_index)
|
|
if bbox_select == "all":
|
|
return bboxes
|
|
elif bbox_select == "first":
|
|
return [bboxes[0]]
|
|
elif bbox_select == "by_index":
|
|
new_bboxes = []
|
|
for i in indexs:
|
|
try:
|
|
new_bboxes.append(bboxes[i])
|
|
except IndexError:
|
|
log(f"Object detector output by_index: invalid bbox index {i}", message_type='warning')
|
|
return new_bboxes
|
|
|
|
|
|
class LS_BBOXES_JOIN:
|
|
|
|
def __init__(self):
|
|
self.NODE_NAME = 'BBoxes Join'
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
|
|
return {
|
|
"required": {
|
|
"bboxes_1": ("BBOXES",),
|
|
},
|
|
"optional": {
|
|
"bboxes_2": ("BBOXES",),
|
|
"bboxes_3": ("BBOXES",),
|
|
"bboxes_4": ("BBOXES",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BBOXES",)
|
|
RETURN_NAMES = ("bboxes",)
|
|
FUNCTION = 'bboxes_join'
|
|
CATEGORY = '😺dzNodes/LayerMask'
|
|
|
|
def bboxes_join(self, bboxes_1, bboxes_2=None, bboxes_3=None, bboxes_4=None):
|
|
if bboxes_2 is not None:
|
|
bboxes_1.extend(bboxes_2)
|
|
if bboxes_3 is not None:
|
|
bboxes_1.extend(bboxes_3)
|
|
if bboxes_4 is not None:
|
|
bboxes_1.extend(bboxes_4)
|
|
return (bboxes_1,)
|
|
|
|
class LS_OBJECT_DETECTOR_FL2:
|
|
|
|
def __init__(self):
|
|
self.NODE_NAME = 'Object Detector Florence2'
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ), #
|
|
"prompt": ("STRING", {"default": "subject"}),
|
|
"florence2_model": ("FLORENCE2",),
|
|
"sort_method": (sort_method_list,),
|
|
"bbox_select": (select_list,),
|
|
"select_index": ("STRING", {"default": "0,"},),
|
|
},
|
|
"optional": {
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BBOXES", "IMAGE",)
|
|
RETURN_NAMES = ("bboxes", "preview",)
|
|
FUNCTION = 'object_detector_fl2'
|
|
CATEGORY = '😺dzNodes/LayerMask'
|
|
|
|
def object_detector_fl2(self, image, prompt, florence2_model, sort_method, bbox_select, select_index):
|
|
|
|
bboxes = []
|
|
ret_previews = []
|
|
max_new_tokens = 512
|
|
num_beams = 3
|
|
do_sample = False
|
|
fill_mask = False
|
|
|
|
model = florence2_model['model']
|
|
processor = florence2_model['processor']
|
|
|
|
img = tensor2pil(image[0]).convert("RGB")
|
|
task = 'caption to phrase grounding'
|
|
from .florence2_ultra import process_image
|
|
results, _ = process_image(model, processor, img, task,
|
|
max_new_tokens, num_beams, do_sample,
|
|
fill_mask, prompt)
|
|
|
|
if isinstance(results, dict):
|
|
results["width"] = img.width
|
|
results["height"] = img.height
|
|
|
|
bboxes = self.fbboxes_to_list(results)
|
|
bboxes = sort_bboxes(bboxes, sort_method)
|
|
bboxes = select_bboxes(bboxes, bbox_select, select_index)
|
|
preview = draw_bounding_boxes(img, bboxes, color="random", line_width=-1)
|
|
ret_previews.append(pil2tensor(preview))
|
|
if len(bboxes) == 0:
|
|
log(f"{self.NODE_NAME} no object found", message_type='warning')
|
|
else:
|
|
log(f"{self.NODE_NAME} found {len(bboxes)} object(s)", message_type='info')
|
|
return (bboxes, torch.cat(ret_previews, dim=0))
|
|
|
|
def fbboxes_to_list(self, F_BBOXES) -> list:
|
|
if isinstance(F_BBOXES, str):
|
|
return None
|
|
ret_bboxes = []
|
|
width = F_BBOXES["width"]
|
|
height = F_BBOXES["height"]
|
|
x1_c = width
|
|
y1_c = height
|
|
x2_c = y2_c = 0
|
|
label = ""
|
|
if "bboxes" in F_BBOXES:
|
|
for idx in range(len(F_BBOXES["bboxes"])):
|
|
bbox = F_BBOXES["bboxes"][idx]
|
|
new_label = F_BBOXES["labels"][idx].removeprefix("</s>")
|
|
if new_label not in label:
|
|
if idx > 0:
|
|
label = label + ", "
|
|
label = label + new_label
|
|
if len(bbox) == 4:
|
|
x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
|
|
elif len(bbox) == 8:
|
|
x1 = int(min(bbox[0::2]))
|
|
x2 = int(max(bbox[0::2]))
|
|
y1 = int(min(bbox[1::2]))
|
|
y2 = int(max(bbox[1::2]))
|
|
else:
|
|
continue
|
|
x1_c = min(x1_c, x1)
|
|
y1_c = min(y1_c, y1)
|
|
x2_c = max(x2_c, x2)
|
|
y2_c = max(y2_c, y2)
|
|
ret_bboxes.append([x1, y1, x2, y2])
|
|
else:
|
|
x1_c = width
|
|
y1_c = height
|
|
x2_c = y2_c = 0
|
|
for polygon in F_BBOXES["polygons"][0]:
|
|
if len(_polygon) < 3:
|
|
print('Invalid polygon:', _polygon)
|
|
continue
|
|
x1_c = min(x1_c, int(min(polygon[0::2])))
|
|
x2_c = max(x2_c, int(max(polygon[0::2])))
|
|
y1_c = min(y1_c, int(min(polygon[1::2])))
|
|
y2_c = max(y2_c, int(max(polygon[1::2])))
|
|
ret_bboxes.append(x1_c, y1_c, x2_c, y2_c)
|
|
return ret_bboxes
|
|
|
|
class LS_OBJECT_DETECTOR_MASK:
|
|
|
|
def __init__(self):
|
|
self.NODE_NAME = 'Object Detector MASK'
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
|
|
return {
|
|
"required": {
|
|
"object_mask": ("MASK",),
|
|
"sort_method": (sort_method_list,),
|
|
"bbox_select": (select_list,),
|
|
"select_index": ("STRING", {"default": "0,"},),
|
|
},
|
|
"optional": {
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BBOXES", "IMAGE",)
|
|
RETURN_NAMES = ("bboxes", "preview",)
|
|
FUNCTION = 'object_detector_mask'
|
|
CATEGORY = '😺dzNodes/LayerMask'
|
|
|
|
def object_detector_mask(self, object_mask, sort_method, bbox_select, select_index):
|
|
|
|
bboxes = []
|
|
if object_mask.dim() == 2:
|
|
object_mask = torch.unsqueeze(object_mask, 0)
|
|
|
|
cv_mask = tensor2cv2(object_mask[0])
|
|
cv_mask = cv2.cvtColor(cv_mask, cv2.COLOR_BGR2GRAY)
|
|
_, binary = cv2.threshold(cv_mask, 127, 255, cv2.THRESH_BINARY)
|
|
# invert mask
|
|
# binary = cv2.bitwise_not(binary)
|
|
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
for contour in contours:
|
|
x, y, w, h = cv2.boundingRect(contour)
|
|
bboxes.append([x, y, x + w, y + h])
|
|
bboxes = sort_bboxes(bboxes, sort_method)
|
|
bboxes = select_bboxes(bboxes, bbox_select, select_index)
|
|
ret_previews = []
|
|
preview = draw_bounding_boxes(tensor2pil(object_mask[0]).convert("RGB"), bboxes, color="random", line_width=-1)
|
|
ret_previews.append(pil2tensor(preview))
|
|
|
|
if len(bboxes) == 0:
|
|
log(f"{self.NODE_NAME} no object found", message_type='warning')
|
|
else:
|
|
log(f"{self.NODE_NAME} found {len(bboxes)} object(s)", message_type='info')
|
|
|
|
return (bboxes, torch.cat(ret_previews, dim=0))
|
|
|
|
|
|
class LS_OBJECT_DETECTOR_YOLO8:
|
|
|
|
def __init__(self):
|
|
self.NODE_NAME = 'Object Detector YOLO8'
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_ext = [".pt"]
|
|
model_path = os.path.join(folder_paths.models_dir, 'yolo')
|
|
FILES_DICT = get_files(model_path, model_ext)
|
|
FILE_LIST = list(FILES_DICT.keys())
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ),
|
|
"yolo_model": (FILE_LIST,),
|
|
"sort_method": (sort_method_list,),
|
|
"bbox_select": (select_list,),
|
|
"select_index": ("STRING", {"default": "0,"},),
|
|
},
|
|
"optional": {
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BBOXES", "IMAGE",)
|
|
RETURN_NAMES = ("bboxes", "preview",)
|
|
FUNCTION = 'object_detector_yolo8'
|
|
CATEGORY = '😺dzNodes/LayerMask'
|
|
|
|
def object_detector_yolo8(self, image, yolo_model, sort_method, bbox_select, select_index):
|
|
|
|
from ultralytics import YOLO
|
|
model_path = os.path.join(folder_paths.models_dir, 'yolo')
|
|
yolo_model = YOLO(os.path.join(model_path, yolo_model))
|
|
|
|
bboxes = []
|
|
ret_previews = []
|
|
|
|
img = torch.unsqueeze(image[0], 0)
|
|
_image = tensor2pil(img)
|
|
results = yolo_model(_image, retina_masks=True)
|
|
for result in results:
|
|
yolo_plot_image = cv2.cvtColor(result.plot(), cv2.COLOR_BGR2RGB)
|
|
|
|
# no mask, if have box, draw box
|
|
if result.boxes is not None and len(result.boxes.xyxy) > 0:
|
|
for box in result.boxes:
|
|
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
|
|
bboxes.append([x1, y1, x2, y2])
|
|
bboxes = sort_bboxes(bboxes, sort_method)
|
|
bboxes = select_bboxes(bboxes, bbox_select, select_index)
|
|
preview = draw_bounding_boxes(_image.convert("RGB"), bboxes, color="random", line_width=-1)
|
|
ret_previews.append(pil2tensor(preview))
|
|
|
|
if len(bboxes) == 0:
|
|
log(f"{self.NODE_NAME} no object found", message_type='warning')
|
|
else:
|
|
log(f"{self.NODE_NAME} found {len(bboxes)} object(s)", message_type='info')
|
|
|
|
return (bboxes, torch.cat(ret_previews, dim=0),)
|
|
|
|
class LS_OBJECT_DETECTOR_YOLOWORLD:
|
|
|
|
def __init__(self):
|
|
self.NODE_NAME = 'Object Detector YOLO-WORLD'
|
|
self.model_path = os.path.join(folder_paths.models_dir, 'yolo-world')
|
|
os.environ['MODEL_CACHE_DIR'] = self.model_path
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_list =['yolo_world/v2-x', 'yolo_world/v2-l', 'yolo_world/v2-m',
|
|
'yolo_world/v2-s', 'yolo_world/l', 'yolo_world/m',
|
|
'yolo_world/s']
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ),
|
|
"yolo_world_model": (model_list,),
|
|
"confidence_threshold": ("FLOAT", {"default": 0.05, "min": 0, "max": 1, "step": 0.01}),
|
|
"nms_iou_threshold": ("FLOAT", {"default": 0.3, "min": 0, "max": 1, "step": 0.01}),
|
|
"prompt": ("STRING", {"default": "subject"}),
|
|
"sort_method": (sort_method_list,),
|
|
"bbox_select": (select_list,),
|
|
"select_index": ("STRING", {"default": "0,"},),
|
|
},
|
|
"optional": {
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BBOXES", "IMAGE",)
|
|
RETURN_NAMES = ("bboxes", "preview",)
|
|
FUNCTION = 'object_detector_yoloworld'
|
|
CATEGORY = '😺dzNodes/LayerMask'
|
|
|
|
def object_detector_yoloworld(self, image, yolo_world_model,
|
|
confidence_threshold, nms_iou_threshold, prompt,
|
|
sort_method, bbox_select, select_index):
|
|
import supervision as sv
|
|
|
|
model=self.load_yolo_world_model(yolo_world_model, prompt)
|
|
infer_outputs = []
|
|
img = (255 * image[0].cpu().numpy()).astype(np.uint8)
|
|
results = model.infer(
|
|
img, confidence=confidence_threshold)
|
|
detections = sv.Detections.from_inference(results)
|
|
detections = detections.with_nms(
|
|
class_agnostic=False,
|
|
threshold=nms_iou_threshold
|
|
)
|
|
infer_outputs.append(detections)
|
|
bboxes = infer_outputs[0].xyxy.tolist()
|
|
bboxes = [[int(value) for value in sublist] for sublist in bboxes]
|
|
bboxes = sort_bboxes(bboxes, sort_method)
|
|
bboxes = select_bboxes(bboxes, bbox_select, select_index)
|
|
ret_previews = []
|
|
preview = draw_bounding_boxes(tensor2pil(image[0]).convert('RGB'), bboxes, color="random", line_width=-1)
|
|
ret_previews.append(pil2tensor(preview))
|
|
|
|
if len(bboxes) == 0:
|
|
log(f"{self.NODE_NAME} no object found", message_type='warning')
|
|
else:
|
|
log(f"{self.NODE_NAME} found {len(bboxes)} object(s)", message_type='info')
|
|
|
|
return (bboxes, torch.cat(ret_previews, dim=0))
|
|
|
|
def process_categories(self, categories: str) -> List[str]:
|
|
return [category.strip().lower() for category in categories.split(',')]
|
|
|
|
def load_yolo_world_model(self,model_id: str, categories: str) -> List[torch.nn.Module]:
|
|
from inference.models import YOLOWorld as YOLOWorldImpl
|
|
model = YOLOWorldImpl(model_id=model_id)
|
|
categories = self.process_categories(categories)
|
|
model.set_classes(categories)
|
|
return model
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LayerMask: BBoxJoin": LS_BBOXES_JOIN,
|
|
"LayerMask: ObjectDetectorFL2": LS_OBJECT_DETECTOR_FL2,
|
|
"LayerMask: ObjectDetectorMask": LS_OBJECT_DETECTOR_MASK,
|
|
"LayerMask: ObjectDetectorYOLO8": LS_OBJECT_DETECTOR_YOLO8,
|
|
"LayerMask: ObjectDetectorYOLOWorld": LS_OBJECT_DETECTOR_YOLOWORLD
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LayerMask: BBoxJoin": "LayerMask: BBox Join",
|
|
"LayerMask: ObjectDetectorFL2": "LayerMask: Object Detector Florence2",
|
|
"LayerMask: ObjectDetectorMask": "LayerMask: Object Detector Mask",
|
|
"LayerMask: ObjectDetectorYOLO8": "LayerMask: Object Detector YOLO8",
|
|
"LayerMask: ObjectDetectorYOLOWorld": "LayerMask: Object Detector YOLO World"
|
|
} |