|
|
| import folder_paths
|
| from .imagefunc import *
|
|
|
| select_list = ["all", "first", "by_index"]
|
| sort_method_list = ["left_to_right", "top_to_bottom", "big_to_small", "confidence"]
|
|
|
|
|
| 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])
|
| elif method == "big_to_small":
|
| sorted_bboxes = sorted(bboxes, key=lambda bbox: (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), reverse=True)
|
| else:
|
| sorted_bboxes = bboxes
|
| 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):
|
| all_inputs = [b for b in [bboxes_2, bboxes_3, bboxes_4] if b is not None]
|
| for other in all_inputs:
|
| for i in range(len(other)):
|
| if i < len(bboxes_1):
|
| bboxes_1[i].extend(other[i])
|
| else:
|
| bboxes_1.append(other[i])
|
| 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):
|
|
|
| ret_bboxes = []
|
| ret_previews = []
|
| max_new_tokens = 512
|
| num_beams = 3
|
| do_sample = False
|
| fill_mask = False
|
|
|
| model = florence2_model['model']
|
| processor = florence2_model['processor']
|
|
|
| for img in image:
|
| bboxes = []
|
| img = tensor2pil(img.unsqueeze(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))
|
| ret_bboxes.append(standardize_bbox(bboxes))
|
|
|
|
|
|
|
|
|
|
|
|
|
| return (ret_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])
|
| if len(ret_bboxes) == 0:
|
| 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):
|
|
|
| ret_bboxes = []
|
| ret_previews = []
|
|
|
| if object_mask.dim() == 2:
|
| object_mask = torch.unsqueeze(object_mask, 0)
|
|
|
| for msk in object_mask:
|
| bboxes = []
|
| cv_mask = tensor2cv2(msk)
|
| cv_mask = cv2.cvtColor(cv_mask, cv2.COLOR_BGR2GRAY)
|
| _, binary = cv2.threshold(cv_mask, 127, 255, cv2.THRESH_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)
|
| preview = draw_bounding_boxes(tensor2pil(msk).convert("RGB"), bboxes, color="random", line_width=-1)
|
| ret_previews.append(pil2tensor(preview))
|
|
|
|
|
|
|
|
|
|
|
|
|
| ret_bboxes.append(standardize_bbox(bboxes))
|
|
|
| return (ret_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))
|
|
|
| ret_bboxes = []
|
| ret_previews = []
|
|
|
| for img in image:
|
| bboxes = []
|
| img = torch.unsqueeze(img.unsqueeze(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)
|
|
|
|
|
| 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))
|
|
|
|
|
|
|
|
|
|
|
|
|
| ret_bboxes.append(standardize_bbox(bboxes))
|
|
|
| return (ret_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):
|
| ret_previews = []
|
| ret_bboxes = []
|
|
|
| try:
|
| import supervision as sv
|
| except ImportError as e:
|
| log(f"{self.NODE_NAME}: {e}", message_type='warning')
|
| return None
|
| model=self.load_yolo_world_model(yolo_world_model, prompt)
|
|
|
| for i in image:
|
| infer_outputs = []
|
|
|
| img = tensor2np(i)
|
| 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)
|
|
|
|
|
| preview = draw_bounding_boxes(tensor2pil(i.unsqueeze(0)).convert('RGB'), bboxes, color="random", line_width=-1)
|
| ret_previews.append(pil2tensor(preview))
|
|
|
|
|
|
|
|
|
|
|
|
|
| ret_bboxes.append(standardize_bbox(bboxes))
|
|
|
| return (ret_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]:
|
| try:
|
| from inference.models import YOLOWorld as YOLOWorldImpl
|
| except ImportError as e:
|
| log(f"{self.NODE_NAME}: {e}", message_type='warning')
|
| return None
|
| model = YOLOWorldImpl(model_id=model_id)
|
| categories = self.process_categories(categories)
|
| model.set_classes(categories)
|
| return model
|
|
|
|
|
| class LS_DrawBBoxMask:
|
|
|
| def __init__(self):
|
| self.NODE_NAME = 'Draw BBOX Mask'
|
| pass
|
|
|
| @classmethod
|
| def INPUT_TYPES(cls):
|
| return {
|
| "required": {
|
| "image": ("IMAGE",),
|
| "bboxes": ("BBOXES",),
|
| "grow_top": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
|
| "grow_bottom": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
|
| "grow_left": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
|
| "grow_right": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
|
| },
|
| "optional": {
|
| }
|
| }
|
|
|
| RETURN_TYPES = ("MASK",)
|
| RETURN_NAMES = ("mask",)
|
| FUNCTION = 'draw_bbox_mask'
|
| CATEGORY = '😺dzNodes/LayerMask'
|
|
|
| def draw_bbox_mask(self, image, bboxes, grow_top, grow_bottom, grow_left, grow_right
|
| ):
|
|
|
| ret_masks = []
|
| for index in range(len(image)):
|
| img = tensor2pil(image[index].unsqueeze(0))
|
| mask = Image.new("L", img.size, color='black')
|
| bboxes_i = bboxes[index]
|
| for bbox in bboxes_i:
|
| try:
|
| if len(bbox) == 0:
|
| continue
|
| else:
|
| x1, y1, x2, y2 = bbox
|
| except ValueError:
|
| if len(bbox) == 0:
|
| continue
|
| else:
|
| x1, y1, x2, y2 = bbox[index]
|
| w = x2 - x1
|
| h = y2 - y1
|
| if grow_top:
|
| y1 = int(y1 - h * grow_top)
|
| if grow_bottom:
|
| y2 = int(y2 + h * grow_bottom)
|
| if grow_left:
|
| x1 = int(x1 - w * grow_left)
|
| if grow_right:
|
| x2 = int(x2 + w * grow_right)
|
| if y1 > y2 or x1 > x2:
|
| log(f"{self.NODE_NAME} Invalid bbox after extend: ({x1},{y1},{x2},{y2})", message_type='warning')
|
| continue
|
| draw = ImageDraw.Draw(mask)
|
| draw.rectangle([x1, y1, x2, y2], fill='white', outline='white', width=0)
|
| ret_masks.append(pil2tensor(mask))
|
|
|
| log(f"{self.NODE_NAME} Processed {len(ret_masks)} mask(s).", message_type='finish')
|
| return (torch.cat(ret_masks, dim=0),)
|
|
|
|
|
| class LS_DrawBBoxMaskV2:
|
|
|
| def __init__(self):
|
| self.NODE_NAME = 'Draw BBOX Mask V2'
|
| pass
|
|
|
| @classmethod
|
| def INPUT_TYPES(cls):
|
| return {
|
| "required": {
|
| "image": ("IMAGE",),
|
| "bboxes": ("BBOXES",),
|
| "grow_top": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
|
| "grow_bottom": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
|
| "grow_left": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
|
| "grow_right": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 0.01}),
|
| "rounded_rect_radius": ("INT", {"default": 50, "min": 0, "max": 100, "step": 1}),
|
| "anti_aliasing": ("INT", {"default": 2, "min": 0, "max": 16, "step": 1}),
|
| },
|
| "optional": {
|
| }
|
| }
|
|
|
| RETURN_TYPES = ("MASK",)
|
| RETURN_NAMES = ("mask",)
|
| FUNCTION = 'draw_bbox_mask_v2'
|
| CATEGORY = '😺dzNodes/LayerMask'
|
|
|
| def draw_bbox_mask_v2(self, image, bboxes, grow_top, grow_bottom, grow_left, grow_right,
|
| rounded_rect_radius, anti_aliasing):
|
|
|
| ret_masks = []
|
| for index in range(len(image)):
|
| img = tensor2pil(image[index].unsqueeze(0))
|
| mask = Image.new("L", img.size, color='black')
|
| bboxes_i = bboxes[index]
|
| if grow_top or grow_bottom or grow_left or grow_right:
|
| new_bboxes_i = []
|
| for bbox in bboxes_i:
|
| try:
|
| if len(bbox) == 0:
|
| continue
|
| else:
|
| x1, y1, x2, y2 = bbox
|
| except ValueError:
|
| if len(bbox) == 0:
|
| continue
|
| else:
|
| x1, y1, x2, y2 = bbox[index]
|
| w = x2 - x1
|
| h = y2 - y1
|
| if grow_top:
|
| y1 = int(y1 - h * grow_top)
|
| if grow_bottom:
|
| y2 = int(y2 + h * grow_bottom)
|
| if grow_left:
|
| x1 = int(x1 - w * grow_left)
|
| if grow_right:
|
| x2 = int(x2 + w * grow_right)
|
| if y1 > y2:
|
| y1, y2 = y2, y1
|
| if x1 > x2:
|
| x1, x2 = x2, x1
|
| if y2 - y1 < 1:
|
| y2 += 1
|
| if x2 - x1 < 1:
|
| x2 += 1
|
| new_bboxes_i.append((x1, y1, x2, y2))
|
| bboxes_i = new_bboxes_i
|
| mask = draw_rounded_rectangle(mask, rounded_rect_radius, bboxes_i, anti_aliasing)
|
| ret_masks.append(pil2tensor(mask))
|
|
|
| log(f"{self.NODE_NAME} Processed {len(ret_masks)} mask(s).", message_type='finish')
|
| return (torch.cat(ret_masks, dim=0),)
|
|
|
|
|
| NODE_CLASS_MAPPINGS = {
|
| "LayerMask: BBoxJoin": LS_BBOXES_JOIN,
|
| "LayerMask: DrawBBoxMaskV2": LS_DrawBBoxMaskV2,
|
| "LayerMask: DrawBBoxMask": LS_DrawBBoxMask,
|
| "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(Advance)",
|
| "LayerMask: DrawBBoxMaskV2": "LayerMask: Draw BBox Mask V2(Advance)",
|
| "LayerMask: DrawBBoxMask": "LayerMask: Draw BBox Mask(Advance)",
|
| "LayerMask: ObjectDetectorFL2": "LayerMask: Object Detector Florence2(Advance)",
|
| "LayerMask: ObjectDetectorMask": "LayerMask: Object Detector Mask(Advance)",
|
| "LayerMask: ObjectDetectorYOLO8": "LayerMask: Object Detector YOLO8(Advance)",
|
| "LayerMask: ObjectDetectorYOLOWorld": "LayerMask: Object Detector YOLO World(Obsolete)"
|
| }
|
|
|
