rpotter6298
352 lines
12 KiB
Python
352 lines
12 KiB
Python
"""
|
||
portable_versions/image_loader.py
|
||
==================================
|
||
A self-contained image loader with in-memory caching, an optional
|
||
preprocessing pipeline (e.g. disc cropping), and composable augmentations.
|
||
|
||
Returns plain NumPy arrays — works with PyTorch, TensorFlow, JAX, or
|
||
anything else that can consume an ndarray.
|
||
|
||
Dependencies: Pillow, numpy (nothing else)
|
||
|
||
Quickstart
|
||
----------
|
||
from portable_versions.image_loader import ImageLoader, RandomHorizontalFlip, RandomRotation, ColorJitter
|
||
|
||
# 1. Build the loader (once per run)
|
||
loader = ImageLoader(
|
||
target_size=(200, 200),
|
||
normalize=True, # float32 in [0, 1] with ImageNet mean/std
|
||
cache=True, # each image decoded from disk only once
|
||
workers=4, # parallel cache warm-up threads
|
||
preprocessor=my_crop_fn, # optional callable(PIL.Image) -> PIL.Image
|
||
)
|
||
|
||
# 2. Attach augmentations (applied randomly and independently per call)
|
||
loader.augmentation = [
|
||
RandomHorizontalFlip(p=0.5),
|
||
RandomRotation(degrees=15),
|
||
ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1, hue=0.05),
|
||
]
|
||
|
||
# 3. Warm the cache up front (optional but fast)
|
||
loader.warm(all_paths)
|
||
|
||
# 4. Fetch images by path list — call as many times as you like
|
||
# Returns ndarray of shape (N, H, W, 3), dtype float32
|
||
imgs = loader.get_img(train_paths)
|
||
|
||
# For TensorFlow:
|
||
import tensorflow as tf
|
||
tensor = tf.constant(imgs) # (N, H, W, 3)
|
||
|
||
# For PyTorch:
|
||
import torch
|
||
tensor = torch.from_numpy(imgs).permute(0, 3, 1, 2) # (N, C, H, W)
|
||
|
||
|
||
Augmentations reference
|
||
-----------------------
|
||
All augmentation classes live in this file and depend only on PIL + numpy.
|
||
|
||
RandomHorizontalFlip(p=0.5)
|
||
RandomVerticalFlip(p=0.5)
|
||
RandomRotation(degrees=15)
|
||
ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1, hue=0.05)
|
||
RandomGrayscale(p=0.1)
|
||
|
||
You can also pass any callable(PIL.Image.Image) -> PIL.Image.Image as an
|
||
augmentation step.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import random
|
||
import threading
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
from typing import Callable, Iterable, List, Optional, Tuple, Union
|
||
|
||
import numpy as np
|
||
from PIL import Image, ImageEnhance, ImageOps
|
||
|
||
# ImageNet channel statistics (RGB)
|
||
_IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||
_IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||
|
||
PathLike = Union[str, Path]
|
||
|
||
|
||
def _call_preprocessor(
|
||
fn: Callable[..., Image.Image],
|
||
img: Image.Image,
|
||
path: Path,
|
||
) -> Image.Image:
|
||
"""Call preprocessor as fn(img, path) if it accepts two args, else fn(img)."""
|
||
try:
|
||
return fn(img, path)
|
||
except TypeError:
|
||
return fn(img)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Core loader
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class ImageLoader:
|
||
"""
|
||
Preprocessing pipeline + in-memory cache + augmentation, returning NumPy.
|
||
|
||
Parameters
|
||
----------
|
||
target_size : (height, width)
|
||
Output spatial dimensions. Applied after ``preprocessor`` (if any).
|
||
Ignored when a ``preprocessor`` already resizes to the right size.
|
||
normalize : bool
|
||
When True, output is float32 with ImageNet mean/std subtraction.
|
||
When False, output is uint8 in [0, 255].
|
||
cache : bool
|
||
Store decoded+preprocessed images in RAM so each file is read from
|
||
disk at most once. The cache persists across ``get_img`` calls.
|
||
workers : int
|
||
Thread count for ``warm()``. 0 or 1 = single-threaded.
|
||
preprocessor : callable, optional
|
||
Called as ``preprocessor(img: PIL.Image) -> PIL.Image`` before
|
||
resizing and caching. Use this for disc cropping, padding, etc.
|
||
augmentation : list of callables
|
||
Each element is called as ``fn(img: PIL.Image) -> PIL.Image``.
|
||
Applied **after** cache retrieval, so augmentations are NOT cached —
|
||
they are re-sampled independently on every ``get_img`` call.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
target_size: Tuple[int, int] = (224, 224),
|
||
*,
|
||
normalize: bool = True,
|
||
cache: bool = True,
|
||
workers: int = 4,
|
||
preprocessor: Optional[Callable[[Image.Image], Image.Image]] = None,
|
||
) -> None:
|
||
self.target_size = target_size
|
||
self.normalize = normalize
|
||
self.workers = workers
|
||
self.preprocessor = preprocessor
|
||
self.augmentation: List[Callable[[Image.Image], Image.Image]] = []
|
||
|
||
self._cache: Optional[dict[str, np.ndarray]] = {} if cache else None
|
||
self._lock = threading.Lock()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Public
|
||
# ------------------------------------------------------------------
|
||
|
||
def warm(self, paths: Iterable[PathLike]) -> None:
|
||
"""
|
||
Pre-load all *paths* into the cache in parallel.
|
||
|
||
Already-cached paths are skipped, so calling ``warm`` multiple
|
||
times (e.g. once per fold) is safe and only loads new images.
|
||
"""
|
||
if self._cache is None:
|
||
return
|
||
|
||
paths = [str(p) for p in paths]
|
||
to_warm = [p for p in paths if p not in self._cache]
|
||
if not to_warm:
|
||
return
|
||
|
||
already = len(paths) - len(to_warm)
|
||
print(
|
||
f"[ImageLoader] warming {len(to_warm)} images"
|
||
+ (f" ({already} already cached)" if already else ""),
|
||
flush=True,
|
||
)
|
||
|
||
def _load_one(path_str: str) -> None:
|
||
arr = self._decode(path_str)
|
||
with self._lock:
|
||
self._cache.setdefault(path_str, arr)
|
||
|
||
if self.workers <= 1:
|
||
for p in to_warm:
|
||
_load_one(p)
|
||
else:
|
||
with ThreadPoolExecutor(max_workers=self.workers) as ex:
|
||
futures = {ex.submit(_load_one, p): p for p in to_warm}
|
||
for fut in as_completed(futures):
|
||
fut.result()
|
||
|
||
def get_img(
|
||
self,
|
||
paths: Iterable[PathLike],
|
||
augment: bool = True,
|
||
) -> np.ndarray:
|
||
"""
|
||
Return images for the given paths as a single NumPy array.
|
||
|
||
Parameters
|
||
----------
|
||
paths : iterable of path-like
|
||
File paths to load. If the cache is enabled and a path has
|
||
been warmed (or loaded before), it is served from RAM.
|
||
augment : bool
|
||
Apply ``self.augmentation`` pipeline. Set to False at eval time.
|
||
|
||
Returns
|
||
-------
|
||
np.ndarray, shape (N, H, W, 3)
|
||
float32 in [0, 1] (or normalised) if ``self.normalize`` is True,
|
||
otherwise uint8 in [0, 255].
|
||
"""
|
||
imgs = []
|
||
for p in paths:
|
||
img = self._get_one(str(p), augment=augment)
|
||
imgs.append(img)
|
||
return np.stack(imgs, axis=0)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Internal
|
||
# ------------------------------------------------------------------
|
||
|
||
def _decode(self, path_str: str) -> np.ndarray:
|
||
"""Open, preprocess, and resize → uint8 HWC ndarray (for the cache)."""
|
||
img = Image.open(path_str).convert("RGB")
|
||
if self.preprocessor is not None:
|
||
img = _call_preprocessor(self.preprocessor, img, Path(path_str))
|
||
# Only resize here if preprocessor didn't already produce target_size
|
||
if img.size != (self.target_size[1], self.target_size[0]):
|
||
img = img.resize((self.target_size[1], self.target_size[0]), Image.BILINEAR)
|
||
return np.asarray(img, dtype=np.uint8)
|
||
|
||
def _get_one(self, path_str: str, augment: bool) -> np.ndarray:
|
||
if self._cache is not None:
|
||
arr = self._cache.get(path_str)
|
||
if arr is None:
|
||
arr = self._decode(path_str)
|
||
with self._lock:
|
||
self._cache.setdefault(path_str, arr)
|
||
img = Image.fromarray(arr, mode="RGB")
|
||
else:
|
||
img = Image.open(path_str).convert("RGB")
|
||
if self.preprocessor is not None:
|
||
img = self.preprocessor(img)
|
||
if img.size != (self.target_size[1], self.target_size[0]):
|
||
img = img.resize((self.target_size[1], self.target_size[0]), Image.BILINEAR)
|
||
|
||
if augment and self.augmentation:
|
||
for fn in self.augmentation:
|
||
img = fn(img)
|
||
|
||
arr = np.asarray(img, dtype=np.float32) / 255.0
|
||
if self.normalize:
|
||
arr = (arr - _IMAGENET_MEAN) / _IMAGENET_STD
|
||
else:
|
||
arr = (arr * 255).clip(0, 255).astype(np.uint8)
|
||
return arr
|
||
|
||
def __len__(self) -> int:
|
||
"""Number of images currently in the cache."""
|
||
return len(self._cache) if self._cache is not None else 0
|
||
|
||
def __repr__(self) -> str:
|
||
return (
|
||
f"ImageLoader(target_size={self.target_size}, "
|
||
f"normalize={self.normalize}, "
|
||
f"cached={len(self)}, "
|
||
f"augmentations={len(self.augmentation)})"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Augmentation primitives (PIL-only, no torch/tf dependencies)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class RandomHorizontalFlip:
|
||
"""Flip image left-right with probability *p*."""
|
||
def __init__(self, p: float = 0.5):
|
||
self.p = p
|
||
|
||
def __call__(self, img: Image.Image) -> Image.Image:
|
||
return ImageOps.mirror(img) if random.random() < self.p else img
|
||
|
||
|
||
class RandomVerticalFlip:
|
||
"""Flip image top-bottom with probability *p*."""
|
||
def __init__(self, p: float = 0.5):
|
||
self.p = p
|
||
|
||
def __call__(self, img: Image.Image) -> Image.Image:
|
||
return ImageOps.flip(img) if random.random() < self.p else img
|
||
|
||
|
||
class RandomRotation:
|
||
"""Rotate by a uniformly-sampled angle in [-degrees, +degrees]."""
|
||
def __init__(self, degrees: float = 15):
|
||
self.degrees = degrees
|
||
|
||
def __call__(self, img: Image.Image) -> Image.Image:
|
||
angle = random.uniform(-self.degrees, self.degrees)
|
||
return img.rotate(angle, resample=Image.BILINEAR, expand=False)
|
||
|
||
|
||
class ColorJitter:
|
||
"""
|
||
Randomly jitter brightness, contrast, saturation, and hue.
|
||
|
||
Each factor is sampled uniformly from [1 - amount, 1 + amount].
|
||
Hue shift is sampled from [-hue, +hue] (range 0–0.5).
|
||
Pass 0 for any channel to leave it unchanged.
|
||
"""
|
||
def __init__(
|
||
self,
|
||
brightness: float = 0.2,
|
||
contrast: float = 0.2,
|
||
saturation: float = 0.1,
|
||
hue: float = 0.05,
|
||
):
|
||
self.brightness = brightness
|
||
self.contrast = contrast
|
||
self.saturation = saturation
|
||
self.hue = hue
|
||
|
||
def __call__(self, img: Image.Image) -> Image.Image:
|
||
ops = []
|
||
if self.brightness:
|
||
ops.append(("brightness", self.brightness))
|
||
if self.contrast:
|
||
ops.append(("contrast", self.contrast))
|
||
if self.saturation:
|
||
ops.append(("saturation", self.saturation))
|
||
if self.hue:
|
||
ops.append(("hue", self.hue))
|
||
random.shuffle(ops)
|
||
|
||
for kind, amount in ops:
|
||
factor = random.uniform(1 - amount, 1 + amount)
|
||
if kind == "brightness":
|
||
img = ImageEnhance.Brightness(img).enhance(factor)
|
||
elif kind == "contrast":
|
||
img = ImageEnhance.Contrast(img).enhance(factor)
|
||
elif kind == "saturation":
|
||
img = ImageEnhance.Color(img).enhance(factor)
|
||
elif kind == "hue":
|
||
# PIL has no direct hue enhancer — shift via HSV in numpy
|
||
arr = np.asarray(img.convert("HSV"), dtype=np.int16)
|
||
shift = int(random.uniform(-self.hue, self.hue) * 255)
|
||
arr[:, :, 0] = (arr[:, :, 0] + shift) % 256
|
||
img = Image.fromarray(arr.astype(np.uint8), mode="HSV").convert("RGB")
|
||
return img
|
||
|
||
|
||
class RandomGrayscale:
|
||
"""Convert to grayscale (keeping 3 channels) with probability *p*."""
|
||
def __init__(self, p: float = 0.1):
|
||
self.p = p
|
||
|
||
def __call__(self, img: Image.Image) -> Image.Image:
|
||
if random.random() < self.p:
|
||
img = ImageOps.grayscale(img).convert("RGB")
|
||
return img
|
||
|
||
|