Spaces:
Sleeping
Sleeping
File size: 3,564 Bytes
0ba38d0 | 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | # AI-assisted (Claude Code, claude.ai) -- https://claude.ai
"""Extract ResNet50 image embeddings for articles in the subset."""
import torch
import numpy as np
import pandas as pd
from torchvision import models, transforms
from PIL import Image
from pathlib import Path
def extract_features(
data_dir: str = "data",
batch_size: int = 32,
full_catalog: bool = True,
):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device.type == "cuda":
torch.cuda.set_per_process_memory_fraction(0.5, 0)
print(f"GPU memory capped at 50%")
print(f"Using device: {device}")
# Load pretrained ResNet50, remove classification head
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model = torch.nn.Sequential(*list(model.children())[:-1]) # -> 2048-dim
model = model.to(device)
model.eval()
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
])
data_dir = Path(data_dir)
if full_catalog:
articles = pd.read_csv(data_dir / "raw" / "articles.csv", dtype={"article_id": str})
print(f"Extracting for FULL catalog: {len(articles):,} articles")
else:
articles = pd.read_csv(data_dir / "processed" / "articles_subset.csv", dtype={"article_id": str})
print(f"Extracting for subset: {len(articles):,} articles")
image_dir = data_dir / "raw" / "images"
article_ids = articles["article_id"].values
embeddings = np.zeros((len(article_ids), 2048), dtype=np.float32)
missing = []
# Process in batches
for start in range(0, len(article_ids), batch_size):
end = min(start + batch_size, len(article_ids))
batch_ids = article_ids[start:end]
batch_tensors = []
batch_indices = []
for i, aid in enumerate(batch_ids):
img_path = image_dir / aid[:3] / f"{aid}.jpg"
if not img_path.exists():
missing.append(aid)
continue
try:
img = Image.open(img_path).convert("RGB")
batch_tensors.append(preprocess(img))
batch_indices.append(start + i)
except Exception as e:
print(f" Error loading {aid}: {e}")
missing.append(aid)
if batch_tensors:
batch = torch.stack(batch_tensors).to(device)
with torch.no_grad():
features = model(batch).squeeze(-1).squeeze(-1).cpu().numpy()
for j, idx in enumerate(batch_indices):
embeddings[idx] = features[j]
if (start // batch_size) % 10 == 0:
print(f" Processed {end}/{len(article_ids)} articles...")
# Save
out_dir = data_dir / "processed"
np.save(out_dir / "image_embeddings.npy", embeddings)
# Save article_id -> index mapping
id_map = pd.DataFrame({"article_id": article_ids, "embedding_idx": range(len(article_ids))})
id_map.to_csv(out_dir / "article_embedding_map.csv", index=False)
print(f"\nDone! Saved {len(article_ids)} embeddings to {out_dir / 'image_embeddings.npy'}")
print(f"Shape: {embeddings.shape}")
if missing:
print(f"Missing images: {len(missing)} articles")
return embeddings, article_ids
if __name__ == "__main__":
import sys
full = "--subset-only" not in sys.argv
extract_features(full_catalog=full)
|