def project_bbox()

in phosa/global_opt.py [0:0]

• 23 lines of code
• 5 McCabe index (conditional complexity)

def project_bbox(vertices, renderer, parts_labels=None, bbox_expansion=0.0):
    """
    Computes the 2D bounding box of the vertices after projected to the image plane.

    TODO(@jason): Batch these operations.

    Args:
        vertices (V x 3).
        renderer: Renderer used to get camera parameters.
        parts_labels (dict): Dictionary mapping a part name to the corresponding vertex
            indices.
        bbox_expansion (float): Amount to expand the bounding boxes.

    Returns:
        If a part_label dict is given, returns a dictionary mapping part name to bbox.
        Else, returns the projected 2D bounding box.
    """
    proj = nr.projection(
        (vertices * torch.tensor([[1, -1, 1.0]]).cuda()).unsqueeze(0),
        K=renderer.K,
        R=renderer.R,
        t=renderer.t,
        dist_coeffs=renderer.dist_coeffs,
        orig_size=1,
    )
    proj = proj.squeeze(0)[:, :2]
    if parts_labels is None:
        parts_labels = {"": torch.arange(len(vertices)).to(vertices.device)}
    bbox_parts = {}
    for part, inds in parts_labels.items():
        bbox = torch.cat((proj[inds].min(0).values, proj[inds].max(0).values), dim=0)
        if bbox_expansion:
            center = (bbox[:2] + bbox[2:]) / 2
            b = (bbox[2:] - bbox[:2]) / 2 * (1 + bbox_expansion)
            bbox = torch.cat((center - b, center + b))
        bbox_parts[part] = bbox
    if "" in parts_labels:
        return bbox_parts[""]
    return bbox_parts