Function draw_png

experiments/human_eval/draw_interest_drift_heatmap.py:119–163

()

def draw_png() -> Image.Image:
    image = Image.new("RGB", (W, H), "white")
    draw = ImageDraw.Draw(image)

    for r, method in enumerate(METHODS):
        y = TOP + r * CELL_H
        method_w, method_h = text_size(draw, method, F_ROW)
        draw.text(
            (LEFT - 22 - method_w, y + (CELL_H - GAP - method_h) / 2 - 1),
            method,
            font=F_ROW,
            fill=TEXT,
        )

        for c, (_metric, direction, values) in enumerate(METRICS):
            x = LEFT + c * CELL_W
            score = score_for_metric(values, r, direction)
            draw.rectangle(
                (x, y, x + CELL_W - GAP, y + CELL_H - GAP),
                fill=color_for_score(score),
            )
            centered_text(draw, x + CELL_W / 2, y + CELL_H / 2, fmt(values[r]), F_CELL, (0, 0, 0))

    label_y = TOP + len(METHODS) * CELL_H + 49
    for c, (metric, direction, _values) in enumerate(METRICS):
        x = LEFT + c * CELL_W + CELL_W / 2
        draw_rotated_text(
            image,
            metric + (" \u2193" if direction == "down" else " \u2191"),
            x,
            label_y,
            -21,
            F_COL,
            TEXT,
        )

    bar_x = LEFT + len(METRICS) * CELL_W + 32
    bar_y = TOP + 4
    bar_h = len(METHODS) * CELL_H - GAP
    for i in range(bar_h):
        score = 1 - i / max(1, bar_h - 1)
        draw.line((bar_x, bar_y + i, bar_x + 14, bar_y + i), fill=color_for_score(score))
    draw.text((bar_x + 23, bar_y - 6), "better", font=F_SMALL, fill=TEXT)
    draw.text((bar_x + 23, bar_y + bar_h - 14), "worse", font=F_SMALL, fill=TEXT)
    return image


def rgb(c: tuple[int, int, int]) -> str: