common/cached.py

"""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')