viz.py 文件源码

def draw_one(imgarr, z_arr):
    # Note that this clipping makes the visualisation somewhat
    # misleading, as it incorrectly suggests objects occlude one
    # another.
    clipped = np.clip(imgarr.data.cpu().numpy(), 0, 1)
    img = arr2img(clipped).convert('RGB')
    draw = ImageDraw.Draw(img)
    for k, z in enumerate(z_arr):
        # It would be better to use z_pres to change the opacity of
        # the bounding boxes, but I couldn't make that work with PIL.
        # Instead this darkens the color, and skips boxes altogether
        # when z_pres==0.
        if z.pres > 0:
            (x, y), w, h = bounding_box(z, imgarr.size(0))
            color = tuple(map(lambda c: int(c * z.pres), colors(k)))
            draw.rectangle([x, y, x + w, y + h], outline=color)
    is_relaxed = any(z.pres != math.floor(z.pres) for z in z_arr)
    fmtstr = '{:.1f}' if is_relaxed else '{:.0f}'
    draw.text((0, 0), fmtstr.format(sum(z.pres for z in z_arr)), fill='white')
    return img2arr(img)