| |
| """Fast served postprocess + response encoding (default; ``MOGE_SERVE_FAST=0`` -> the torch/numpy path). |
| |
| Same results as ``postprocess.canvas_to_camera`` + ``postprocess.to_original`` + ``np.savez_compressed`` |
| + ``base64.b64encode`` (``tests/test_serve_fastpath.py`` checks parity for every request-flag |
| combination), with three changes in HOW the work is done: |
| |
| * ``camera``: the focal / z-shift solve runs on the 64x64 nearest-neighbour samples only. Upstream |
| ``recover_focal_shift`` builds the same samples with ``F.interpolate(mode='nearest')`` over the full |
| 1080x1920 maps; here the identical rows / columns are gathered directly (indices from torch's own |
| 'nearest' kernel), and the same scipy solver runs on the same arrays -> bit-identical focal / shift. |
| * ``finish``: shift, metric scale, letterbox crop, bilinear resize to the original size, normal |
| re-normalize, nearest mask resize, re-projection (force_projection) and masking in ONE multithreaded |
| C++ pass (``host_ext.moge_finish``) that writes the response arrays directly (normal already float16 |
| for npz). Arithmetic follows torch's CPU kernels operation by operation (see the C++ comment). |
| * ``npz_b64``: a valid ``.npz`` (zip of ``.npy`` members, ``np.load``-compatible, same keys / dtypes / |
| shapes / bytes) whose deflate streams are produced in parallel: each member is cut into 1 MiB pieces, |
| each piece is raw-deflated independently (``Z_SYNC_FLUSH`` between pieces, ``Z_FINISH`` on the last: |
| the pigz construction, one valid deflate stream per member) by a thread pool. Compressor (round 4): |
| ISA-L's ``isal_deflate`` called through ctypes from python-isal's own shared object (``_IsalCt``, |
| level 2): ctypes releases the GIL, python-isal's ``compressobj`` does not (its pool saturated at |
| ~1.2 GB/s); output pieces go to one persistent scratch buffer and base64 runs over the segment list |
| (``moge_b64_segs``) straight into the JSON body buffer. Fallbacks: python-isal ``isal_zlib`` (level 2) -- ~10x faster than |
| zlib on these float maps, payload +4% (35.8 vs 34.5 MB for a 1080p response) -- else stdlib zlib at |
| level 1 (as small as ``np.savez_compressed``'s zlib level 6, 2.5x faster). ``MOGE_NPZ_DEFLATE`` / |
| ``MOGE_NPZ_LEVEL`` override. Any valid deflate stream inflates to the same bytes, so ``np.load`` |
| returns identical arrays. base64 is a multithreaded C++ loop. |
| """ |
| from __future__ import annotations |
|
|
| import ctypes |
| import io |
| import math |
| import os |
| import struct |
| import threading |
| import zlib |
| from concurrent.futures import ThreadPoolExecutor |
| from typing import Dict, Optional, Tuple |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| from tt_moge.tt import host_ext |
|
|
| Region = Tuple[int, int, int, int] |
| _CHUNK = 1 << 20 |
| _pool: Optional[ThreadPoolExecutor] = None |
| _pool_lock = threading.Lock() |
| _cache: Dict[tuple, object] = {} |
|
|
|
|
| def enabled() -> bool: |
| if os.environ.get("MOGE_SERVE_FAST", "1").strip().lower() in ("0", "false", "no", "off"): |
| return False |
| return host_ext.enabled() |
|
|
|
|
| def nthreads() -> int: |
| return int(os.environ.get("MOGE_SERVE_THREADS", "0")) or host_ext.nthreads() |
|
|
|
|
| def _get_pool() -> ThreadPoolExecutor: |
| global _pool |
| with _pool_lock: |
| if _pool is None: |
| _pool = ThreadPoolExecutor(max_workers=max(1, nthreads()), thread_name_prefix="moge-npz") |
| return _pool |
|
|
|
|
| def _nearest_idx(n_in: int, n_out: int) -> torch.Tensor: |
| """Source indices of torch's 'nearest' resize n_in -> n_out (from the kernel itself).""" |
| key = ("nearest", n_in, n_out) |
| if key not in _cache: |
| src = torch.arange(n_in, dtype=torch.float32)[None, None] |
| idx = F.interpolate(src, size=n_out, mode="nearest")[0, 0] |
| _cache[key] = idx.to(torch.int64) |
| return _cache[key] |
|
|
|
|
| def _view_plane_uv_axes(width: int, height: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| """The u (width) / v (height) linspaces of moge's normalized_view_plane_uv (same expressions).""" |
| key = ("uv", width, height) |
| if key not in _cache: |
| aspect_ratio = width / height |
| span_x = aspect_ratio / (1 + aspect_ratio ** 2) ** 0.5 |
| span_y = 1 / (1 + aspect_ratio ** 2) ** 0.5 |
| u = torch.linspace(-span_x * (width - 1) / width, span_x * (width - 1) / width, width, dtype=torch.float32) |
| v = torch.linspace(-span_y * (height - 1) / height, span_y * (height - 1) / height, height, dtype=torch.float32) |
| _cache[key] = (u, v) |
| return _cache[key] |
|
|
|
|
| @torch.inference_mode() |
| def camera(out: Dict[str, torch.Tensor], region: Optional[Region], fov_x: Optional[float] = None, |
| downsample: Tuple[int, int] = (64, 64)) -> Dict[str, torch.Tensor]: |
| """``postprocess.canvas_to_camera`` without the full-resolution work: returns ``focal``, ``shift``, |
| ``metric_scale`` (0-d float32 tensors) and the canvas ``intrinsics`` [3,3]. ``region`` is the |
| valid (non-letterbox) rectangle, None = the whole canvas (fit='stretch').""" |
| from moge.utils.geometry_numpy import solve_optimal_focal_shift, solve_optimal_shift |
| import utils3d |
|
|
| points = out["points"][0] |
| height, width = points.shape[:2] |
| aspect_ratio = width / height |
| iy, ix = _nearest_idx(height, downsample[0]), _nearest_idx(width, downsample[1]) |
| p_lr = points.index_select(0, iy).index_select(1, ix).float() |
| m_lr = out["mask"][0].index_select(0, iy).index_select(1, ix).float() > 0.5 |
| if region is not None: |
| x0, y0, x1, y1 = region |
| m_lr &= ((iy >= y0) & (iy < y1))[:, None] & ((ix >= x0) & (ix < x1))[None, :] |
| u, v = _view_plane_uv_axes(width, height) |
| uv_lr = torch.stack(torch.meshgrid(u[ix], v[iy], indexing="xy"), dim=-1) |
|
|
| m_np = m_lr.numpy() |
| p_np = p_lr.numpy()[m_np] |
| uv_np = uv_lr.numpy()[m_np] |
| if fov_x is not None: |
| focal_t = torch.tensor( |
| [aspect_ratio / (1 + aspect_ratio ** 2) ** 0.5 / math.tan(math.radians(fov_x) / 2)], dtype=torch.float32) |
| if uv_np.shape[0] < 2: |
| f_val, s_val = 1, 0 |
| elif fov_x is None: |
| s, f = solve_optimal_focal_shift(uv_np, p_np) |
| f_val, s_val = float(f), float(s) |
| else: |
| s_val = float(solve_optimal_shift(uv_np, p_np, focal_t.numpy()[0])) |
| shift = torch.tensor([s_val], dtype=torch.float32).reshape(()) |
| focal = (torch.tensor([f_val], dtype=torch.float32) if fov_x is None else focal_t).reshape(-1)[0] |
|
|
| fx = focal / 2 * (1 + aspect_ratio ** 2) ** 0.5 / aspect_ratio |
| fy = focal / 2 * (1 + aspect_ratio ** 2) ** 0.5 |
| half = torch.tensor(0.5, dtype=torch.float32) |
| intrinsics = utils3d.pt.intrinsics_from_focal_center(fx[None], fy[None], half, half)[0] |
| metric_scale = out["metric_scale"].float().reshape(-1)[0] |
| return {"focal": focal, "shift": shift, "metric_scale": metric_scale, "intrinsics": intrinsics} |
|
|
|
|
| def _kinv(intrinsics: torch.Tensor) -> torch.Tensor: |
| """Rows 0..2 of inv(4x4 intrinsics) exactly as utils3d.unproject_cv builds and inverts it.""" |
| k = intrinsics[None] |
| k4 = torch.cat([ |
| torch.cat([k, torch.zeros((*k.shape[:-2], 3, 1), dtype=k.dtype)], dim=-1), |
| torch.tensor([[0, 0, 0, 1]], dtype=k.dtype).expand(*k.shape[:-2], 1, 4)], dim=-2) |
| return torch.linalg.inv(k4)[0, :3, :].contiguous() |
|
|
|
|
| def _uv_map_axes(height: int, width: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| """utils3d uv_map linspaces (u: width, v: height).""" |
| key = ("uvmap", height, width) |
| if key not in _cache: |
| u = torch.linspace(0.5 / width, 1 - 0.5 / width, width, dtype=torch.float32) |
| v = torch.linspace(0.5 / height, 1 - 0.5 / height, height, dtype=torch.float32) |
| _cache[key] = (u.contiguous(), v.contiguous()) |
| return _cache[key] |
|
|
|
|
| class _Finish: |
| """Arguments of one ``host_ext.moge_finish_rows`` call sequence (all rows or row bands).""" |
|
|
| @torch.inference_mode() |
| def __init__(self, out, cam, canvas, region, size, force_projection, apply_mask, normal_half): |
| from tt_moge.server import postprocess as pp |
| H, W = size |
| self.pts = out["points"][0].float().contiguous() |
| self.nrm = out["normal"][0].float().contiguous() |
| self.msk = out["mask"][0].float().contiguous() |
| self.Hc, self.Wc = self.pts.shape[:2] |
| self.region, self.size = region, size |
| self.intrinsics = pp.crop_intrinsics(cam["intrinsics"], canvas, region) |
| self.kinv = _kinv(self.intrinsics) |
| self.u, self.v = _uv_map_axes(H, W) |
| x0, y0, x1, y1 = region |
| self.ny = _nearest_idx(y1 - y0, H).to(torch.int32).contiguous() |
| self.nx = _nearest_idx(x1 - x0, W).to(torch.int32).contiguous() |
| self.shift, self.ms = float(cam["shift"]), float(cam["metric_scale"]) |
| self.fp, self.am, self.half = int(force_projection), int(apply_mask), int(normal_half) |
| self.depth = np.empty((H, W), np.float32) |
| self.points = np.empty((H, W, 3), np.float32) |
| self.normal = np.empty((H, W, 3), np.float16 if normal_half else np.float32) |
| self.mask = np.empty((H, W), np.uint8) |
| self.cnt, self.mn, self.mx = 0, math.inf, -math.inf |
|
|
| def rows(self, ya: int, yb: int, nthreads_: int): |
| H, W = self.size |
| x0, y0, x1, y1 = self.region |
| stats = np.zeros(3, np.float64) |
| rc = host_ext.lib().moge_finish_rows( |
| self.pts.data_ptr(), self.nrm.data_ptr(), self.msk.data_ptr(), self.Hc, self.Wc, x0, y0, x1, y1, |
| self.shift, self.ms, H, W, self.ny.data_ptr(), self.nx.data_ptr(), self.u.data_ptr(), self.v.data_ptr(), |
| self.kinv.data_ptr(), self.fp, self.am, self.depth.ctypes.data, self.points.ctypes.data, |
| self.normal.ctypes.data, self.half, self.mask.ctypes.data, stats.ctypes.data, nthreads_, ya, yb) |
| if rc != 0: |
| raise ValueError(f"moge_finish: bad region {self.region} for canvas {(self.Wc, self.Hc)}") |
| self.cnt += int(stats[0]) |
| self.mn, self.mx = min(self.mn, float(stats[1])), max(self.mx, float(stats[2])) |
|
|
| def result(self) -> Dict[str, object]: |
| return {"depth": self.depth, "points": self.points, "normal": self.normal, "mask": self.mask.view(np.bool_), |
| "intrinsics": self.intrinsics, "valid_count": self.cnt, "depth_min": self.mn, "depth_max": self.mx} |
|
|
|
|
| @torch.inference_mode() |
| def finish(out: Dict[str, torch.Tensor], cam: Dict[str, torch.Tensor], canvas: Tuple[int, int], region: Region, |
| size: Tuple[int, int], force_projection: bool = True, apply_mask: bool = True, |
| normal_half: bool = True) -> Dict[str, object]: |
| """``postprocess.to_original(canvas_to_camera(...))`` as numpy arrays: ``depth`` f32 [H,W], |
| ``points`` f32 [H,W,3], ``normal`` [H,W,3] (f16 if ``normal_half`` else f32), ``mask`` bool [H,W], |
| ``intrinsics`` torch [3,3] (cropped, normalized), plus ``valid_count`` / ``depth_min`` / ``depth_max``.""" |
| f = _Finish(out, cam, canvas, region, size, force_projection, apply_mask, normal_half) |
| f.rows(0, size[0], nthreads()) |
| return f.result() |
|
|
|
|
| def gather_segments(segs) -> np.ndarray: |
| """``b"".join(segs)`` into one uint8 numpy buffer (served as a memoryview body, no bytes copy).""" |
| views = [np.frombuffer(memoryview(x).cast("B"), np.uint8) for x in segs] |
| out = np.empty(sum(v.size for v in views), np.uint8) |
| o = 0 |
| for v in views: |
| out[o:o + v.size] = v |
| o += v.size |
| return out |
|
|
|
|
| def finish_npz_b64(out, cam, canvas, region, size, force_projection: bool, apply_mask: bool, |
| metric_scale: float, headroom: int = 0, bands: Optional[int] = None, |
| timing: Optional[Dict[str, float]] = None, raw_npz: bool = False): |
| """``finish`` + ``npz_b64_bytes`` of the npz response arrays, overlapped: the finish pass runs in row |
| bands and each band's rows of points / depth / normal / mask are deflated (thread pool) while the |
| next band is computed. Same arrays, same npz members (bit-identical decoded values). Returns |
| (finish result dict, base64) -- or with ``raw_npz`` the npz file itself as a uint8 array (npz_binary). |
| ``timing`` (optional) gets ``finish`` = seconds until the last band.""" |
| import time |
| t0 = time.perf_counter() |
| bands = bands or int(os.environ.get("MOGE_FINISH_BANDS", "4")) |
| f = _Finish(out, cam, canvas, region, size, force_projection, apply_mask, True) |
| H = size[0] |
| intr = f.intrinsics.numpy().astype(np.float32) |
| names = ("points", "depth", "normal", "mask", "intrinsics", "metric_scale") |
| arrs = {"points": f.points, "depth": f.depth, "normal": f.normal, "mask": f.mask, |
| "intrinsics": intr, "metric_scale": np.float32(metric_scale)} |
| mod, default_level = _deflate_backend() |
| level = int(os.environ.get("MOGE_NPZ_LEVEL", str(default_level))) |
| ct = isinstance(mod, _IsalCt) |
| edges = [round(H * b / bands) for b in range(bands + 1)] |
| edges = sorted(set(edges)) |
| |
| hdrs, plan = {}, {} |
| for nm in names: |
| a = np.asanyarray(arrs[nm]) |
| hb = io.BytesIO() |
| np.lib.format.write_array_header_1_0(hb, np.lib.format.header_data_from_array_1_0(a)) |
| hdrs[nm] = hb.getvalue() |
| raw = memoryview(a.reshape(-1).view(np.uint8)) if a.size else memoryview(b"") |
| per_band = [] |
| if nm in ("points", "depth", "normal", "mask"): |
| rowb = len(raw) // H |
| for b in range(len(edges) - 1): |
| lo, hi = edges[b] * rowb, edges[b + 1] * rowb |
| per_band.append([raw[i:min(i + _CHUNK, hi)] for i in range(lo, hi, _CHUNK)]) |
| else: |
| per_band.append([raw[i:i + _CHUNK] for i in range(0, len(raw), _CHUNK)]) |
| plan[nm] = (raw, per_band) |
| with _scratch_lock: |
| scratch = None |
| if ct: |
| need = sum(mod.bound(len(hdrs[nm])) + sum(mod.bound(len(p)) for ps in plan[nm][1] for p in ps) |
| for nm in names) |
| scratch = _scratch(need) |
| soff = 0 |
| pool = _get_pool() |
| futs = {nm: [] for nm in names} |
|
|
| def submit(nm, piece, last): |
| nonlocal soff |
| dst = None |
| if ct: |
| cap = mod.bound(len(piece)) |
| dst, soff = scratch[soff:soff + cap], soff + cap |
| futs[nm].append(pool.submit(_deflate_piece, mod, piece, level, last, dst)) |
|
|
| def n_pieces(nm): |
| return sum(len(ps) for ps in plan[nm][1]) |
|
|
| for nm in names: |
| submit(nm, hdrs[nm], n_pieces(nm) == 0) |
| if nm in ("intrinsics", "metric_scale"): |
| ps = plan[nm][1][0] |
| for j, p in enumerate(ps): |
| submit(nm, p, j == len(ps) - 1) |
| for b in range(len(edges) - 1): |
| f.rows(edges[b], edges[b + 1], nthreads()) |
| for nm in ("points", "depth", "normal", "mask"): |
| ps = plan[nm][1][b] |
| for j, p in enumerate(ps): |
| submit(nm, p, b == len(edges) - 2 and j == len(ps) - 1) |
| if timing is not None: |
| timing["finish"] = time.perf_counter() - t0 |
| stats_fut = None |
| if _STATS_OVERLAP: |
| stats_fut = _get_aux_pool().submit(depth_stats, f.result()) |
| crcs = {nm: pool.submit(_crc, mod, [hdrs[nm]] + [p for ps in plan[nm][1] for p in ps]) for nm in names} |
| out_segs, central = [], [] |
| off = 0 |
| for nm in names: |
| comp = [x.result() for x in futs[nm]] |
| csize = sum(memoryview(c).nbytes for c in comp) |
| usize = len(hdrs[nm]) + len(plan[nm][0]) |
| crc = crcs[nm].result() |
| fb = (nm + ".npy").encode() |
| lh = struct.pack("<IHHHHHIIIHH", 0x04034B50, 20, 0, 8, 0, 0x21, crc, csize, usize, len(fb), 0) + fb |
| central.append(struct.pack("<IHHHHHHIIIHHHHHII", 0x02014B50, 20, 20, 0, 8, 0, 0x21, crc, csize, usize, |
| len(fb), 0, 0, 0, 0, 0o600 << 16, off) + fb) |
| out_segs.append(lh) |
| out_segs.extend(comp) |
| off += len(lh) + csize |
| if timing is not None: |
| timing["deflated"] = time.perf_counter() - t0 |
| cd = b"".join(central) |
| out_segs.append(cd) |
| out_segs.append(struct.pack("<IHHHHIIH", 0x06054B50, 0, 0, len(central), len(central), len(cd), off, 0)) |
| b64v = gather_segments(out_segs) if raw_npz else b64_segments(out_segs, headroom) |
| if timing is not None: |
| timing["b64"] = time.perf_counter() - t0 |
| res = f.result() |
| if stats_fut is not None: |
| res["depth_stats"] = stats_fut.result() |
| return res, b64v |
|
|
|
|
| _STATS_OVERLAP = os.environ.get("MOGE_STATS_OVERLAP", "1").strip().lower() not in ("0", "false", "no", "off") |
| _aux_pool: Optional[ThreadPoolExecutor] = None |
|
|
|
|
| def _get_aux_pool() -> ThreadPoolExecutor: |
| global _aux_pool |
| with _pool_lock: |
| if _aux_pool is None: |
| _aux_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="moge-aux") |
| return _aux_pool |
|
|
|
|
| def depth_stats(res: Dict[str, object]) -> Optional[Dict[str, float]]: |
| """{"min", "median", "max"} of the valid depths, == the numpy expressions in app.predict (precomputed |
| by finish_npz_b64 when overlapped).""" |
| if "depth_stats" in res: |
| return res["depth_stats"] |
| n = res["valid_count"] |
| if n == 0: |
| return None |
| depth, mask = res["depth"], res["mask"] |
| L = host_ext.lib() |
| kth = lambda k: np.float32(L.moge_masked_kth(depth.ctypes.data, mask.ctypes.data, depth.size, k, nthreads())) |
| if n % 2: |
| med = kth(n // 2) |
| else: |
| med = np.float32(kth(n // 2 - 1) + kth(n // 2)) / np.float32(2) |
| return {"min": res["depth_min"], "median": float(med), "max": res["depth_max"]} |
|
|
|
|
| |
|
|
| class _IsalCt: |
| """ISA-L ``isal_deflate`` called through ctypes from the python-isal extension's own shared object. |
| ctypes releases the GIL for the whole call (python-isal's ``compressobj.compress`` holds it: its pool |
| throughput saturated at ~1.2 GB/s on this host vs ~5 GB/s here with 8 threads). Same compressor, same |
| level semantics; the produced raw-deflate pieces are valid deflate streams (checked by a round trip at |
| load). Thread-local stream state + level buffer.""" |
|
|
| STRM = 1 << 19 |
| LBUF = 1 << 20 |
| Z_FINISH = Z_SYNC_FLUSH = None |
|
|
| def __init__(self): |
| import ctypes |
| from isal import _isal, isal_zlib |
| if int(_isal.ISAL_MAJOR_VERSION) != 2: |
| raise ImportError(f"isal_deflate layout unverified for ISA-L {_isal.ISAL_VERSION}") |
| self.ct = ctypes |
| L = ctypes.CDLL(isal_zlib.__file__) |
| L.isal_deflate_init.argtypes = [ctypes.c_void_p] |
| L.isal_deflate_init.restype = None |
| L.isal_deflate.argtypes = [ctypes.c_void_p] |
| L.isal_deflate.restype = ctypes.c_int |
| self.L = L |
| self.crc32 = isal_zlib.crc32 |
| self.tl = threading.local() |
|
|
| class Stream(ctypes.Structure): |
| _fields_ = [("next_in", ctypes.c_void_p), ("avail_in", ctypes.c_uint32), ("total_in", ctypes.c_uint32), |
| ("next_out", ctypes.c_void_p), ("avail_out", ctypes.c_uint32), ("total_out", ctypes.c_uint32), |
| ("hufftables", ctypes.c_void_p), ("level", ctypes.c_uint32), |
| ("level_buf_size", ctypes.c_uint32), ("level_buf", ctypes.c_void_p), |
| ("end_of_stream", ctypes.c_uint16), ("flush", ctypes.c_uint16), |
| ("gzip_flag", ctypes.c_uint16), ("hist_bits", ctypes.c_uint16), |
| ("internal_state", ctypes.c_uint8 * (self.STRM - 64))] |
| assert Stream.end_of_stream.offset == 56 and Stream.internal_state.offset == 64 |
| self.Stream = Stream |
| rng = np.random.default_rng(0) |
| data = (rng.standard_normal(300_000).astype(np.float32) * 10).view(np.uint8) |
| for level in (0, 1, 2, 3): |
| pieces = [memoryview(data)[i:i + 400_000] for i in range(0, data.size, 400_000)] |
| comp = b"".join(self.deflate(p, level, j == len(pieces) - 1) for j, p in enumerate(pieces)) |
| d = zlib.decompressobj(-15) |
| if d.decompress(comp) != data.tobytes() or not d.eof: |
| raise ImportError("isal_deflate (ctypes) self test failed") |
|
|
| @staticmethod |
| def bound(n: int) -> int: |
| return n + n // 2 + 65536 |
|
|
| def deflate(self, data, level: int, last: bool, out: Optional[np.ndarray] = None) -> np.ndarray: |
| ct, tl = self.ct, self.tl |
| if not hasattr(tl, "s"): |
| tl.s = self.Stream() |
| tl.lbuf = (ct.c_uint8 * self.LBUF)() |
| s = tl.s |
| self.L.isal_deflate_init(ct.byref(s)) |
| src = np.frombuffer(data, np.uint8) if len(data) else np.zeros(1, np.uint8)[:0] |
| n = src.size |
| cap = self.bound(n) |
| if out is None: |
| out = np.empty(cap, np.uint8) |
| assert out.size >= cap and out.flags.c_contiguous |
| s.next_in = src.ctypes.data if n else None |
| s.avail_in = n |
| s.next_out = out.ctypes.data |
| s.avail_out = cap |
| s.level = level |
| s.level_buf = ct.addressof(tl.lbuf) if level > 0 else None |
| s.level_buf_size = self.LBUF if level > 0 else 0 |
| s.end_of_stream = 1 if last else 0 |
| s.flush = 1 |
| s.gzip_flag = 0 |
| s.hist_bits = 0 |
| rc = self.L.isal_deflate(ct.byref(s)) |
| if rc != 0 or s.avail_in != 0 or s.avail_out == 0: |
| raise RuntimeError(f"isal_deflate rc={rc} avail_in={s.avail_in} avail_out={s.avail_out}") |
| return out[:cap - s.avail_out] |
|
|
|
|
| _isalct_obj = None |
|
|
|
|
| def _isalct(): |
| global _isalct_obj |
| if _isalct_obj is None: |
| _isalct_obj = _IsalCt() |
| return _isalct_obj |
|
|
|
|
| def _deflate_backend(): |
| """(module, default level): ISA-L through ctypes (GIL released; ``isalct``), python-isal's |
| ``isal_zlib`` (``isal``), else stdlib zlib. Override with MOGE_NPZ_DEFLATE=isalct|isal|zlib and |
| MOGE_NPZ_LEVEL.""" |
| want = os.environ.get("MOGE_NPZ_DEFLATE", "auto").strip().lower() |
| if want in ("auto", "isalct"): |
| try: |
| return _isalct(), 2 |
| except Exception: |
| if want == "isalct": |
| raise |
| if want in ("auto", "isal", "isalct"): |
| try: |
| from isal import isal_zlib |
| return isal_zlib, 2 |
| except ImportError: |
| if want == "isal": |
| raise |
| return zlib, 1 |
|
|
|
|
| def deflate_backend_name() -> str: |
| mod, level = _deflate_backend() |
| name = "isal (ctypes)" if isinstance(mod, _IsalCt) else mod.__name__ |
| return f"{name} level {int(os.environ.get('MOGE_NPZ_LEVEL', str(level)))}" |
|
|
|
|
| def _inflate_backend(): |
| """isal_zlib (python-isal) when installed, else zlib: decompress + crc32 for the request decode.""" |
| try: |
| from isal import isal_zlib |
| return isal_zlib |
| except ImportError: |
| return zlib |
|
|
|
|
| def _deflate_piece(mod, data, level: int, last: bool, out=None): |
| if isinstance(mod, _IsalCt): |
| return mod.deflate(data, level, last, out) |
| co = mod.compressobj(level, mod.DEFLATED, -15) |
| return co.compress(data) + co.flush(mod.Z_FINISH if last else mod.Z_SYNC_FLUSH) |
|
|
|
|
| def _crc(mod, pieces) -> int: |
| c = 0 |
| for p in pieces: |
| c = mod.crc32(p, c) |
| return c |
|
|
|
|
| def npz_segments(arrays: Dict[str, np.ndarray], level: Optional[int] = None) -> list: |
| """A .npz (zip, deflate) equivalent to ``np.savez_compressed(**arrays)`` for ``np.load``, as the list |
| of byte segments whose concatenation is the file (``npz_bytes`` joins them). The segments may view a |
| shared scratch buffer: hold ``_scratch_lock`` until they are consumed (``npz_bytes`` / ``npz_b64_bytes``).""" |
| mod, default_level = _deflate_backend() |
| if level is None: |
| level = int(os.environ.get("MOGE_NPZ_LEVEL", str(default_level))) |
| pool = _get_pool() |
| ct = isinstance(mod, _IsalCt) |
| plan = [] |
| for name, a in arrays.items(): |
| a = np.asanyarray(a) |
| if not a.flags.c_contiguous: |
| a = a.copy(order="C") |
| hdr = io.BytesIO() |
| np.lib.format.write_array_header_1_0(hdr, np.lib.format.header_data_from_array_1_0(a)) |
| raw = memoryview(a.reshape(-1).view(np.uint8)) if a.size else memoryview(b"") |
| plan.append((name, [hdr.getvalue()] + [raw[i:i + _CHUNK] for i in range(0, len(raw), _CHUNK)])) |
| |
| |
| scratch = _scratch(sum(mod.bound(len(p)) for _, ps in plan for p in ps)) if ct else None |
| soff = 0 |
| members = [] |
| for name, pieces in plan: |
| futs = [] |
| for j, p in enumerate(pieces): |
| dst = None |
| if ct: |
| cap = mod.bound(len(p)) |
| dst, soff = scratch[soff:soff + cap], soff + cap |
| futs.append(pool.submit(_deflate_piece, mod, p, level, j == len(pieces) - 1, dst)) |
| members.append((name + ".npy", sum(len(p) for p in pieces), pool.submit(_crc, mod, pieces), futs)) |
| out, central = [], [] |
| off = 0 |
| for fname, usize, crc_f, futs in members: |
| comp = [f.result() for f in futs] |
| csize = sum(memoryview(c).nbytes for c in comp) |
| crc = crc_f.result() |
| nb = fname.encode() |
| |
| lh = struct.pack("<IHHHHHIIIHH", 0x04034B50, 20, 0, 8, 0, 0x21, crc, csize, usize, len(nb), 0) + nb |
| central.append(struct.pack("<IHHHHHHIIIHHHHHII", 0x02014B50, 20, 20, 0, 8, 0, 0x21, crc, csize, usize, |
| len(nb), 0, 0, 0, 0, 0o600 << 16, off) + nb) |
| out.append(lh) |
| out.extend(comp) |
| off += len(lh) + csize |
| cd = b"".join(central) |
| out.append(cd) |
| out.append(struct.pack("<IHHHHIIH", 0x06054B50, 0, 0, len(central), len(central), len(cd), off, 0)) |
| return out |
|
|
|
|
| _scratch_buf = np.empty(0, np.uint8) |
| _scratch_lock = threading.RLock() |
|
|
|
|
| def _scratch(n: int) -> np.ndarray: |
| """Persistent output buffer of the parallel deflate (callers hold ``_scratch_lock`` until the |
| segments it backs are consumed).""" |
| global _scratch_buf |
| if _scratch_buf.size < n: |
| _scratch_buf = np.empty(n, np.uint8) |
| return _scratch_buf |
|
|
|
|
| def npz_bytes(arrays: Dict[str, np.ndarray], level: Optional[int] = None) -> bytes: |
| with _scratch_lock: |
| return b"".join(npz_segments(arrays, level)) |
|
|
|
|
| class Spliced: |
| """A base64 value written at ``buf[start:end]`` of a larger buffer with ``start`` bytes of headroom in |
| front and ``buf.size - end`` behind (the JSON body is assembled around it in place).""" |
|
|
| def __init__(self, buf: np.ndarray, start: int, end: int): |
| self.buf, self.start, self.end = buf, start, end |
|
|
| def tobytes(self) -> bytes: |
| return self.buf[self.start:self.end].tobytes() |
|
|
|
|
| _body_bufs: list = [] |
| _BODY_POOL = int(os.environ.get("MOGE_BODY_POOL", "3")) |
|
|
|
|
| def _body_buffer(n: int) -> np.ndarray: |
| """uint8 [n] for a response body: a view of a pooled buffer that no live response references any |
| more (a 39 MB fresh allocation is mmap'd and page-faults on first touch: ~1.7 ms of the base64 pass). |
| A buffer is reused only when its refcount shows no view of it alive (bodies are memoryviews of it).""" |
| import sys |
| with _pool_lock: |
| for a in _body_bufs: |
| if a.size >= n and sys.getrefcount(a) <= 3: |
| return a[:n] |
| a = np.empty(max(n + n // 16, 1 << 20), np.uint8) |
| if len(_body_bufs) < _BODY_POOL: |
| _body_bufs.append(a) |
| else: |
| for i in range(len(_body_bufs)): |
| if sys.getrefcount(_body_bufs[i]) <= 2: |
| _body_bufs[i] = a |
| break |
| return a[:n] |
|
|
|
|
| def b64_segments(segs, headroom: int = 0): |
| """``b64(b"".join(segs))`` without the join (multithreaded C++ over the segment list). With |
| ``headroom`` > 0 the result is a ``Spliced`` with that many free bytes before and after the value.""" |
| import ctypes |
| views = [np.frombuffer(memoryview(x).cast("B"), np.uint8) for x in segs] |
| views = [v for v in views if v.size] |
| n = sum(v.size for v in views) |
| m = 4 * ((n + 2) // 3) |
| res = _body_buffer(m + 2 * headroom) if headroom else np.empty(m, np.uint8) |
| if n: |
| ptrs = (ctypes.c_void_p * len(views))(*[v.ctypes.data for v in views]) |
| lens = (ctypes.c_int64 * len(views))(*[v.size for v in views]) |
| host_ext.lib().moge_b64_segs(ptrs, lens, len(views), res.ctypes.data + headroom, nthreads()) |
| return Spliced(res, headroom, headroom + m) if headroom else res |
|
|
|
|
| def b64(data) -> np.ndarray: |
| """base64 (standard alphabet, padded) as an ASCII uint8 array; multithreaded C++ (the output buffer is |
| np.empty: a zero-filled bytearray of ~48 MB alone costs ~20 ms).""" |
| mv = memoryview(data).cast("B") |
| n = len(mv) |
| res = np.empty(4 * ((n + 2) // 3), np.uint8) |
| if n: |
| src = np.frombuffer(mv, np.uint8) |
| host_ext.lib().moge_b64(src.ctypes.data, n, res.ctypes.data, nthreads()) |
| return res |
|
|
|
|
| def npz_b64_bytes(arrays: Dict[str, np.ndarray], level: Optional[int] = None, headroom: int = 0): |
| with _scratch_lock: |
| return b64_segments(npz_segments(arrays, level), headroom) |
|
|
|
|
| def npz_b64(arrays: Dict[str, np.ndarray], level: Optional[int] = None) -> str: |
| return npz_b64_bytes(arrays, level).tobytes().decode("ascii") |
|
|
|
|
| def tune_malloc() -> str: |
| """MOGE_MALLOPT=1 (default): keep freed multi-MB blocks in the glibc heap instead of mmap/munmap per |
| allocation (M_MMAP_THRESHOLD 32 MB = glibc's max, M_TRIM_THRESHOLD 256 MB): the served path allocates |
| and frees tens of MB per request (decoded image, response arrays, npz / base64 buffers) and every fresh |
| mmap'd page is a page fault on first touch (~0.25 us per 4 KB). Only allocation placement changes.""" |
| if os.environ.get("MOGE_MALLOPT", "1").strip().lower() in ("0", "false", "no", "off"): |
| return "off" |
| try: |
| libc = ctypes.CDLL("libc.so.6") |
| ok = libc.mallopt(-3, 32 << 20) == 1 and libc.mallopt(-1, 256 << 20) == 1 |
| return "mmap 32 MB / trim 256 MB" if ok else "mallopt refused" |
| except OSError as e: |
| return f"unavailable ({e})" |
|
|
|
|
| |
|
|
| _PNG_SIG = b"\x89PNG\r\n\x1a\n" |
|
|
|
|
| def decode_png_rgb(raw: bytes) -> Optional[np.ndarray]: |
| """uint8 [H, W, 3] of an 8-bit, non-interlaced RGB / RGBA PNG (== PIL ``open().convert('RGB')``: |
| PNG is lossless and the alpha channel is dropped), or None for anything else (the caller then uses |
| PIL). Chunk CRCs are verified; IDAT is inflated with ISA-L (zlib fallback), unfiltered in C++.""" |
| if bytes(memoryview(raw)[:8]) != _PNG_SIG: |
| return None |
| try: |
| from PIL import Image |
| mod = _inflate_backend() |
| p, n = 8, len(raw) |
| mv = memoryview(raw) |
| ihdr, idat, idat_off, end = None, [], [], False |
| while p + 12 <= n: |
| ln, = struct.unpack_from(">I", raw, p) |
| typ = bytes(mv[p + 4:p + 8]) |
| if p + 12 + ln > n: |
| return None |
| data = mv[p + 8:p + 8 + ln] |
| crc, = struct.unpack_from(">I", raw, p + 8 + ln) |
| if mod.crc32(data, mod.crc32(typ)) & 0xFFFFFFFF != crc: |
| return None |
| if typ == b"IHDR": |
| ihdr = struct.unpack(">IIBBBBB", data) |
| elif typ == b"IDAT": |
| idat.append(data) |
| idat_off.append((p + 8, ln)) |
| elif typ == b"IEND": |
| end = True |
| break |
| elif typ == b"PLTE" or not (typ[0] & 0x20): |
| return None |
| p += 12 + ln |
| if ihdr is None or not idat or not end: |
| return None |
| W, H, depth, ctype, comp, filt, interlace = ihdr |
| if depth != 8 or ctype not in (2, 6) or comp or filt or interlace or W < 1 or H < 1: |
| return None |
| if W * H > (Image.MAX_IMAGE_PIXELS or 1 << 62): |
| return None |
| bpp = 3 if ctype == 2 else 4 |
| rowbytes = W * bpp |
| out = np.empty((H, W, 3), np.uint8) |
| isal = _isal_inflate_fns() |
| if isal is not None: |
| rc = _png_isal_ct(isal, raw, idat_off, H, rowbytes, bpp, out) |
| return out if rc == 0 else None |
| if _png_unfilter_pipelined(mod, idat, H, rowbytes, bpp, out): |
| return out |
| return None |
| except Exception: |
| return None |
|
|
|
|
| _unf_pool: Optional[ThreadPoolExecutor] = None |
| _isal_fns = None |
| _png_tls = threading.local() |
|
|
|
|
| def _png_isal_ct(fns, raw, idat_off, H: int, rowbytes: int, bpp: int, out: np.ndarray) -> int: |
| """IDAT inflate (ISA-L, C++, no GIL) pipelined with the C++ unfilter: moge_png_inflate_unfilter.""" |
| L = host_ext.lib() |
| total = H * (rowbytes + 1) |
| z = getattr(_png_tls, "z", None) |
| if z is None or z.size < total: |
| z = _png_tls.z = np.empty(total, np.uint8) |
| base = np.frombuffer(raw, np.uint8).ctypes.data |
| k = len(idat_off) |
| ptrs = (ctypes.c_void_p * k)(*[base + o for o, _ in idat_off]) |
| lens = (ctypes.c_int64 * k)(*[ln for _, ln in idat_off]) |
| times = _png_times if _PNG_TIMES else None |
| return L.moge_png_inflate_unfilter(fns[0], fns[1], ptrs, lens, k, H, rowbytes, bpp, z.ctypes.data, |
| out.ctypes.data, times.ctypes.data if times is not None else None) |
|
|
|
|
| _PNG_TIMES = os.environ.get("MOGE_DECODE_LOG") == "1" |
| _png_times = np.zeros(2, np.float64) |
|
|
|
|
| def _isal_inflate_fns(): |
| """(isal_inflate_init, isal_inflate) addresses from python-isal's shared object for the C++ PNG decode, |
| or None (MOGE_PNG_ISALCT=0, no python-isal / host_ext, ISA-L other than 2.31, or the self test fails).""" |
| global _isal_fns |
| if _isal_fns is None: |
| _isal_fns = False |
| try: |
| if os.environ.get("MOGE_PNG_ISALCT", "1").strip().lower() in ("0", "false", "no", "off"): |
| raise ImportError("disabled") |
| from isal import _isal, isal_zlib |
| if not str(_isal.ISAL_VERSION).startswith("2.31."): |
| raise ImportError(f"inflate_state layout unverified for ISA-L {_isal.ISAL_VERSION}") |
| if not host_ext.enabled(): |
| raise ImportError("host_ext unavailable") |
| so = ctypes.CDLL(isal_zlib.__file__) |
| fns = (ctypes.cast(so.isal_inflate_init, ctypes.c_void_p).value, |
| ctypes.cast(so.isal_inflate, ctypes.c_void_p).value) |
| if not _isal_selftest(fns): |
| raise ImportError("self test failed") |
| _isal_fns = fns |
| except Exception: |
| _isal_fns = False |
| return _isal_fns or None |
|
|
|
|
| def _isal_selftest(fns) -> bool: |
| """All 5 filter types, RGB + RGBA, multi-chunk IDAT, vs zlib + the serial C++ unfilter; a corrupted |
| adler32 must be rejected.""" |
| rng = np.random.default_rng(0) |
| L = host_ext.lib() |
| for bpp in (3, 4): |
| H, W = 37, 53 |
| rowbytes = W * bpp |
| rows = rng.integers(0, 256, (H, rowbytes + 1), dtype=np.uint8) |
| rows[:, 0] = np.arange(H) % 5 |
| rows[5:9, 1:] = rows[4, 1:] |
| comp = zlib.compress(rows.tobytes(), 6) |
| ref = np.empty((H, W, 3), np.uint8) |
| if L.moge_png_unfilter(rows.ctypes.data, H, rowbytes, bpp, 3, ref.ctypes.data) != 0: |
| return False |
| cut = [0, 7, 100, len(comp) // 2, len(comp)] |
| raw = bytes(comp) |
| offs = [(cut[i], cut[i + 1] - cut[i]) for i in range(len(cut) - 1)] |
| out = np.empty((H, W, 3), np.uint8) |
| if _png_isal_ct(fns, raw, offs, H, rowbytes, bpp, out) != 0 or not np.array_equal(out, ref): |
| return False |
| bad = bytearray(comp) |
| bad[-1] ^= 1 |
| if _png_isal_ct(fns, bytes(bad), [(0, len(bad))], H, rowbytes, bpp, out) == 0: |
| return False |
| return True |
| _PNG_ROWS = 96 |
|
|
|
|
| def _png_unfilter_pipelined(mod, idat, H: int, rowbytes: int, bpp: int, out: np.ndarray) -> bool: |
| """Inflate the IDAT stream piecewise on this thread while a worker thread unfilters the rows that are |
| complete (ctypes releases the GIL); same bytes as inflating everything first and unfiltering after |
| (the unfilter is sequential over rows and keeps its previous row in its own state).""" |
| global _unf_pool |
| if _unf_pool is None: |
| _unf_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="moge-png") |
| L = host_ext.lib() |
| stride = rowbytes + 1 |
| total = H * stride |
| z = np.empty(total, np.uint8) |
| st = L.moge_unf_new(rowbytes, bpp, 3) |
| futs = [] |
| try: |
| d = mod.decompressobj() |
| have, done = 0, 0 |
| rows_out = out.reshape(H, -1) |
|
|
| def flush(upto: int): |
| nonlocal done |
| if upto > done: |
| futs.append(_unf_pool.submit(L.moge_unf_rows, st, z[done * stride:].ctypes.data, upto - done, |
| rows_out[done].ctypes.data)) |
| done = upto |
| for chunk in idat: |
| piece = d.decompress(chunk) |
| n = len(piece) |
| if have + n > total: |
| n = total - have |
| z[have:have + n] = np.frombuffer(piece, np.uint8, n) |
| have += n |
| if have // stride - done >= _PNG_ROWS: |
| flush(have // stride) |
| piece = d.flush() |
| n = min(len(piece), total - have) |
| z[have:have + n] = np.frombuffer(piece, np.uint8, n) |
| have += n |
| if have < total or not d.eof: |
| return False |
| flush(H) |
| return all(f.result() == 0 for f in futs) |
| finally: |
| for f in futs: |
| f.result() |
| L.moge_unf_free(st) |
|
|
|
|
| _PyUnicode_AsUTF8AndSize = None |
|
|
|
|
| def _ascii_view(s: str): |
| """(address, length) of an ASCII str's own UTF-8 data (== its compact ASCII storage, no copy).""" |
| global _PyUnicode_AsUTF8AndSize |
| if _PyUnicode_AsUTF8AndSize is None: |
| f = ctypes.pythonapi.PyUnicode_AsUTF8AndSize |
| f.argtypes = [ctypes.py_object, ctypes.POINTER(ctypes.c_ssize_t)] |
| f.restype = ctypes.c_void_p |
| _PyUnicode_AsUTF8AndSize = f |
| n = ctypes.c_ssize_t() |
| addr = _PyUnicode_AsUTF8AndSize(s, ctypes.byref(n)) |
| return addr, n.value |
|
|
|
|
| def b64decode_view(s): |
| """``b64decode`` into a per-thread persistent buffer: a memoryview valid until this thread's next call |
| (no 2-3 MB fresh allocations per request; an ASCII str is read in place). Same bytes as b64decode.""" |
| if isinstance(s, str) and s.isascii() and len(s) >= 4096 and host_ext.enabled(): |
| addr, n = _ascii_view(s) |
| if addr and n == len(s): |
| buf = getattr(_png_tls, "b64", None) |
| need = n // 4 * 3 |
| if buf is None or buf.size < need: |
| buf = _png_tls.b64 = np.empty(max(need, 1 << 20), np.uint8) |
| k = host_ext.lib().moge_b64dec(addr, n, buf.ctypes.data, nthreads()) |
| if k >= 0: |
| return memoryview(buf)[:k] |
| return memoryview(b64decode(s)) |
|
|
|
|
| def b64decode(s) -> bytes: |
| """``base64.b64decode(s, validate=False)``: multithreaded C++ for a clean standard-alphabet string, |
| Python's decoder for anything else (which skips foreign characters / raises on bad padding).""" |
| import base64 |
| data = s.encode("ascii") if isinstance(s, str) and s.isascii() else None |
| if data is not None and len(data) >= 4096 and host_ext.enabled(): |
| buf = np.empty(len(data) // 4 * 3, np.uint8) |
| src = np.frombuffer(data, np.uint8) |
| n = host_ext.lib().moge_b64dec(src.ctypes.data, src.size, buf.ctypes.data, nthreads()) |
| if n >= 0: |
| return buf[:n].tobytes() |
| return base64.b64decode(s, validate=False) |
|
|
|
|
| def _deflate_parallel(mod, data: memoryview, level: int): |
| """Raw-deflate pieces of one stream (futures, in order).""" |
| pool = _get_pool() |
| pieces = [data[i:i + _CHUNK] for i in range(0, len(data), _CHUNK)] or [data] |
| return [pool.submit(_deflate_piece, mod, p, level, j == len(pieces) - 1) for j, p in enumerate(pieces)] |
|
|
|
|
| def png_bytes(arr: np.ndarray, level: Optional[int] = None) -> bytes: |
| """Lossless PNG of a uint8 [H,W] / [H,W,3] or uint16 [H,W] array (decodes to exactly ``arr``, like |
| PIL's ``Image.fromarray(arr).save(PNG)``): 'Up' filter on every row (vectorized), the zlib stream |
| deflated in parallel pieces (pigz construction), adler32 trailer.""" |
| mod, default_level = _deflate_backend() |
| if level is None: |
| level = int(os.environ.get("MOGE_NPZ_LEVEL", str(default_level))) |
| if arr.dtype == np.uint16 and arr.ndim == 2: |
| depth, ctype, rows = 16, 0, np.ascontiguousarray(arr.astype(">u2")).view(np.uint8) |
| elif arr.dtype == np.uint8 and arr.ndim == 2: |
| depth, ctype, rows = 8, 0, np.ascontiguousarray(arr) |
| elif arr.dtype == np.uint8 and arr.ndim == 3 and arr.shape[2] == 3: |
| depth, ctype, rows = 8, 2, np.ascontiguousarray(arr).reshape(arr.shape[0], -1) |
| else: |
| raise ValueError(f"png_bytes: unsupported {arr.dtype} {arr.shape}") |
| H, W = arr.shape[:2] |
| rows = rows.reshape(H, -1) |
| filt = np.empty((H, rows.shape[1] + 1), np.uint8) |
| filt[:, 0] = 2 |
| filt[0, 1:] = rows[0] |
| np.subtract(rows[1:], rows[:-1], out=filt[1:, 1:]) |
| raw = memoryview(filt.reshape(-1)) |
| futs = _deflate_parallel(mod, raw, level) |
| adler = zlib.adler32(raw) |
| comp = b"".join(memoryview(f.result()).cast("B") for f in futs) |
|
|
| def chunk(typ: bytes, data: bytes) -> bytes: |
| return struct.pack(">I", len(data)) + typ + data + struct.pack(">I", zlib.crc32(data, zlib.crc32(typ)) & 0xFFFFFFFF) |
|
|
| ihdr = struct.pack(">IIBBBBB", W, H, depth, ctype, 0, 0, 0) |
| idat = b"\x78\x01" + comp + struct.pack(">I", adler & 0xFFFFFFFF) |
| return _PNG_SIG + chunk(b"IHDR", ihdr) + chunk(b"IDAT", idat) + chunk(b"IEND", b"") |
|
|
|
|
| def fit_to_canvas(rgb: np.ndarray, canvas: Tuple[int, int], fit: str) -> Tuple[torch.Tensor, Region]: |
| """``postprocess.fit_to_canvas`` on a uint8 [H, W, 3] array in one multithreaded C++ pass (Pillow's |
| BILINEAR resize arithmetic, paste on the gray canvas, / 255 as float32, planar NCHW output); |
| bit-identical (tests/test_serve_fastpath.py).""" |
| from tt_moge.server import postprocess as pp |
| cw, ch = canvas |
| H, W = rgb.shape[:2] |
| if fit == "stretch": |
| rw, rh, x0, y0 = cw, ch, 0, 0 |
| elif fit == "pad": |
| s = min(cw / W, ch / H) |
| rw, rh = max(1, round(W * s)), max(1, round(H * s)) |
| x0, y0 = (cw - rw) // 2, (ch - rh) // 2 |
| else: |
| raise ValueError(f"unknown fit mode {fit!r}") |
| rgb = np.ascontiguousarray(rgb, dtype=np.uint8) |
| out = torch.empty(1, 3, ch, cw, dtype=torch.float32) |
| pad = float(np.float32(pp.PAD_GRAY[0]) / np.float32(255.0)) |
| assert pp.PAD_GRAY[0] == pp.PAD_GRAY[1] == pp.PAD_GRAY[2] |
| rc = host_ext.lib().moge_fit_canvas(rgb.ctypes.data, H, W, rgb.strides[0], rw, rh, x0, y0, cw, ch, pad, |
| out.data_ptr(), nthreads()) |
| if rc != 0: |
| raise ValueError("moge_fit_canvas: bad geometry") |
| return out, (x0, y0, x0 + rw, y0 + rh) |
|
|
|
|
| def fit_to_canvas_u8(rgb: np.ndarray, canvas: Tuple[int, int], fit: str) -> Tuple[torch.Tensor, Region]: |
| """Round 12 (MOGE_SERVE_U8): ``fit_to_canvas`` without the float32 canvas: the uint8 canvas [1, Ch, Cw, 3] |
| (NHWC) whose u8 / 255 IS fit_to_canvas's float32 output (TtMoGe's input prep reads it that way: the same |
| model input bit for bit). A canvas-sized image is passed through without a copy.""" |
| from tt_moge.server import postprocess as pp |
| cw, ch = canvas |
| H, W = rgb.shape[:2] |
| if fit == "stretch": |
| rw, rh, x0, y0 = cw, ch, 0, 0 |
| elif fit == "pad": |
| s = min(cw / W, ch / H) |
| rw, rh = max(1, round(W * s)), max(1, round(H * s)) |
| x0, y0 = (cw - rw) // 2, (ch - rh) // 2 |
| else: |
| raise ValueError(f"unknown fit mode {fit!r}") |
| rgb = np.ascontiguousarray(rgb, dtype=np.uint8) |
| region = (x0, y0, x0 + rw, y0 + rh) |
| if (rw, rh, x0, y0) == (W, H, 0, 0) and (W, H) == (cw, ch): |
| return torch.from_numpy(rgb)[None], region |
| out = torch.empty(1, ch, cw, 3, dtype=torch.uint8) |
| assert pp.PAD_GRAY[0] == pp.PAD_GRAY[1] == pp.PAD_GRAY[2] |
| rc = host_ext.lib().moge_fit_canvas_u8(rgb.ctypes.data, H, W, rgb.strides[0], rw, rh, x0, y0, cw, ch, |
| float(pp.PAD_GRAY[0]), out.data_ptr(), nthreads()) |
| if rc != 0: |
| raise ValueError("moge_fit_canvas_u8: bad geometry") |
| return out, region |
|
|
