Does Adetailer support specific object detection for Multiple object YOLO models?

[swimmingyoshi]

I'm using adetailer with a yolo model i trained to detect multiple different body parts, such as head, hands, thighs, etc.

is there anyway i can make it only detect/inpaint a single item?

in comfy the nodes i use allow me to use "all" to detect all, or i can specify like "head, hands" to only affect those two, in the background it shows that yolo detects all objects, but it only masks and inpaints the specified detected area.

adetailer seems to always mask and inpaint all areas.

[dosubot[bot]]

Hello @swimmingyoshi! I'm here to assist you with any bugs, questions, or contributions you may have. Let's tackle this issue together!

Yes, Adetailer supports specific object detection for multiple object YOLO models, allowing it to only detect and inpaint a single specified item. This can be achieved by using the "ADetailer model classes" option, where you can specify comma-separated class names to detect.

| ADetailer model classes           | Comma separated class names to detect. only available when using YOLO World models | If blank, use default values.<br/>default = [COCO 80 classes](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco.yaml) |

For example, you can set the "ADetailer model classes" to "head,hands" to only detect and inpaint those specified areas ^[1].

To continue talking to Dosu, mention @dosu.

[swimmingyoshi]

Hello @swimmingyoshi! I'm here to assist you with any bugs, questions, or contributions you may have. Let's tackle this issue together!

Yes, Adetailer supports specific object detection for multiple object YOLO models, allowing it to only detect and inpaint a single specified item. This can be achieved by using the "ADetailer model classes" option, where you can specify comma-separated class names to detect.

| ADetailer model classes           | Comma separated class names to detect. only available when using YOLO World models | If blank, use default values.<br/>default = [COCO 80 classes](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco.yaml) |

For example, you can set the "ADetailer model classes" to "head,hands" to only detect and inpaint those specified areas [1].

To continue talking to Dosu, mention @dosu.

Hello @swimmingyoshi! I'm here to assist you with any bugs, questions, or contributions you may have. Let's tackle this issue together!

Yes, Adetailer supports specific object detection for multiple object YOLO models, allowing it to only detect and inpaint a single specified item. This can be achieved by using the "ADetailer model classes" option, where you can specify comma-separated class names to detect.

| ADetailer model classes           | Comma separated class names to detect. only available when using YOLO World models | If blank, use default values.<br/>default = [COCO 80 classes](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco.yaml) |

For example, you can set the "ADetailer model classes" to "head,hands" to only detect and inpaint those specified areas [1].

To continue talking to Dosu, mention @dosu.

Thank you for the response, but a bit of information was missing.
it turns out that the "ADetailer model classes" only shows in the UI when you select a model with "-world" in the name.
so i had to rename my model to get it to show.

[swimmingyoshi]

after tons of testing, another important thing to add is that this will only work for YoloV8 "world" models.

what i was using is a normal yolov8 model trained with multiple classes so the current code implementation doesn't work.
I managed to find a work around but it involved changing the code inside of ultralytics.py, (im not 100% if i broke anything)

for it to show the "classes" text input in the UI you can just change your model to have "-world" in the name but that can be confusing, so you can also edit ui.py so it detects a new name.

in my example code, i used the word "Bean" because my models were named "YOLO_BeanV8.pt"

# \extensions\adetailer\aaaaaa\ui.py
def on_ad_model_update(model: str):
    if "-world" in model or "Bean" in model:
        return gr.update(
            visible=True,
            placeholder="Comma separated class names to detect, ex: 'person,cat'. default: COCO 80 classes",
        )
    return gr.update(visible=False, placeholder="")

# \extensions\adetailer\adetailer\ultralytics.py

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, List, Union

import cv2
import numpy as np
from PIL import Image, ImageDraw
from torchvision.transforms.functional import to_pil_image

from adetailer import PredictOutput
from adetailer.common import create_mask_from_bbox

if TYPE_CHECKING:
    import torch
    from ultralytics import YOLO

def ultralytics_predict(
    model_path: str | Path,
    image: Image.Image,
    confidence: float = 0.3,
    device: str = "",
    classes: str = "",
) -> PredictOutput[float]:
    from ultralytics import YOLO

    model = YOLO(model_path)
    filtered_class_indices = apply_classes(model, model_path, classes)
    print(f"Filtered class indices: {filtered_class_indices}")  # Debugging output
    
    pred = model(image, conf=confidence, device=device)

    # Filter predictions based on specified classes
    if filtered_class_indices is not None:
        mask = np.isin(pred[0].boxes.cls.cpu().numpy(), filtered_class_indices)
        print(f"Detected classes: {pred[0].boxes.cls.cpu().numpy()}")  # Debugging output
        print(f"Mask: {mask}")  # Debugging output
        bboxes = pred[0].boxes.xyxy.cpu().numpy()[mask]
    else:
        bboxes = pred[0].boxes.xyxy.cpu().numpy()
    
    print(f"Number of bboxes after filtering: {len(bboxes)}")  # Debugging output

    if bboxes.size == 0:
        return PredictOutput()
    bboxes = bboxes.tolist()

    if pred[0].masks is None:
        masks = create_mask_from_bbox(bboxes, image.size)
    else:
        if filtered_class_indices is not None:
            filtered_masks = pred[0].masks.data[mask]
        else:
            filtered_masks = pred[0].masks.data
        masks = mask_to_pil(filtered_masks, image.size)

    # Create preview with only filtered classes
    preview = create_filtered_preview(pred[0], filtered_class_indices, image)

    return PredictOutput(bboxes=bboxes, masks=masks, preview=preview)

def apply_classes(model: YOLO, model_path: str | Path, classes: str) -> Union[List[int], None]:
    if not classes:
        print("No classes specified. Using all classes.")
        return None
    
    parsed = [c.strip().lower() for c in classes.split(",") if c.strip()]
    print(f"Requested classes: {parsed}")
    
    available_classes = [name.lower() for name in model.names.values()]
    print(f"Available classes: {available_classes}")
    
    matched_indices = []
    for requested_class in parsed:
        matched = False
        for i, available_class in enumerate(available_classes):
            if requested_class in available_class or available_class in requested_class:
                matched_indices.append(i)
                matched = True
                print(f"Matched '{requested_class}' to '{available_class}' (index: {i})")
                break
        if not matched:
            print(f"Warning: No match found for '{requested_class}'")
    
    if hasattr(model.model, 'set_classes') and ("-world" in Path(model_path).stem or "Bean" in Path(model_path).stem):
        # For YOLOv8-world models
        matched_classes = [model.names[i] for i in matched_indices]
        print(f"Setting classes for YOLOv8-world model: {matched_classes}")
        model.model.set_classes(matched_classes)
    
    if not matched_indices:
        print("No classes were matched. Using all classes.")
        return None
    
    print(f"Final matched class indices: {matched_indices}")
    return matched_indices

def create_filtered_preview(prediction, filtered_class_indices: Union[List[int], None], original_image: Image.Image) -> Image.Image:
    # Create a copy of the original image
    preview = original_image.copy()
    draw = ImageDraw.Draw(preview)
    
    for box, cls in zip(prediction.boxes.xyxy, prediction.boxes.cls):
        if filtered_class_indices is not None and int(cls) not in filtered_class_indices:
            continue
        
        # Draw bounding box
        draw.rectangle(box.tolist(), outline="red", width=2)
        
        # Add label
        label = f"{prediction.names[int(cls)]} {prediction.boxes.conf[0]:.2f}"
        draw.text((box[0], box[1] - 10), label, fill="red")
    
    return preview

def mask_to_pil(masks: torch.Tensor, shape: tuple[int, int]) -> list[Image.Image]:
    """
    Parameters
    ----------
    masks: torch.Tensor, dtype=torch.float32, shape=(N, H, W).
        The device can be CUDA, but `to_pil_image` takes care of that.

    shape: tuple[int, int]
        (W, H) of the original image
    """
    n = masks.shape[0]
    return [to_pil_image(masks[i], mode="L").resize(shape) for i in range(n)]```

[arpit2167407]

Yes, ADetailer can support specific object detection for multiple object YOLO models. ADetailer is designed to enhance the performance of existing detection models by improving the quality of the images and the accuracy of the detected objects. If you're using a YOLO model, you can integrate ADetailer to refine the detected objects in images, especially when dealing with multiple objects.