visualize_features.py

"""
Visualize the 3680-dim feature space using PCA + t-SNE.
Extract features from a subset of training images, build pair features,
reduce to 2D, color by predicate label.
"""

import os, json, random, time
import numpy as np
import torch
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from PIL import Image
from transformers import DetrForObjectDetection, DetrImageProcessor

random.seed(2026); np.random.seed(2026); torch.manual_seed(2026)

DATA = 'data_final'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

with open(f'{DATA}/train.json') as f: train_ann = json.load(f)
with open(f'{DATA}/predicates.json') as f: pred_info = json.load(f)
predicates = pred_info['predicates']

# Use 500 images for speed
sample = random.sample(train_ann, min(500, len(train_ann)))
print(f"Using {len(sample)} train images")

# Load DETR
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
detr = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50").to(device)
detr.eval()

def compute_iou(box1, box2):
    x1, y1 = max(box1[0], box2[0]), max(box1[1], box2[1])
    x2, y2 = min(box1[2], box2[2]), min(box1[3], box2[3])
    inter = max(0, x2-x1) * max(0, y2-y1)
    a1 = max(0, (box1[2]-box1[0]) * (box1[3]-box1[1]))
    a2 = max(0, (box2[2]-box2[0]) * (box2[3]-box2[1]))
    return inter / (a1 + a2 - inter + 1e-6)

# Extract pair features with labels
all_feats = []
all_labels = []
all_pred_names = []

t0 = time.time()
for idx, ann in enumerate(sample):
    img_path = os.path.join(DATA, 'images', 'train', ann['file_name'])
    if not os.path.exists(img_path):
        continue
    try:
        image = Image.open(img_path).convert('RGB')
    except:
        continue
    
    inputs = processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = detr(**inputs, output_attentions=True, output_hidden_states=True)
    
    target_sizes = torch.tensor([image.size[::-1]]).to(device)
    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.5)[0]
    det_boxes = results['boxes'].cpu()
    n_det = len(det_boxes)
    if n_det < 2:
        continue
    
    logits = outputs.logits[0].cpu()
    probs_all = logits.softmax(-1)[..., :-1].max(-1).values
    keep = logits.argmax(-1) != detr.config.num_labels
    scores_all = probs_all.clone(); scores_all[~keep] = -1
    _, top_idx = scores_all.topk(min(n_det, 100))
    top_idx = top_idx[:n_det].sort().values
    
    decoder_hidden = [h[0, top_idx].cpu() for h in outputs.decoder_hidden_states]
    self_attn = [a[0, :, top_idx][:, :, top_idx].cpu() for a in outputs.decoder_attentions]
    
    # Match relationships and build features
    for rel in ann.get('relationships', []):
        p_idx = rel['predicate']
        best_s, best_si = -1, 0.3
        best_o, best_oi = -1, 0.3
        for d in range(n_det):
            si = compute_iou(det_boxes[d].tolist(), rel['subject_box'])
            oi = compute_iou(det_boxes[d].tolist(), rel['object_box'])
            if si > best_si: best_s, best_si = d, si
            if oi > best_oi: best_o, best_oi = d, oi
        
        if best_s >= 0 and best_o >= 0 and best_s != best_o:
            parts = []
            for h in decoder_hidden:
                hf = h.float()
                parts.extend([hf[best_s], hf[best_o]])
            for a in self_attn:
                af = a.float()
                parts.extend([af[:, best_s, best_o], af[:, best_o, best_s]])
            
            feat = torch.cat(parts).numpy()
            all_feats.append(feat)
            all_labels.append(p_idx)
            all_pred_names.append(rel['predicate_name'])
    
    if (idx+1) % 100 == 0:
        print(f"  {idx+1}/{len(sample)} ({len(all_feats)} pairs)")

print(f"Extracted {len(all_feats)} pair features in {time.time()-t0:.0f}s")

X = np.array(all_feats)
y = np.array(all_labels)
print(f"Feature matrix: {X.shape}")

# ============================================================
# PCA to 50 dims, then t-SNE to 2D
# ============================================================
print("PCA 3680 -> 50...")
pca = PCA(n_components=50, random_state=2026)
X_pca = pca.fit_transform(X)
print(f"PCA variance explained: {pca.explained_variance_ratio_.sum()*100:.1f}%")

print("t-SNE 50 -> 2...")
tsne = TSNE(n_components=2, random_state=2026, perplexity=30, max_iter=1000)
X_2d = tsne.fit_transform(X_pca)
print("t-SNE done")

# ============================================================
# Plot: Top 15 predicates colored, rest in gray
# ============================================================
from collections import Counter
label_counts = Counter(all_pred_names)
top15 = [p for p, _ in label_counts.most_common(15)]

# Color palette
colors_palette = [
    '#e6194B', '#3cb44b', '#ffe119', '#4363d8', '#f58231',
    '#911eb4', '#42d4f4', '#f032e6', '#bfef45', '#fabed4',
    '#469990', '#dcbeff', '#9A6324', '#800000', '#aaffc3',
]

fig, axes = plt.subplots(1, 2, figsize=(24, 10))

# Plot 1: All predicates, top 15 colored
ax = axes[0]
# First plot gray background (non-top-15)
mask_other = np.array([p not in top15 for p in all_pred_names])
if mask_other.any():
    ax.scatter(X_2d[mask_other, 0], X_2d[mask_other, 1], c='#cccccc', s=8, alpha=0.3, label='other')

for i, pred in enumerate(top15):
    mask = np.array([p == pred for p in all_pred_names])
    if mask.any():
        ax.scatter(X_2d[mask, 0], X_2d[mask, 1], c=colors_palette[i], s=15, alpha=0.6, label=f'{pred} ({mask.sum()})')

ax.legend(fontsize=7, loc='upper right', markerscale=2)
ax.set_title('t-SNE of DETR Pair Features (3680d -> PCA 50d -> t-SNE 2d)\nColored by top-15 predicates', fontsize=12)
ax.set_xlabel('t-SNE 1')
ax.set_ylabel('t-SNE 2')

# Plot 2: Grouped into categories
ax2 = axes[1]
category_map = {
    'Spatial': ['on', 'in', 'on top of', 'under', 'above', 'behind', 'near', 'next to', 'in front of', 'beside', 'between', 'along', 'over'],
    'Possessive': ['has', 'of', 'with', 'has a', 'of a', 'for'],
    'Action': ['holding', 'riding', 'eating', 'playing', 'carrying', 'walking on', 'sitting on', 'standing on', 'laying on', 'sitting in', 'standing in', 'hanging on', 'hanging from', 'covering'],
    'Wearing': ['wearing', 'wears', 'wearing a'],
}

cat_colors = {'Spatial': '#4363d8', 'Possessive': '#3cb44b', 'Action': '#e6194B', 'Wearing': '#f58231'}

def get_category(pred_name):
    for cat, preds in category_map.items():
        if pred_name in preds:
            return cat
    return 'Other'

cats = [get_category(p) for p in all_pred_names]
for cat, color in cat_colors.items():
    mask = np.array([c == cat for c in cats])
    if mask.any():
        ax2.scatter(X_2d[mask, 0], X_2d[mask, 1], c=color, s=15, alpha=0.5, label=f'{cat} ({mask.sum()})')

mask_oth = np.array([c == 'Other' for c in cats])
if mask_oth.any():
    ax2.scatter(X_2d[mask_oth, 0], X_2d[mask_oth, 1], c='#cccccc', s=8, alpha=0.3, label=f'Other ({mask_oth.sum()})')

ax2.legend(fontsize=10, markerscale=2)
ax2.set_title('Same t-SNE, colored by relationship category\n(Spatial / Possessive / Action / Wearing)', fontsize=12)
ax2.set_xlabel('t-SNE 1')
ax2.set_ylabel('t-SNE 2')

plt.tight_layout()
plt.savefig('tsne_features.png', dpi=150, bbox_inches='tight')
print("Saved tsne_features.png")
plt.close()