File size: 2,882 Bytes
e8b8483 | 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 | """COCO loading and input normalization.
Images are resized to a square `resolution` with bilinear interpolation and
normalized with ImageNet statistics, matching the protocol every stage was
measured under.
"""
from pathlib import Path
from typing import Iterable, List, Sequence, Tuple, Union
import numpy as np
import torch
from PIL import Image
from .paths import COCO_ROOT
MEAN = (0.485, 0.456, 0.406)
STD = (0.229, 0.224, 0.225)
PERSON_CATEGORY_ID = 1
def _stats(device: str) -> Tuple[torch.Tensor, torch.Tensor]:
mean = torch.tensor(MEAN).view(1, 3, 1, 1).to(device)
std = torch.tensor(STD).view(1, 3, 1, 1).to(device)
return mean, std
def normalize(img: Image.Image, resolution: int, device: str) -> torch.Tensor:
"""PIL image -> (1, 3, R, R) normalized float tensor."""
img = img.convert('RGB').resize((resolution, resolution), Image.BILINEAR)
arr = np.asarray(img, dtype=np.uint8).copy()
x = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(device).float() / 255.0
mean, std = _stats(device)
return (x - mean) / std
def load_image(image: Union[str, Path, Image.Image, np.ndarray, torch.Tensor],
resolution: int, device: str) -> torch.Tensor:
"""Accept a path, PIL image, HWC array, or CHW tensor; return a batch of 1."""
if isinstance(image, (str, Path)):
img = Image.open(image)
elif isinstance(image, Image.Image):
img = image
elif isinstance(image, np.ndarray):
img = Image.fromarray(image)
elif isinstance(image, torch.Tensor):
arr = image.cpu().numpy() if image.ndim == 3 else image[0].cpu().numpy()
if arr.shape[0] == 3:
arr = arr.transpose(1, 2, 0)
img = Image.fromarray((arr * 255).astype('uint8'))
else:
raise TypeError(f'unsupported image type: {type(image)}')
return normalize(img, resolution, device)
def coco_split(split: str = 'val2017'):
"""Return (COCO handle, image-file lookup) for a COCO split."""
from pycocotools.coco import COCO
coco = COCO(str(COCO_ROOT / 'annotations' / f'instances_{split}.json'))
id_to_file = {i['id']: i['file_name'] for i in coco.loadImgs(coco.getImgIds())}
return coco, id_to_file
def person_labels(coco, img_ids: Sequence[int], device: str = 'cpu') -> torch.Tensor:
"""Image-level person presence for each id, as a bool tensor."""
labels = [
any(a['category_id'] == PERSON_CATEGORY_ID
for a in coco.loadAnns(coco.getAnnIds(imgIds=i, iscrowd=False)))
for i in img_ids
]
return torch.tensor(labels, dtype=torch.bool, device=device)
def image_paths(id_to_file: dict, img_ids: Iterable[int],
split: str = 'val2017') -> List[Path]:
"""Absolute paths for a sequence of image ids within a split."""
root = COCO_ROOT / split
return [root / id_to_file[i] for i in img_ids]
|