Description

JIRAs:

• https://wandb.atlassian.net/browse/WB-5502
• https://wandb.atlassian.net/browse/WB-7361

The wandb.Image class can contain bounding box and mask annotations as well as class definitions. Until now, the Python type system was annotating Images incorrectly, which forces the Weave type system to attempt to infer the class (and annotation types) from the image content itself. This is particularly bad for two reasons: 1) inferring types from the data slows Weave down and breaks the fundamental pattern that the graph can be operated on without inspecting the data; 2) if a table has an image column, it would need to inspect every single image to get a comprehensive type for the column. Currently, we have some Python rules that enforce a single annotation type, so in the Weave system, we just inspect the first Image to determine the type. This was always a buggy part of the system that we (I?) was delaying fixing - and in particular were relying on the stricter constraints that enforced homogenous columns. However, the Growth team (@AyushExel and @morganmcg1) found some compelling use cases which need different annotation class sets within the same column (for detectron2 integation). So, finally it comes time to get this right.

First, let's understanding the high level data structure (annotated typescript style)

type Image {
     boxes: {[box_layer_id: string]: {box_data: Array<BOX_DATA_TYPE>, class_labels:{[class_id:string]: string}}}
     masks: {[mask_layer_id: string]: {mask_data: MASK_DATA_TYPE, class_labels:{[class_id:string]: string}}}
     classes: Array<{id: string, name: string}>
}

There are a number of odd parts of this:

• Classes are defined at both the layer level and the Image level
• Class definitions are in different formats (I didn't even include the "Classes" class here...)
• Boxes will autogenerate a class definition if one does not exist, and masks will do so if and only if the data is in the form of a numpy array, not a file.
  .... uggg

Something else: individual boxes (the BOX_DATA_TYPE) above can contain a dictionary of scores - which can be used to filter the boxes down.

Ok, so the goal of this PR is to:

1. Have Python fully capture Image type information such that the Weave system does not need to inspect the data
2. As part of 1, enable the Images in a column (type assignment) to incrementally build up the Image Type to be the full set of available classes - allowing for heterogeneous columns
3. As part of the above, make the Image level class object be fully representative of the entire set of classes below.

Note: this PR should not effect any user scripts - we only loosen some constraints and add additional metadata to the artifact logging. There is another PR https://github.com/wandb/core/pull/8169 which will contain the frontend changes which can consume such additional type metadata. Note: these can be deployed in any order as the current UI should not be effected by this change, and the future UI still needs to work on the old CLI.

Testing

How was this PR tested?

Unit tests + manual tests:

import wandb
import numpy as np

def make_image(num_masks=3, num_boxes=3, size = (256, 256)):
    data = np.random.randint(255, size=size)
    classes = [{"id": 0, "name": "person"}, {"id": 1, "name": "cat"}, {"id": 2, "name": "dog"}, {"id": 3, "name": "bird"}]
    masks = {}
    if num_masks > 0:
        masks["mask_layer_0_1"] = {
            "mask_data": np.random.randint(2, size=size),
        }
    if num_masks > 1:
        masks["mask_layer_2_3"] = {
            "mask_data": np.random.randint(2, size=size) + 2,
        }
    if num_masks > 2:
        masks["mask_layer_4_5"] = {
            "mask_data": np.random.randint(2, size=size) + 4,
        }
    
    boxes = {}
    if num_boxes > 0:
        boxes["box_layer_0_1"] = {
            "box_data": [
                {
                    "position": {"minX": 0.1, "minY": 0.1, "maxX": 0.2, "maxY": 0.2},
                    "class_id": 0,
                    "scores": {
                        "loss_0": 0.1,
                    }
                },
                {
                    "position": {"minX": 0.3, "minY": 0.3, "maxX": 0.4, "maxY": 0.4},
                    "class_id": 1,
                    "scores": {
                        "loss_1": 0.1,
                    }
                }
            ]
        }
    if num_boxes > 1:
        boxes["box_layer_2_3"] = {
            "box_data": [
                {
                    "position": {"minX": 0.5, "minY": 0.5, "maxX": 0.6, "maxY": 0.6},
                    "class_id": 2,
                    "scores": {
                        "loss_2": 0.1,
                    }
                },
                {
                    "position": {"minX": 0.7, "minY": 0.7, "maxX": 0.8, "maxY": 0.8},
                    "class_id": 3,
                    "scores": {
                        "loss_3": 0.1,
                    }
                }
            ]
        }
    if num_boxes > 2:
        boxes["box_layer_100"] = {
            "box_data": [
                {
                    "position": {"minX": 0.9, "minY": 0.9, "maxX": 1.0, "maxY": 1.0},
                    "class_id": 100,
                    "scores": {
                        "loss_100": 0.1,
                    }
                },
            ]
        }

    if num_masks > 0 and num_boxes > 0:
        return wandb.Image(data, classes=classes, masks=masks, boxes=boxes)
    elif num_masks > 0:
        return wandb.Image(data, classes=classes, masks=masks)
    elif num_boxes > 0:
        return wandb.Image(data, classes=classes, boxes=boxes)
    else:
        return wandb.Image(data)

wandb.init(project="image_classes")
wandb.log({"image": make_image()})
image = make_image()
table = wandb.Table(data=[
        [make_image(num_masks, num_boxes)]
        for num_masks in range(4)
        for num_boxes in range(4)
    ], 
    columns=["image"])
wandb.log({"table": table})
wandb.finish()

Release Notes

Below, please enter user-facing release notes as one or more bullet points.
If your change is not user-visible, write NO RELEASE NOTES instead, with no bullet points.

------------- BEGIN RELEASE NOTES ------------------

• Weave Type System: Supports mixed classes in bounding box and image mask annotation layers

------------- END RELEASE NOTES --------------------