File size: 2,451 Bytes
1ea7ba6 | 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 | """Extract backbone features on test set, project to 2D with t-SNE."""
import sys, os
sys.path.insert(0, "/mnt/d/SpiceNet" if os.path.exists("/mnt/d/SpiceNet") else "D:/SpiceNet")
import numpy as np
import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import config
from src.dataset import get_dataloaders
from src.model import load_checkpoint
@torch.no_grad()
def extract_features(model, loader, device):
model.eval()
feats, labels = [], []
for imgs, tex, col, lbl in loader:
imgs = imgs.to(device)
f = model.backbone(imgs) # (B, 1792) pre-classifier features
feats.append(f.cpu().numpy())
labels.extend(lbl.tolist())
return np.concatenate(feats), np.array(labels)
def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model, *_ = load_checkpoint(str(config.CHECKPOINT_DIR / "best.pth"), device)
_, _, test_loader, _, _ = get_dataloaders(multimodal=True)
print(f"Extracting features on {len(test_loader.dataset)} test samples...")
feats, labels = extract_features(model, test_loader, device)
print(f"Features: {feats.shape}, labels: {labels.shape}")
print("Running t-SNE (perplexity=30, ~30s)...")
proj = TSNE(n_components=2, perplexity=30, init="pca",
learning_rate="auto", random_state=config.RANDOM_SEED).fit_transform(feats)
fig, ax = plt.subplots(figsize=(11, 9))
cmap = plt.get_cmap("tab20")
for i, cls in enumerate(config.CLASSES):
mask = labels == i
ax.scatter(proj[mask, 0], proj[mask, 1], s=8, c=[cmap(i)], label=cls, alpha=0.7)
ax.set_xlabel("t-SNE 1"); ax.set_ylabel("t-SNE 2")
ax.set_title("EfficientNet-B4 backbone features (test set) — strong-aug model")
ax.legend(loc="best", fontsize=9, markerscale=1.8, framealpha=0.85)
ax.grid(alpha=0.2)
out_path = config.OUTPUT_DIR / "feature_tsne.png"
plt.tight_layout()
plt.savefig(out_path, dpi=150)
plt.close()
print(f"Saved -> {out_path}")
# Per-class centroid distances (just for the analysis table)
centroids = np.stack([proj[labels == i].mean(0) for i in range(len(config.CLASSES))])
np.savez(config.OUTPUT_DIR / "feature_tsne_data.npz",
proj=proj, labels=labels, centroids=centroids)
if __name__ == "__main__":
main()
|