common/pools.py

"""Named pools, cited by name in every artifact's provenance block.

Dims are selected on TRAIN2017 and reported on VAL5000. The two never overlap,
so no figure is measured on data that chose it.
"""
from typing import NamedTuple, Optional


class Pool(NamedTuple):
    name: str
    split: str
    n: Optional[int]
    selection: str


TRAIN2017 = Pool(
    'TRAIN2017', 'train2017', None,
    'all 118287 train2017 images, used only to select dims')

VAL5000 = Pool(
    'VAL5000', 'val2017', None,
    'all 5000 val2017 images, used only to report')

POOLS = {p.name: p for p in (TRAIN2017, VAL5000)}


def by_name(name: str) -> Pool:
    if name not in POOLS:
        raise ValueError(f'unknown pool {name!r}; expected one of {sorted(POOLS)}')
    return POOLS[name]