fix batch image bug of ObjectDetector and SAM2Ultra nodes

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')

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':

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"]