diff --git a/py/object_detector.py b/py/object_detector.py index f58de97..4e2745e 100644 --- a/py/object_detector.py +++ b/py/object_detector.py @@ -1,4 +1,3 @@ - from .imagefunc import * select_list = ["all", "first", "by_index"] @@ -106,7 +105,6 @@ class LS_OBJECT_DETECTOR_FL2: def object_detector_fl2(self, image, prompt, florence2_model, sort_method, bbox_select, select_index): ret_bboxes = [] - bboxes = [] ret_previews = [] max_new_tokens = 512 num_beams = 3 @@ -117,6 +115,7 @@ class LS_OBJECT_DETECTOR_FL2: 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 @@ -216,11 +215,12 @@ class LS_OBJECT_DETECTOR_MASK: ret_bboxes = [] ret_previews = [] - bboxes = [] + 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) @@ -279,10 +279,10 @@ class LS_OBJECT_DETECTOR_YOLO8: yolo_model = YOLO(os.path.join(model_path, yolo_model)) ret_bboxes = [] - 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) @@ -362,13 +362,19 @@ class LS_OBJECT_DETECTOR_YOLOWORLD: ) infer_outputs.append(detections) - if len(infer_outputs[0].xyxy) > 0: - 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) - else: - bboxes = [[0, 0, i.shape[1], i.shape[0]]] + # if len(infer_outputs[0].xyxy) > 0: + # 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) + # else: + # bboxes = [] + + 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)) @@ -423,11 +429,21 @@ class LS_DrawBBoxMask: ): ret_masks = [] - for img in image: - img = tensor2pil(img) + for index in range(len(image)): + img = tensor2pil(image[index].unsqueeze(0)) mask = Image.new("L", img.size, color='black') - for bbox in bboxes: - x1, y1, x2, y2 = bbox + 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: @@ -443,7 +459,6 @@ class LS_DrawBBoxMask: continue draw = ImageDraw.Draw(mask) draw.rectangle([x1, y1, x2, y2], fill='white', outline='white', width=0) - del draw ret_masks.append(pil2tensor(mask)) log(f"{self.NODE_NAME} Processed {len(ret_masks)} mask(s).", message_type='finish') diff --git a/py/sam_2_ultra.py b/py/sam_2_ultra.py index 66cf9da..494a5b3 100644 --- a/py/sam_2_ultra.py +++ b/py/sam_2_ultra.py @@ -325,11 +325,16 @@ class LS_SAM2_ULTRA: for index in range(len(image)): img = image[index].unsqueeze(0) + orig_image = tensor2pil(img) # Handle possible bboxes - if len(bboxes) == 0: - log(f"{self.NODE_NAME} skipped, because bboxes is empty.", message_type='error') - return (image, None) + if len(bboxes[index]) == 0: + log(f"{self.NODE_NAME} bboxes index {index} is empty, output black mask.", message_type='warning') + _mask = Image.new("L", orig_image.size, color="black") + ret_image = RGB2RGBA(orig_image, _mask.convert('L')) + ret_images.append(pil2tensor(ret_image)) + ret_masks.append(image2mask(_mask)) + continue else: boxes_np_batch = [] for bbox_list in bboxes[index]: @@ -401,7 +406,7 @@ class LS_SAM2_ULTRA: local_files_only = True else: local_files_only = False - orig_image = tensor2pil(img) + detail_range = detail_erode + detail_dilate if process_detail: if detail_method == 'GuidedFilter': diff --git a/pyproject.toml b/pyproject.toml index bdf6d4a..9c751c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_layerstyle" description = "A set of nodes for ComfyUI it generate image like Adobe Photoshop's Layer Style. the Drop Shadow is first completed node, and follow-up work is in progress." -version = "1.0.85" +version = "1.0.86" license = "MIT" dependencies = ["numpy", "pillow", "torch", "matplotlib", "Scipy", "scikit_image", "scikit_learn", "opencv-contrib-python", "pymatting", "segment_anything", "timm", "addict", "yapf", "colour-science", "wget", "mediapipe", "loguru", "typer_config", "fastapi", "rich", "google-generativeai", "diffusers", "omegaconf", "tqdm", "transformers", "kornia", "image-reward", "ultralytics", "blend_modes", "blind-watermark", "qrcode", "pyzbar", "transparent-background", "huggingface_hub", "accelerate", "bitsandbytes", "torchscale", "wandb", "hydra-core", "psd-tools", "inference-cli[yolo-world]", "inference-gpu[yolo-world]", "onnxruntime", "peft", "iopath"]