| |
| """Round 14 (P1): native host preprocess -- ``resize.c`` (crop + grey pad + PIL BILINEAR box resize + mirror in one call, |
| bit-identical to the PIL path), built lazily with the host C compiler into a per-user cache and loaded with ctypes. |
| |
| Nothing here is required: if there is no compiler, the build fails, or the load-time self-check (random crops vs the PIL |
| path, bit for bit) fails, :func:`crop_resize` is ``None`` and preprocess keeps the PIL path (same bytes, slower). |
| ``HAMER_PRE_NATIVE=0`` disables it. |
| """ |
| from __future__ import annotations |
|
|
| import ctypes |
| import hashlib |
| import os |
| import subprocess |
| import tempfile |
| import threading |
| from pathlib import Path |
| from typing import Optional |
|
|
| import numpy as np |
|
|
| _SRCS = [Path(__file__).with_name("resize.c")] |
| |
| _CFLAGS = ["-O3", "-march=native", "-ffp-contract=off", "-fno-fast-math", "-shared", "-fPIC"] |
| _LOCK = threading.Lock() |
| _FN = None |
| _FN_PS = None |
| _IM2COL = None |
| _TRIED = False |
| STATUS = "not loaded" |
|
|
|
|
| def _cache_dir() -> Path: |
| d = os.environ.get("HAMER_NATIVE_DIR") |
| if d: |
| return Path(d) |
| base = os.environ.get("XDG_CACHE_HOME") or str(Path.home() / ".cache") |
| return Path(base) / "hamer-native" |
|
|
|
|
| def _build() -> Optional[Path]: |
| src = b"".join(p.read_bytes() for p in _SRCS) |
| cc = os.environ.get("CC", "cc") |
| tag = hashlib.sha1(src + " ".join(_CFLAGS).encode() + os.uname().machine.encode()).hexdigest()[:16] |
| for d in (_cache_dir(), Path(tempfile.gettempdir()) / f"hamer-native-{os.getuid()}"): |
| so = d / f"hamer_resize_{tag}.so" |
| if so.exists(): |
| return so |
| try: |
| d.mkdir(parents=True, exist_ok=True) |
| tmp = d / f".{so.name}.{os.getpid()}" |
| subprocess.run([cc, *_CFLAGS, "-o", str(tmp), *map(str, _SRCS)], check=True, capture_output=True, timeout=120) |
| os.replace(tmp, so) |
| return so |
| except Exception: |
| continue |
| return None |
|
|
|
|
| def _self_check(fn) -> bool: |
| """The native path must reproduce the PIL path bit for bit (random frames / boxes / both hands, edge-crossing boxes).""" |
| from PIL import Image |
|
|
| from hamer import preprocess as pp |
|
|
| rng = np.random.default_rng(1234) |
| for i in range(24): |
| H, W = (int(v) for v in rng.integers(48, 400, 2)) |
| a = rng.integers(0, 256, (H, W, 3), dtype=np.uint8) |
| img = Image.fromarray(a) |
| w, h = rng.uniform(4, W), rng.uniform(4, H) |
| x1, y1 = rng.uniform(-w / 2, W - w / 2), rng.uniform(-h / 2, H - h / 2) |
| bbox, rf, right = (x1, y1, x1 + w, y1 + h), float(rng.choice([2.0, 1.0, 3.3])), bool(i % 2) |
| if pp.hamer_bbox_size(bbox, rf) < 2: |
| continue |
| crop, meta = pp._square_crop(img, bbox, rf) |
| ref = pp.resize_crop_hw(crop, right, bands=1) |
| got = _call(fn, a, bbox, rf, right) |
| got_chw = _call(fn, a, bbox, rf, right, chw=True) |
| if got is None or not np.array_equal(ref, got[0]) or got[1] != meta: |
| return False |
| if got_chw is None or not np.array_equal(ref.transpose(2, 0, 1), got_chw[0]) or got_chw[1] != meta: |
| return False |
| return True |
|
|
|
|
| def _call(fn, a: np.ndarray, bbox, rescale_factor: float, is_right: bool, chw: bool = False): |
| """(u8 HWC [256, 192, 3] or planar CHW [3, 256, 192] when chw, (cx, cy, bbox_size)) or None when the PIL path must handle it (degenerate boxes).""" |
| if a.ndim != 3 or a.shape[2] != 3 or a.dtype != np.uint8 or a.strides[1:] != (3, 1): |
| return None |
| H, W = a.shape[:2] |
| return _call_ptr(fn, a.ctypes.data, H, W, a.strides[0], 3, bbox, rescale_factor, is_right, chw) |
|
|
|
|
| def _call_ptr(fn, ptr: int, H: int, W: int, row_stride: int, ps: int, bbox, rescale_factor: float, is_right: bool, |
| chw: bool = False): |
| """_call on a raw u8 image pointer with ``ps`` bytes per pixel (3: packed RGB; 4: Pillow's RGBX rows).""" |
| from hamer.preprocess import hamer_bbox_size |
|
|
| x1, y1, x2, y2 = bbox |
| cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0 |
| bs = hamer_bbox_size(bbox, rescale_factor) |
| half = bs / 2.0 |
| S = int(round(bs)) |
| |
| l, t = max(0, int(round(cx - half))), max(0, int(round(cy - half))) |
| r, b = min(W, int(round(cx + half))), min(H, int(round(cy + half))) |
| px, py = max(0, int(round(-(cx - half)))), max(0, int(round(-(cy - half)))) |
| if S < 2 or r <= l or b <= t: |
| return None |
| out = np.empty((3, 256, 192) if chw else (256, 192, 3), np.uint8) |
| args = (S, l, t, r - l, b - t, px, py, 0 if is_right else 1, 1 if chw else 0, out.ctypes.data) |
| rc = fn(ptr, H, W, row_stride, *args) if ps == 3 and fn is not _FN_PS else fn(ptr, H, W, row_stride, ps, *args) |
| if rc != 0: |
| return None |
| return out, (cx, cy, bs) |
|
|
|
|
| def load(): |
| """The ctypes function (or None). Built / checked once per process.""" |
| global _FN, _FN_PS, _IM2COL, _TRIED, STATUS |
| with _LOCK: |
| if _TRIED: |
| return _FN |
| _TRIED = True |
| pre_off = os.environ.get("HAMER_PRE_NATIVE", "1").strip().lower() in ("0", "false", "no", "off") |
| try: |
| so = _build() |
| if so is None: |
| STATUS = "no compiler / build failed: PIL path" |
| return None |
| lib = ctypes.CDLL(str(so)) |
| fn = lib.hamer_crop_resize |
| fn.restype = ctypes.c_int |
| fn.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_long] + [ctypes.c_int] * 9 + [ctypes.c_void_p] |
| ic = lib.hamer_im2col_tile_bf16 |
| ic.restype = ctypes.c_int |
| ic.argtypes = [ctypes.c_void_p, ctypes.c_long, ctypes.c_long, ctypes.c_long, ctypes.c_int, ctypes.c_int, |
| ctypes.c_int, ctypes.c_void_p] |
| _IM2COL = ic |
| if pre_off: |
| STATUS = "off (HAMER_PRE_NATIVE=0)" |
| return None |
| if not _self_check(fn): |
| STATUS = "self-check FAILED (not bit-identical to PIL): PIL path" |
| return None |
| _FN = fn |
| STATUS = f"on ({so.name}, self-check bit-identical)" |
| if os.environ.get("HAMER_PRE_ZEROCOPY", "1").strip().lower() not in ("0", "false", "no", "off"): |
| fps = lib.hamer_crop_resize_ps |
| fps.restype = ctypes.c_int |
| fps.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_long, ctypes.c_int] \ |
| + [ctypes.c_int] * 9 + [ctypes.c_void_p] |
| _FN_PS = fps |
| if _self_check_zc(): |
| STATUS += "; zero-copy PIL buffer on" |
| else: |
| _FN_PS = None |
| STATUS += "; zero-copy PIL buffer off (self-check failed)" |
| except Exception as e: |
| STATUS = f"load failed ({e!r}): PIL path" |
| _FN = None |
| return _FN |
|
|
|
|
| def im2col_tile_fn(): |
| """Round 14: ``hamer_im2col_tile_bf16`` (u8 crop -> bf16 TILE im2col of the centred crop) or None.""" |
| load() |
| return _IM2COL |
|
|
|
|
| def crop_resize(a: np.ndarray, bbox, rescale_factor: float, is_right: bool, chw: bool = False): |
| fn = load() |
| if fn is None: |
| return None |
| return _call(fn, a, bbox, rescale_factor, is_right, chw) |
|
|
|
|
| |
| |
| |
| class _ArrowArray(ctypes.Structure): |
| pass |
|
|
|
|
| _ArrowArray._fields_ = [("length", ctypes.c_int64), ("null_count", ctypes.c_int64), ("offset", ctypes.c_int64), |
| ("n_buffers", ctypes.c_int64), ("n_children", ctypes.c_int64), |
| ("buffers", ctypes.POINTER(ctypes.c_void_p)), |
| ("children", ctypes.POINTER(ctypes.POINTER(_ArrowArray))), ("dictionary", ctypes.c_void_p), |
| ("release", ctypes.c_void_p), ("private_data", ctypes.c_void_p)] |
| _CAP_GET = ctypes.pythonapi.PyCapsule_GetPointer |
| _CAP_GET.restype = ctypes.c_void_p |
| _CAP_GET.argtypes = [ctypes.py_object, ctypes.c_char_p] |
|
|
|
|
| def _rgbx_view(img): |
| """(capsules to keep alive, pointer, H, W) of an RGB PIL image's RGBX pixel buffer, or None.""" |
| if img.mode != "RGB" or not hasattr(img, "__arrow_c_array__"): |
| return None |
| try: |
| img.load() |
| sch, arr = img.__arrow_c_array__() |
| a = ctypes.cast(_CAP_GET(arr, b"arrow_array"), ctypes.POINTER(_ArrowArray)).contents |
| W, H = img.size |
| if a.n_children != 1 or a.length != H * W or a.offset != 0: |
| return None |
| c = a.children[0].contents |
| if c.n_buffers != 2 or c.length != H * W * 4 or c.offset != 0 or not c.buffers[1]: |
| return None |
| return (sch, arr), c.buffers[1], H, W |
| except Exception: |
| return None |
|
|
|
|
| def _self_check_zc() -> bool: |
| """The zero-copy path must give the bytes of the np.asarray path (random frames / boxes / both hands).""" |
| from PIL import Image |
|
|
| rng = np.random.default_rng(4321) |
| for i in range(16): |
| H, W = (int(v) for v in rng.integers(48, 400, 2)) |
| a = rng.integers(0, 256, (H, W, 3), dtype=np.uint8) |
| img = Image.fromarray(a) |
| w, h = rng.uniform(4, W), rng.uniform(4, H) |
| x1, y1 = rng.uniform(-w / 2, W - w / 2), rng.uniform(-h / 2, H - h / 2) |
| bbox, rf, right = (x1, y1, x1 + w, y1 + h), float(rng.choice([2.0, 1.0, 3.3])), bool(i % 2) |
| ref = _call(_FN, a, bbox, rf, right, chw=True) |
| v = _rgbx_view(img) |
| if v is None: |
| return False |
| got = _call_ptr(_FN_PS, v[1], v[2], v[3], v[3] * 4, 4, bbox, rf, right, chw=True) |
| if (ref is None) != (got is None) or (ref is not None and (not np.array_equal(ref[0], got[0]) or ref[1] != got[1])): |
| return False |
| return True |
|
|
|
|
| def crop_resize_image(img, bbox, rescale_factor: float, is_right: bool, chw: bool = False): |
| """crop_resize on a decoded PIL image, reading its pixel buffer in place when possible (round 15).""" |
| fn = load() |
| if fn is None: |
| return None |
| if _FN_PS is not None: |
| v = _rgbx_view(img) |
| if v is not None: |
| keep, ptr, H, W = v |
| res = _call_ptr(_FN_PS, ptr, H, W, W * 4, 4, bbox, rescale_factor, is_right, chw) |
| del keep |
| return res |
| return _call(fn, np.asarray(img), bbox, rescale_factor, is_right, chw) |
|
|
