def add_annotation_bound()

in mujoco_worldgen/util/obj_util.py [0:0]

• 30 lines of code
• 4 McCabe index (conditional complexity)

def add_annotation_bound(xml_dict, annotation_name, bound):
    '''
    Add an annotation bounding box to and XML dictionary.
    Annotation name will be "annotation:" + annotation_name
    Bound is given as a 2 x 3 np.ndarray, and represents:
        [[min_x, min_y, min_z], [max_x, max_y, max_z]]
    '''
    if bound is None:
        return xml_dict  # Nothing to do here
    assert bound.shape == (2, 3), "Bound must be 2 x 3 (see docstring)."
    assert 'worldbody' in xml_dict, "XML must have worldbody"
    worldbody = xml_dict['worldbody']
    assert 'body' in worldbody, "XML worldbody must have bodies"
    name = 'annotation:' + annotation_name
    # Remove old annotations with the same name before inserting new annotation
    bodies = [body for body in worldbody['body'] if body.get('@name') != name]
    rgba = np.random.uniform(size=4)
    rgba[3] = 0.1  # annotation is almost transparent.
    geom = OrderedDict([('@conaffinity', 0),
                        ('@contype', 0),
                        ('@mass', 0.0),
                        ('@pos', np.zeros(3)),
                        ('@rgba', rgba),
                        ('@size', (bound[1] - bound[0]) / 2),  # halfsize
                        ('@type', 'box')])
    annotation = OrderedDict([('@name', name),
                              ('@pos', bound.mean(axis=0)),  # center pos
                              ('geom', [geom])])
    bodies.append(annotation)
    worldbody['body'] = bodies
    print('adding annotation bound (size)', name, bound[1] - bound[0])