File size: 737 Bytes
e8b8483 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | """Reading the pooled feature caches that cache.py writes."""
import json
from pathlib import Path
from typing import Tuple
import numpy as np
import torch
from .data import coco_split, person_labels
def load_pooled(cache, split: str) -> Tuple[torch.Tensor, torch.Tensor]:
"""(N, 768) pooled vectors and their person labels, in cache row order."""
cache = Path(cache)
pooled = np.load(cache / 'pooled.npy').astype(np.float32)
img_ids = json.loads((cache / 'img_ids.json').read_text())
if pooled.shape[0] != len(img_ids):
raise ValueError(f'{cache} has {pooled.shape[0]} rows and {len(img_ids)} ids')
coco, _ = coco_split(split)
return torch.from_numpy(pooled), person_labels(coco, img_ids, 'cpu')
|