Using a Custom Evaluator (for custom tasks)¶

Even though 3DB focuses on computer vision, it can also be used to solve a myriad of possible tasks. We enable users to describe their own evaluation tasks in 3DB.

Out of the box, the frameworks supports:

Classification (threedb.evaluators.classification)
Object Detection (threedb.evaluators.detection)

The procedure for extending evaluators in 3DB is as follows:

Implement a subclass of the base class threedb.evaluators.base_evaluator.
Have your file export the class under the name Evaluator.
Make sure it can be imported in your current environment (via $PYTHONPATH, pip install -e ..., …).
Update the configuration file to use the newly defined evaluator.

Implementation¶

Below, we provide an example implementation of custom evaluator for an image segmentation task. We’ll briefly outline each of the required steps below—for more detailed documentation of each abstract function, see the threedb.evaluators.base_evaluator module.

We’ll start by importing the required modules and subclassing the BaseEvaluator class. In order for the subclass to be valid, we need to declare two variables, KEYS and output_type. At a high level, the purpose of any evaluator is to map from a (model prediction, label) pair to an arbitrary dictionary summarizing the results—the KEYS variable declares what the keys of this dictionary will be. The output_type variable is only used by the dashboard to visualize the results, and can be any string.

from typing import Dict, List, Tuple, Any
import torch as ch
import json
from threedb.evaluators.base_evaluator import BaseEvaluator

class ImageSegmentation(BaseEvaluator):
    # The output_type variable needs to be defined; it can be any string and
    # is only used by the dashboard to decide how the results are displayed.
    output_type = 'segmentation'

    # The KEYS variable needs to match the keys in declare_outputs
    # and summary_stats.
    KEYS = ['is_correct', 'loss', 'seg_map']

Next, we implement the init function: this can take arbitrary arguments and should set up everything that the evaluator needs to generate metadata about model predictions. For our segmentation evaluator, we’ll need (a) a file mapping model UIDs to classes; (b) an error threshold at which to call a segmentation “correct;” and (c) the size of the segmentation maps produced.

def __init__(self, model_uid_to_class_file, l2_thresh, im_size):
    # In order to implement methods in this class you probably
    # will need some metadata. Feel free to accept any argument
    # you need in your constructor. You will be able to populate
    # them in your config file.

    self.model_to_class = json.load(open(model_uid_to_class_file))
    self.l2_thresh = l2_thresh
    self.im_size = im_size

The next abstract function we have to implement is get_segmentation_label(), which takes in a model uid and should return a label to be used in the segmentation map (i.e., the segmentation map will be -1 wherever the backgroiunds visible, and X wherever the model is visible, where X is the output of this function). Below is a function that uses the provided JSON file to return this label:
def get_segmentation_label(self, model_uid: str) -> int: # The output of this function will be the value associated # to pixels that belongs to the object of interest. label = self.uid_to_targets[model_uid] return label[0] if isinstance(label, list) else label

Next, we implement the declare_outputs() function—this must return a dictionary with keys equal to the KEYS variable declared earlier, and values equal to tuples of the form (shape, type). In particular, shape should be a list describing the shape of the tensor that will be returned, and type should be a PyTorch dtype:

def declare_outputs(self) -> Dict[str, Tuple[List[int], str]]:
    # The goal of this method is to declare what kinds of metrics
    # the evaluator will generate.
    return {
        'is_correct': ([], 'bool'),
        'loss': ([], 'float32'),
        'seg_map': (self.im_size, 'float32')
        # Any other metrics you want to report!
    }

The next step is to declare the get_target function, which takes in (a) the UID of the rendered 3D model and (b) the render_output dictionary consisting of the render output (for the built-in Blender engine, this thankfully already comes with a “segmentation” key!), and returns the desired ground-truth that our segmentation model’s output will be compared to:
def get_target(self, model_uid: str, render_output: Dict[str, Tensor]) -> LabelType: # returns the expected label (whatever label means for this specific task) return render_output['segmentation']

The last and most important step is the summary_stats function, which takes in a model prediction and a label (the output of get_target) and returns a dictionary that has the same keys as KEYS, and values that are PyTorch tensors with the declared shapes and types from declare_outputs. For example, assuming the model outputs a segmentation map, we might return something like the following:

def summary_stats(self, pred: ch.Tensor, label: LabelType, input_shape: List[int]) -> Dict[str, Output]:
    # This method is used to generate the metrics declared in
    # declare_outputs() using the output of to_tensor() and
    # get_target().
    loss = (pred - label).norm()
    is_correct = ch.tensor(loss < self.l2_thresh)
    return {
        'is_correct': is_correct,
        'loss': loss,
        'seg_map': pred
        # You can add as many metrics as you want as long
        # as you match declare_outputs
    }

Finally, don’t forget to export your class under the Evaluator variable!

# IMPORTANT! Needed so that threedb is able to import your custom evaluator
# (since it can't know how you named your class).
Evaluator = ImageSegmentation

That’s it! We’ve implemented all the functions needed to add a custom task to 3DB.

Updating the configuration file¶

You should update the evaluation section of your configuration file:

# ... rest of YAML file
evaluation:
    module: "path.to.your.newly.created.module"
    args:
        model_uid_to_class_file: "/path/to/mapping/file"
        l2_thresh: 10.
        im_size: [1, 224, 224]
render_args:
    engine: 'threedb.rendering.render_blender'
    resolution: 256
    samples: 16
    with_segmentation: true