Spaces:
Running on Zero
Running on Zero
File size: 3,300 Bytes
f66bbd0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """visualize.py — Plotting / overlay helpers.
``save_prediction`` reproduces the 4-panel (image / gaze / GT / prediction)
grid used during batch evaluation. ``overlay_heatmap`` and ``overlay_mask``
are lighter-weight, matplotlib-free PIL helpers meant for the Gradio demo
and the example notebook, where you typically want a single composited
image rather than a saved subplot file.
"""
from __future__ import annotations
import os
import numpy as np
import torch
from PIL import Image
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from .constants import IMG_MEAN, IMG_STD
def _denormalize(img: torch.Tensor) -> np.ndarray:
"""(3, H, W) normalized tensor -> (H, W, 3) uint8-range float array in [0, 1]."""
mean = torch.tensor(IMG_MEAN, device=img.device).view(3, 1, 1)
std = torch.tensor(IMG_STD, device=img.device).view(3, 1, 1)
return (img * std + mean).clamp(0, 1).permute(1, 2, 0).cpu().numpy()
@torch.no_grad()
def save_prediction(
img: torch.Tensor,
gaze_heatmap: torch.Tensor,
gt: torch.Tensor,
pred: torch.Tensor,
name: str,
out_dir: str = "predictions",
) -> None:
"""Save a 4-panel [image | gaze heatmap | GT mask | predicted mask] figure."""
os.makedirs(out_dir, exist_ok=True)
img_np = _denormalize(img)
gaze_np = gaze_heatmap.squeeze().cpu().numpy()
gt_np = gt.squeeze().cpu().numpy()
pred_np = pred.squeeze().cpu().numpy()
fig, ax = plt.subplots(1, 4, figsize=(20, 5))
ax[0].imshow(img_np); ax[0].set_title("Image")
ax[1].imshow(gaze_np, cmap="jet"); ax[1].set_title("Gaze Heatmap")
ax[2].imshow(gt_np, cmap="gray"); ax[2].set_title("GT mask")
ax[3].imshow(pred_np, cmap="gray"); ax[3].set_title("Predicted mask")
for a in ax:
a.axis("off")
plt.tight_layout()
plt.savefig(os.path.join(out_dir, name), dpi=100, bbox_inches="tight")
plt.close()
def overlay_heatmap(image: Image.Image, heatmap: np.ndarray, alpha: float = 0.45) -> Image.Image:
"""Composite a [0, 1] heatmap (any H'xW', will be resized) onto a PIL image
using a jet colormap. Used to show the gaze prior over the original image."""
heatmap = np.asarray(heatmap, dtype=np.float32)
heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
heat_img = Image.fromarray((cm.jet(heatmap)[:, :, :3] * 255).astype(np.uint8))
heat_img = heat_img.resize(image.size, resample=Image.BILINEAR)
base = image.convert("RGB")
return Image.blend(base, heat_img, alpha=alpha)
def overlay_mask(
image: Image.Image, mask: np.ndarray, color: tuple[int, int, int] = (255, 60, 60), alpha: float = 0.45
) -> Image.Image:
"""Composite a binary/soft [0, 1] mask onto a PIL image as a solid color wash."""
mask = np.asarray(mask, dtype=np.float32)
mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
mask_img = Image.fromarray((mask * 255).astype(np.uint8)).resize(image.size, resample=Image.NEAREST)
base = image.convert("RGB")
color_layer = Image.new("RGB", base.size, color)
composited = Image.composite(color_layer, base, mask_img.point(lambda p: int(p * alpha)))
return composited
|