[Feature Request] Separate Mask Output as a List of BBOXes

[nomadoor]

My apologies if I'm misunderstanding the current functionality.
When the Florence2 node generates BBOXes, multiple BBOXes are currently merged into a single mask. This creates an issue where, if BBOXes overlap, outputting the mask causes some BBOXes to be swallowed by others.

Proposed Features

• A function to output individual masks as a list for each BBOX when two or more objects are detected.
• Alternatively (though I believe this was suggested in another issue), adding an output for Florence2 as a bbox_detector in Impact Pack.

Thanks for your continued great work.

[mnpenner]

I'm having the same issue. I can't figure out how to get each face as a separate mask.

[kijai]

There's a separate masks node in KJNodes and probably in other node packs too, that you can use for now.

[nomadoor]

@kijai
Thanks, that node does work, but only if the objects to be detected are completely separated.
For example, in this case, the two women on the right are close together, so their masks are merged into one.

I tried passing a list of values from 0 to N into output_mask_select, and it successfully produced separate masks for each BBOX individually.

[HerlanderCarvalho]

I'm having a similar issue, but related to another aspect, which is:

When using "ocr_to_region", the BOXES OUTPUT BORKS, because it is (I'm assuming) a completely different format. In addition, when I try to get the masks only (which is really just what I need) to make a text inpaint replacement with a detailer, the masks return as a single mask, preventing me from specifying which text to replace and/or with what text.

I briefly looked at the code and for me the solution was simple. I moved the lines 596-598 to the beginning of the FOR loop, and then lines 636-642 to the end of the same loop. This will output all the OCR masks found into split ones. But I took the extra effort, and make a code that would fill the above users the needs to have the functionality they require when using "caption_to_phrase_grounding".

I don't want to break anything or give the impression of having the presumption that I know what I'm doing and whether other people's code should be changed or not, so I'm just doing this for my personal use. However, I will leave the appropriate changes for the "ocr_to_region" and "caption_to phrase_grounding", to output the result as a list of masks, instead of a single one. The "vanilla" behavior is still present, the only difference is that if you include "split," in the "output_mask_select" it will split the masks into a list, instead of outputting a single mask. For example, in an image with 3 people, if I put "split,0,2" it will only return the first and 3rd person found, but also as split masks. If you don't specify indexes, then the whichever is found is also split. I only mde this for "ocr_with_region" and "caption_to_phrase_grounding", I can add the same behavior to the rest if anyone needs, but don't count on me actively watching this thread, other than for maybe a week or so.

Also a reminder that until Kijai commits his own code, updating the nodes in ComfyUI will revert any change you decide to do manually. Anywho... should ppl want to do the appropriate manual changes for their own personal use.

ALSO A REMINDER FOR THOSE WHO ARE UNFMALIAR WITH PYTHON! Indentation matters! if you don't use the correct amount of spaces (sort of) as the provided code, it will break!

This will affect 'region_caption', 'dense_region_caption', 'caption_to_phrase_grounding' and 'region_proposal'. On nodes.py replace lines 420-510 (inclusive) with:

                mask_indexes = []
                # Determine mask indexes outside the loop

                split = False
                # Determine if results should be split into a list of masks or not
                if output_mask_select.split(",")[0] == "split":
                    output_mask_select = output_mask_select.replace("split", "")
                    split = True

                if output_mask_select != "":
                    mask_indexes = [n for n in output_mask_select.split(",")]
                    print(mask_indexes)
                else:
                    mask_indexes = [str(i) for i in range(len(bboxes))]

                for index, (bbox, label) in enumerate(zip(bboxes, labels)):
                    # if output_mask_select begins with "split", then create a list, otherwise just return a single mask
                    if index == 0 or split:
                        # Initialize mask_layer only if needed
                        if fill_mask and str(index) in mask_indexes:
                            print("mask_draw: " + str(index))
                            mask_layer = Image.new('RGB', image_pil.size, (0, 0, 0))
                            mask_draw = ImageDraw.Draw(mask_layer)

                    # Modify the label to include the index
                    indexed_label = f"{index}.{label}"
                    
                    if fill_mask:
                        # Ensure y1 is greater than or equal to y0 for mask drawing
                        x0, y0, x1, y1 = bbox[0], bbox[1], bbox[2], bbox[3]
                        if y1 < y0:
                            y0, y1 = y1, y0
                        if x1 < x0:
                            x0, x1 = x1, x0
                            
                        if str(index) in mask_indexes:
                            print("match index:", str(index), "in mask_indexes:", mask_indexes)
                            mask_draw.rectangle([x0, y0, x1, y1], fill=(255, 255, 255))
                        if label in mask_indexes:
                            print("match label")
                            mask_draw.rectangle([x0, y0, x1, y1], fill=(255, 255, 255))

                    # Create a Rectangle patch
                    # Ensure y1 is greater than or equal to y0
                    y0, y1 = bbox[1], bbox[3]
                    if y1 < y0:
                        y0, y1 = y1, y0
                    
                    rect = patches.Rectangle(
                        (bbox[0], y0),  # (x,y) - lower left corner
                        bbox[2] - bbox[0],   # Width
                        y1 - y0,   # Height
                        linewidth=1,
                        edgecolor='r',
                        facecolor='none',
                        label=indexed_label
                    )
                     # Calculate text width with a rough estimation
                    text_width = len(label) * 6  # Adjust multiplier based on your font size
                    text_height = 12  # Adjust based on your font size

                    # Get corrected coordinates
                    x0, y0, x1, y1 = bbox[0], bbox[1], bbox[2], bbox[3]
                    if y1 < y0:
                        y0, y1 = y1, y0
                    if x1 < x0:
                        x0, x1 = x1, x0

                    # Initial text position
                    text_x = x0
                    text_y = y0 - text_height  # Position text above the top-left of the bbox

                    # Adjust text_x if text is going off the left or right edge
                    if text_x < 0:
                        text_x = 0
                    elif text_x + text_width > W:
                        text_x = W - text_width

                    # Adjust text_y if text is going off the top edge
                    if text_y < 0:
                        text_y = y1  # Move text below the bottom-left of the bbox if it doesn't overlap with bbox

                    # Add the rectangle to the plot
                    ax.add_patch(rect)
                    facecolor = random.choice(colormap) if len(image) == 1 else 'red'
                    # Add the label
                    plt.text(
                        text_x,
                        text_y,
                        indexed_label,
                        color='white',
                        fontsize=12,
                        bbox=dict(facecolor=facecolor, alpha=0.5)
                    )

                    if index+1 == len(bboxes) or split:
                        # Initialize mask_layer only if needed
                        if fill_mask and str(index) in mask_indexes:
                            print("mask_tensor: " + str(index))
                            mask_tensor = F.to_tensor(mask_layer)
                            mask_tensor = mask_tensor.unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
                            mask_tensor = mask_tensor.mean(dim=0, keepdim=True)
                            mask_tensor = mask_tensor.repeat(1, 1, 1, 3)
                            mask_tensor = mask_tensor[:, :, :, 0]
                            out_masks.append(mask_tensor)           

This will affect 'ocr_with_region'. On nodes.py replace lines 596-642 (inclusive) with:

            elif task == 'ocr_with_region':
                try:
                    font = ImageFont.load_default().font_variant(size=24)
                except:
                    font = ImageFont.load_default()
                predictions = parsed_answer[task_prompt]
                scale = 1
                image_pil = image_pil.convert('RGBA')
                overlay = Image.new('RGBA', image_pil.size, (255, 255, 255, 0))
                draw = ImageDraw.Draw(overlay)
                bboxes, labels = predictions['quad_boxes'], predictions['labels']
                              
                for box, label in zip(bboxes, labels):
                    # Create a new black image for the mask
                    mask_image = Image.new('RGB', (W, H), 'black')
                    mask_draw = ImageDraw.Draw(mask_image)

                    scaled_box = [v / (width if idx % 2 == 0 else height) for idx, v in enumerate(box)]
                    out_data.append({"label": label, "box": scaled_box})
                    
                    color = random.choice(colormap)
                    new_box = (np.array(box) * scale).tolist()
                    
                    # Ensure polygon coordinates are valid
                    # For polygons, we need to make sure the points form a valid shape
                    # This is a simple check to ensure the polygon has at least 3 points
                    if len(new_box) >= 6:  # At least 3 points (x,y pairs)
                        if fill_mask:
                            color_with_opacity = ImageColor.getrgb(color) + (180,)
                            draw.polygon(new_box, outline=color, fill=color_with_opacity, width=3)
                        else:
                            draw.polygon(new_box, outline=color, width=3)
                        
                        # Get the first point for text positioning
                        text_x, text_y = new_box[0]+8, new_box[1]+2
                        
                        draw.text((text_x, text_y),
                                  "{}".format(label),
                                  align="right",
                                  font=font,
                                  fill=color)
                        
                        # Draw the mask
                        mask_draw.polygon(new_box, outline="white", fill="white")
                
                    # Process the mask
                    mask_tensor = F.to_tensor(mask_image)
                    mask_tensor = mask_tensor.unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
                    mask_tensor = mask_tensor.mean(dim=0, keepdim=True)
                    mask_tensor = mask_tensor.repeat(1, 1, 1, 3)
                    mask_tensor = mask_tensor[:, :, :, 0]
                    out_masks.append(mask_tensor)

                image_pil = Image.alpha_composite(image_pil, overlay)
                image_pil = image_pil.convert('RGB')
                
                image_tensor = F.to_tensor(image_pil)
                image_tensor = image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
                out.append(image_tensor)

I would maybe do a commit, but not only I'm a noob at Python, I also never used GitHub to change other people's code, and last thing I want to do, is break anything! Last but not least, I want to once again express my gratitude to Kijai and all the projects he has been developed for the community! Thank you!!!