clip.report_marsclip_embeddings

report_marsclip_embeddings

Generate Stage A embedding sanity artifacts from a trained MarsCLIP MAE.

load_trained_mae_from_checkpoint

load_trained_mae_from_checkpoint(checkpoint_path: Path | str, *, map_location: str | device = 'cpu') -> tuple[torch.nn.Module, dict[str, Any]]

Load a trained Stage A MAE and the associated checkpoint state.

Source code in src/clip/report_marsclip_embeddings.py

def load_trained_mae_from_checkpoint(
    checkpoint_path: pathlib.Path | str,
    *,
    map_location: str | torch.device = "cpu",
) -> tuple[torch.nn.Module, dict[str, Any]]:
    """Load a trained Stage A MAE and the associated checkpoint state."""
    checkpoint = torch.load(checkpoint_path, map_location=resolve_map_location(map_location))
    config = dict(checkpoint.get("config", {}))
    config.update(_infer_mae_config_from_state_dict(checkpoint.get("model_state", {}), config))
    model = build_mae_model_from_config(config)
    state = load_mae_checkpoint(checkpoint_path, model, map_location=map_location)
    state["config"] = config
    return model, state

collect_mae_embeddings

collect_mae_embeddings(model: Module, dataset: Sequence[dict[str, Any]], *, batch_size: int = 4, device: str | device | None = None, mask_ratio: float = 0.0, max_items: int | None = None) -> tuple[torch.Tensor, list[dict[str, Any]], list[dict[str, Any]]]

Collect pooled Stage A embeddings plus serializable metadata records.

Source code in src/clip/report_marsclip_embeddings.py

def collect_mae_embeddings(
    model: torch.nn.Module,
    dataset: Sequence[dict[str, Any]],
    *,
    batch_size: int = 4,
    device: str | torch.device | None = None,
    mask_ratio: float = 0.0,
    max_items: int | None = None,
) -> tuple[torch.Tensor, list[dict[str, Any]], list[dict[str, Any]]]:
    """Collect pooled Stage A embeddings plus serializable metadata records."""
    count = len(dataset) if max_items is None else min(len(dataset), int(max_items))
    if count <= 0:
        raise ValueError("Dataset must contain at least one sample.")

    samples = [dataset[i] for i in range(count)]
    dataloader = build_mae_dataloader(samples, batch_size=batch_size, shuffle=False)
    resolved_device = _resolve_device(device, model)
    model.to(resolved_device)
    model.eval()

    embeddings: list[torch.Tensor] = []
    records: list[dict[str, Any]] = []
    offset = 0
    with torch.no_grad():
        for batch in dataloader:
            batch_size_actual = len(batch["metadata"])
            batch_on_device = _move_batch_to_device(batch, resolved_device)
            output = model.forward_patch_batch(batch_on_device, mask_ratio=mask_ratio)
            embeddings.append(output.pooled_embedding.detach().cpu())
            for idx in range(batch_size_actual):
                records.append(
                    _serialize_embedding_record(
                        samples[offset + idx],
                        index=offset + idx,
                    )
                )
            offset += batch_size_actual

    return torch.cat(embeddings, dim=0), records, samples

compute_topk_neighbors

compute_topk_neighbors(embeddings: Tensor, *, top_k: int = 3) -> tuple[torch.Tensor, torch.Tensor]

Return cosine-nearest neighbor indices and scores for each embedding.

Source code in src/clip/report_marsclip_embeddings.py

def compute_topk_neighbors(
    embeddings: torch.Tensor,
    *,
    top_k: int = 3,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return cosine-nearest neighbor indices and scores for each embedding."""
    if embeddings.ndim != 2:
        raise ValueError("embeddings must have shape (N, D)")
    num_items = embeddings.shape[0]
    if num_items == 0:
        raise ValueError("embeddings must not be empty.")
    if num_items == 1:
        empty = torch.empty((1, 0), dtype=torch.long)
        return empty, empty.to(dtype=torch.float32)

    k = min(int(top_k), num_items - 1)
    normed = torch.nn.functional.normalize(embeddings.to(torch.float32), dim=1)
    similarity = normed @ normed.T
    similarity.fill_diagonal_(-float("inf"))
    scores, indices = torch.topk(similarity, k=k, dim=1)
    return indices.cpu(), scores.cpu()

project_embeddings_pca

project_embeddings_pca(embeddings: Tensor) -> torch.Tensor

Project embeddings to 2D with a lightweight PCA for qualitative inspection.

Source code in src/clip/report_marsclip_embeddings.py

def project_embeddings_pca(embeddings: torch.Tensor) -> torch.Tensor:
    """Project embeddings to 2D with a lightweight PCA for qualitative inspection."""
    if embeddings.ndim != 2:
        raise ValueError("embeddings must have shape (N, D)")
    num_items, dim = embeddings.shape
    if num_items == 0:
        raise ValueError("embeddings must not be empty.")
    if num_items == 1:
        return torch.zeros((1, 2), dtype=torch.float32)

    centered = embeddings.to(torch.float32) - embeddings.to(torch.float32).mean(dim=0, keepdim=True)
    rank = min(2, num_items, dim)
    _, _, right = torch.pca_lowrank(centered, q=rank, center=False)
    projected = centered @ right[:, :rank]
    if rank == 1:
        projected = torch.cat([projected, torch.zeros_like(projected)], dim=1)
    return projected[:, :2].cpu()

save_embedding_neighbor_gallery

save_embedding_neighbor_gallery(samples: Sequence[dict[str, Any]], neighbor_indices: Tensor, neighbor_scores: Tensor, out_path: Path | str, *, num_queries: int = 4) -> pathlib.Path

Save a query-plus-neighbors gallery from embedding similarity results.

Source code in src/clip/report_marsclip_embeddings.py

def save_embedding_neighbor_gallery(
    samples: Sequence[dict[str, Any]],
    neighbor_indices: torch.Tensor,
    neighbor_scores: torch.Tensor,
    out_path: pathlib.Path | str,
    *,
    num_queries: int = 4,
) -> pathlib.Path:
    """Save a query-plus-neighbors gallery from embedding similarity results."""
    if not samples:
        raise ValueError("samples must not be empty.")
    rows = min(int(num_queries), len(samples))
    top_k = int(neighbor_indices.shape[1]) if neighbor_indices.ndim == 2 else 0
    cols = 1 + max(top_k, 0)

    fig, axes = plt.subplots(rows, cols, figsize=(3.6 * cols, 3.8 * rows))
    if rows == 1:
        axes = np.array([axes])
    if cols == 1:
        axes = axes.reshape(rows, 1)

    for row_idx in range(rows):
        query = samples[row_idx]
        q_md = query.get("metadata", {})
        ax = axes[row_idx, 0]
        ax.imshow(_to_display_rgb(query["image"]), interpolation="nearest")
        ax.axis("off")
        ax.set_title(
            f"query\n{q_md.get('obs_id', 'unknown')}  valid={q_md.get('overall_valid_fraction', 0.0):.0%}",
            fontsize=9,
        )
        for col_idx in range(top_k):
            neighbor_idx = int(neighbor_indices[row_idx, col_idx])
            neighbor = samples[neighbor_idx]
            n_md = neighbor.get("metadata", {})
            ax = axes[row_idx, col_idx + 1]
            ax.imshow(_to_display_rgb(neighbor["image"]), interpolation="nearest")
            ax.axis("off")
            ax.set_title(
                (
                    f"nn{col_idx + 1}  sim={float(neighbor_scores[row_idx, col_idx]):.3f}\n"
                    f"{n_md.get('obs_id', 'unknown')}"
                ),
                fontsize=9,
            )

    fig.suptitle("MarsCLIP Stage A nearest-neighbor gallery", fontsize=12)
    fig.tight_layout()
    out = pathlib.Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out, dpi=160)
    plt.close(fig)
    return out

save_embedding_scatter

save_embedding_scatter(projection: Tensor, records: Sequence[dict[str, Any]], out_path: Path | str) -> pathlib.Path

Save a simple 2D embedding scatter colored by valid-pixel fraction.

Source code in src/clip/report_marsclip_embeddings.py

def save_embedding_scatter(
    projection: torch.Tensor,
    records: Sequence[dict[str, Any]],
    out_path: pathlib.Path | str,
) -> pathlib.Path:
    """Save a simple 2D embedding scatter colored by valid-pixel fraction."""
    if projection.ndim != 2 or projection.shape[1] != 2:
        raise ValueError("projection must have shape (N, 2)")
    if len(records) != projection.shape[0]:
        raise ValueError("records length must match projection rows.")

    colors = [float(record.get("overall_valid_fraction", 0.0)) for record in records]
    fig, ax = plt.subplots(figsize=(6.5, 5.0))
    scatter = ax.scatter(
        projection[:, 0].numpy(),
        projection[:, 1].numpy(),
        c=colors,
        cmap="viridis",
        s=45,
        alpha=0.9,
        edgecolors="black",
        linewidths=0.2,
    )
    for idx in range(min(8, len(records))):
        record = records[idx]
        label = record.get("patch_id") or record.get("obs_id") or f"item_{idx}"
        ax.text(
            float(projection[idx, 0]),
            float(projection[idx, 1]),
            str(label),
            fontsize=7,
            alpha=0.7,
        )
    ax.set_title("MarsCLIP Stage A embedding PCA")
    ax.set_xlabel("PC 1")
    ax.set_ylabel("PC 2")
    ax.grid(alpha=0.25)
    fig.colorbar(scatter, ax=ax, label="valid pixel fraction")
    fig.tight_layout()
    out = pathlib.Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out, dpi=160)
    plt.close(fig)
    return out

save_embedding_report

save_embedding_report(model: Module, dataset: Sequence[dict[str, Any]], out_dir: Path | str, *, checkpoint_path: Path | str | None = None, config: dict[str, Any] | None = None, batch_size: int = 4, top_k: int = 3, num_queries: int = 4, device: str | device | None = None, mask_ratio: float = 0.0, max_items: int | None = None) -> dict[str, Any]

Save Stage A embedding tensors, metadata, gallery, scatter, and summary.

Source code in src/clip/report_marsclip_embeddings.py

def save_embedding_report(
    model: torch.nn.Module,
    dataset: Sequence[dict[str, Any]],
    out_dir: pathlib.Path | str,
    *,
    checkpoint_path: pathlib.Path | str | None = None,
    config: dict[str, Any] | None = None,
    batch_size: int = 4,
    top_k: int = 3,
    num_queries: int = 4,
    device: str | torch.device | None = None,
    mask_ratio: float = 0.0,
    max_items: int | None = None,
) -> dict[str, Any]:
    """Save Stage A embedding tensors, metadata, gallery, scatter, and summary."""
    out_path = pathlib.Path(out_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    embeddings, records, samples = collect_mae_embeddings(
        model,
        dataset,
        batch_size=batch_size,
        device=device,
        mask_ratio=mask_ratio,
        max_items=max_items,
    )
    neighbor_indices, neighbor_scores = compute_topk_neighbors(embeddings, top_k=top_k)
    projection = project_embeddings_pca(embeddings)

    embeddings_path = out_path / "embeddings.pt"
    torch.save({"embeddings": embeddings, "records": records}, embeddings_path)
    metadata_path = out_path / "embedding_metadata.json"
    metadata_path.write_text(json.dumps(records, indent=2))

    gallery_path = save_embedding_neighbor_gallery(
        samples,
        neighbor_indices,
        neighbor_scores,
        out_path / "nearest_neighbors.png",
        num_queries=num_queries,
    )
    scatter_path = save_embedding_scatter(projection, records, out_path / "embedding_scatter.png")

    first_neighbor_mean = None
    if neighbor_scores.numel() > 0:
        first_neighbor_mean = float(neighbor_scores[:, 0].mean().item())
    summary = {
        "checkpoint": str(checkpoint_path) if checkpoint_path is not None else None,
        "num_embeddings": int(embeddings.shape[0]),
        "embedding_dim": int(embeddings.shape[1]),
        "batch_size": int(batch_size),
        "mask_ratio": float(mask_ratio),
        "top_k": int(min(top_k, max(0, embeddings.shape[0] - 1))),
        "num_queries": int(min(num_queries, embeddings.shape[0])),
        "mean_embedding_norm": float(embeddings.norm(dim=1).mean().item()),
        "mean_first_neighbor_similarity": first_neighbor_mean,
        "embeddings_path": str(embeddings_path),
        "metadata_path": str(metadata_path),
        "gallery_path": str(gallery_path),
        "scatter_path": str(scatter_path),
        "config": dict(config or {}),
    }
    summary_path = out_path / "embedding_summary.json"
    summary_path.write_text(json.dumps(summary, indent=2))
    summary["summary_path"] = str(summary_path)
    return summary