Wandb ultralytics integration get_ground_truth_bbox_annotations returns error when no bounding box

[jyuteo]

Description

This error is related to wandb ultralytics integration wandb/integration/ultralytics/bbox_utils.py When get_ground_truth_bbox_annotations is called with background image (no class label), it returns IndexError

from wandb.integration.ultralytics import add_wandb_callback
from ultralytics import YOLO

model = YOLO(..)
add_wandb_callback(model, enable_model_checkpointing=True)
results = model.train(
        project=config["project"],
        data=config["dataset"],
        epochs=config["epochs"],
        imgsz=config["image_size"],
        batch=config["batch_size"],
        resume=config["resume"],
)
# error happened during validation step when the entire validation data are background images

Error

File "/home/xxx/miniconda3/envs/vac_detector/lib/python3.9/site-packages/wandb/integration/ultralytics/bbox_utils.py", line 41, in get_ground_truth_bbox_annotations
    cls_labels = batch["cls"][indices].squeeze(1).tolist()
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

Quick fix

I have tried to comment the repeated line cls_labels = batch["cls"][indices].squeeze(1).tolist() at line 41 in wandb/integration/ultralytics/bbox_utils.py as the line will be executed at line 51 after checking if len(bboxes) == 0

def get_ground_truth_bbox_annotations(
    img_idx: int, image_path: str, batch: Dict, class_name_map: Dict = None
) -> List[Dict[str, Any]]:
    """Get ground truth bounding box annotation data in the form required for `wandb.Image` overlay system."""
    indices = batch["batch_idx"] == img_idx
    bboxes = batch["bboxes"][indices]
    # cls_labels = batch["cls"][indices].squeeze(1).tolist()

    class_name_map_reverse = {v: k for k, v in class_name_map.items()}

    if len(bboxes) == 0:
        wandb.termwarn(
            f"Image: {image_path} has no bounding boxes labels", repeat=False
        )
        return None

    cls_labels = batch["cls"][indices].squeeze(1).tolist()

Additional Files

No response

Environment

WandB version: 0.16.4

Ultralytics version: 8.0.203

Python version: 3.9

Additional Context

No response

[ArtsiomWB]

Hi @jyuteo! Thank you for writing in, reporting this, as well as trying the fix on your side. We will work on reproducing this on my machine.

[ArtsiomWB]

Hi @jyuteo, quick follow-up regarding this, would you be able to send us your full training script in here? Not having much luck reproducing this on my side yet.

[ArtsiomWB]

Hi! wanted to follow up with you regarding this thread!

[ArtsiomWB]

Hi, since we have not heard back from you we are going to close this request. If you would like to re-open the conversation, please let us know!

[jyuteo]

Hey sorry for the late follow up. Here is the training code we use

from ultralytics import YOLO
import wandb
from wandb.integration.ultralytics import add_wandb_callback

if __name__ == "__main__":
    run = wandb.init(project=config["project"], group=config["group"], tags=config["tags"])

    model = YOLO(config["yolov8_yaml"])

    add_wandb_callback(model, enable_model_checkpointing=True)
    
    results = model.train(
        project=config["project"],
        data=config["dataset"],
        epochs=config[" "],
        imgsz=config["image_size"],
        batch=config["batch_size"],
        resume=config["resume"],
    )

    wandb.finish()

[ArtsiomWB]

Thank you so much for sending it over. I have reported it to our engineering team!