Guanc27/check_fonts

1"""2Visualize font embeddings from Phase 3 using PCA/TSNE/UMAP.3"""4 5import argparse6import json7from pathlib import Path8 9import numpy as np10 11try:12    import matplotlib.pyplot as plt13except ImportError as exc:14    raise ImportError(15        "matplotlib is required. Install it with: pip install matplotlib"16    ) from exc17 18try:19    from sklearn.decomposition import PCA20    from sklearn.manifold import TSNE21except ImportError as exc:22    raise ImportError(23        "scikit-learn is required. Install it with: pip install scikit-learn"24    ) from exc25 26try:27    import umap  # type: ignore28    HAS_UMAP = True29except ImportError:30    HAS_UMAP = False31 32try:33    import faiss  # type: ignore34    HAS_FAISS = True35except ImportError:36    HAS_FAISS = False37 38 39def load_metadata(metadata_path):40    with open(metadata_path, "r", encoding="utf-8") as f:41        return json.load(f)42 43 44def load_embeddings(vector_db_dir):45    embeddings_path = Path(vector_db_dir) / "embeddings.npy"46    if embeddings_path.exists():47        return np.load(embeddings_path)48 49    index_path = Path(vector_db_dir) / "faiss.index"50    if not index_path.exists():51        raise FileNotFoundError(52            "No embeddings.npy or faiss.index found. "53            "Run Phase 3 with --save_embeddings or ensure faiss.index exists."54        )55 56    if not HAS_FAISS:57        raise ImportError(58            "faiss-cpu is required to reconstruct embeddings from faiss.index. "59            "Install it with: pip install faiss-cpu"60        )61 62    index = faiss.read_index(str(index_path))63    if not hasattr(index, "reconstruct"):64        raise RuntimeError("FAISS index does not support vector reconstruction.")65 66    vectors = np.zeros((index.ntotal, index.d), dtype="float32")67    for i in range(index.ntotal):68        vectors[i] = index.reconstruct(i)69    return vectors70 71 72def choose_method(method):73    if method == "umap":74        if not HAS_UMAP:75            raise ImportError(76                "umap-learn is required for --method umap. "77                "Install it with: pip install umap-learn"78            )79        return "umap"80    if method in {"pca", "tsne"}:81        return method82    raise ValueError("method must be one of: pca, tsne, umap")83 84 85def reduce_embeddings(embeddings, method, seed):86    if method == "pca":87        reducer = PCA(n_components=2, random_state=seed)88        return reducer.fit_transform(embeddings)89    if method == "tsne":90        reducer = TSNE(n_components=2, random_state=seed, init="pca")91        return reducer.fit_transform(embeddings)92    reducer = umap.UMAP(n_components=2, random_state=seed)93    return reducer.fit_transform(embeddings)94 95 96def main():97    parser = argparse.ArgumentParser(description="Visualize font embeddings")98    parser.add_argument("--vector_db", type=str, default="vector_db",99                        help="Path to vector_db directory")100    parser.add_argument("--method", type=str, default="umap",101                        help="Dimensionality reduction: pca, tsne, umap")102    parser.add_argument("--output", type=str, default="embedding_plot.png",103                        help="Output image path")104    parser.add_argument("--show", action="store_true",105                        help="Display plot interactively")106    parser.add_argument("--max_points", type=int, default=2000,107                        help="Max points to plot (randomly sampled)")108    parser.add_argument("--legend_max", type=int, default=12,109                        help="Show legend only if font count <= legend_max")110    parser.add_argument("--seed", type=int, default=42,111                        help="Random seed for sampling and reducers")112    args = parser.parse_args()113 114    vector_db_dir = Path(args.vector_db)115    metadata_path = vector_db_dir / "metadata.json"116    if not metadata_path.exists():117        raise FileNotFoundError(f"metadata.json not found: {metadata_path}")118 119    metadata = load_metadata(metadata_path)120    samples = metadata.get("samples", [])121 122    embeddings = load_embeddings(vector_db_dir)123    if len(samples) != embeddings.shape[0]:124        raise RuntimeError(125            f"Metadata samples ({len(samples)}) do not match embeddings "126            f"({embeddings.shape[0]}). Re-run Phase 3 with --save_embeddings."127        )128 129    rng = np.random.default_rng(args.seed)130    total = embeddings.shape[0]131    if total > args.max_points:132        idx = rng.choice(total, size=args.max_points, replace=False)133        embeddings = embeddings[idx]134        samples = [samples[i] for i in idx]135 136    method = choose_method(args.method)137    reduced = reduce_embeddings(embeddings, method, args.seed)138 139    font_names = [s["font_name"] for s in samples]140    unique_fonts = sorted(set(font_names))141    font_to_color = {name: i for i, name in enumerate(unique_fonts)}142    colors = [font_to_color[name] for name in font_names]143 144    plt.figure(figsize=(10, 8))145    scatter = plt.scatter(reduced[:, 0], reduced[:, 1], c=colors, cmap="tab20", s=12, alpha=0.8)146    plt.title(f"Font Embeddings ({method.upper()})")147    plt.xlabel("Component 1")148    plt.ylabel("Component 2")149 150    if len(unique_fonts) <= args.legend_max:151        handles, _ = scatter.legend_elements(num=len(unique_fonts))152        plt.legend(handles, unique_fonts, title="Fonts", loc="best", fontsize=8)153 154    plt.tight_layout()155    output_path = Path(args.output)156    output_path.parent.mkdir(parents=True, exist_ok=True)157    plt.savefig(output_path, dpi=200)158    print(f"Saved plot: {output_path}")159 160    if args.show:161        plt.show()162 163 164if __name__ == "__main__":165    main()166