File size: 5,229 Bytes
12acbba | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Patch packing and canvas saving utilities."""
from pathlib import Path
from typing import Tuple, List, Dict, Any, Optional
import numpy as np
import cv2
# Pillow for reliable JPEG writing
try:
from PIL import Image # type: ignore
HAS_PIL = True
except Exception:
Image = None
HAS_PIL = False
def iter_blocks_in_raster(hb: int, wb: int):
"""Iterate 2x2 blocks in raster order of blocks.
hb/wb are patch-grid sizes (even).
Yields (bh, bw) in block-grid.
"""
for bh in range(hb // 2):
for bw in range(wb // 2):
yield bh, bw
def block_to_4_patches(bh: int, bw: int) -> List[Tuple[int, int]]:
"""Convert block coord to 4 patch coords in the required contiguous order.
Returns [(h0, w0), (h0, w0+1), (h0+1, w0), (h0+1, w0+1)]
"""
h0 = 2 * int(bh)
w0 = 2 * int(bw)
return [(h0, w0), (h0, w0 + 1), (h0 + 1, w0), (h0 + 1, w0 + 1)]
def extract_patch_rgb(frame_rgb: np.ndarray, ph: int, pw: int, patch: int = 16) -> np.ndarray:
"""Extract a single patch from RGB frame."""
p = int(patch)
y0 = int(ph) * p
x0 = int(pw) * p
return frame_rgb[y0:y0 + p, x0:x0 + p, :]
def pack_patches_to_canvases(
patches: np.ndarray,
hb: int,
wb: int,
patch: int,
placement_order: str = "block_raster",
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Pack patches into one or more full canvases.
Packing order is raster over 2x2 blocks, and within each block the 4 patches
are placed in the order: (0,0),(0,1),(1,0),(1,1). This guarantees that
`image.reshape(-1)` will have consecutive 4 tokens corresponding to a 2x2 block.
Args:
patches: uint8 array (N, patch, patch, 3) where N is multiple of hb*wb
hb: Number of patches in height direction
wb: Number of patches in width direction
patch: Patch size in pixels
placement_order: "block_raster" (default) or "wh_raster"
Returns:
images_rgb: uint8 (num_images, H, W, 3)
patch_position: int32 (N, 3) [img_idx, patch_h, patch_w] aligned 1-1 with patches
img_ptr: int32 (num_images+1,) prefix-sum boundaries (each image has hb*wb patches)
"""
hb = int(hb)
wb = int(wb)
p = int(patch)
S_full = hb * wb
placement_order = str(placement_order).lower().strip()
if patches.size == 0:
images = np.zeros((0, hb * p, wb * p, 3), dtype=np.uint8)
patch_pos = np.zeros((0, 3), dtype=np.int32)
img_ptr = np.zeros((1,), dtype=np.int32)
return images, patch_pos, img_ptr
assert patches.ndim == 4 and patches.shape[1] == p and patches.shape[2] == p and patches.shape[3] == 3
assert patches.shape[0] % S_full == 0, f"patches must be multiple of S_full={S_full}, got {patches.shape[0]}"
num_images = int(patches.shape[0] // S_full)
H = hb * p
W = wb * p
images = np.zeros((num_images, H, W, 3), dtype=np.uint8)
patch_pos = np.zeros((patches.shape[0], 3), dtype=np.int32)
idx = 0
for img_i in range(num_images):
if placement_order == "wh_raster":
# Raster order by patch position (h, w)
for ph in range(hb):
for pw in range(wb):
y0 = int(ph) * p
x0 = int(pw) * p
images[img_i, y0:y0 + p, x0:x0 + p, :] = patches[idx]
patch_pos[idx, :] = (int(img_i), int(ph), int(pw))
idx += 1
else:
# Default: block raster order (2x2 blocks)
for bh in range(hb // 2):
for bw in range(wb // 2):
coords = block_to_4_patches(bh, bw)
for (ph, pw) in coords:
y0 = int(ph) * p
x0 = int(pw) * p
images[img_i, y0:y0 + p, x0:x0 + p, :] = patches[idx]
patch_pos[idx, :] = (int(img_i), int(ph), int(pw))
idx += 1
img_ptr = (np.arange(0, num_images + 1, dtype=np.int32) * int(S_full)).astype(np.int32)
return images, patch_pos, img_ptr
def save_canvases_as_jpg(images_rgb: np.ndarray, out_dir: str, quality: int = 95) -> List[str]:
"""Save (num_images, H, W, 3) RGB uint8 canvases into JPEG files.
Returns list of written filenames (basenames).
"""
out: List[str] = []
out_p = Path(out_dir)
out_p.mkdir(parents=True, exist_ok=True)
if images_rgb is None or images_rgb.size == 0:
return out
q = int(max(1, min(100, int(quality))))
for i in range(int(images_rgb.shape[0])):
fn = f"canvas_{i:03d}.jpg"
fp = out_p / fn
arr = images_rgb[i]
# Prefer Pillow for JPEG reliability; fallback to OpenCV.
if HAS_PIL and Image is not None:
Image.fromarray(arr).save(str(fp), format="JPEG", quality=q, subsampling=0, optimize=True)
else:
bgr = arr[:, :, ::-1]
ok = cv2.imwrite(str(fp), bgr, [int(cv2.IMWRITE_JPEG_QUALITY), q])
if not ok:
raise RuntimeError("Failed to write JPEG. Please install pillow: pip install pillow")
out.append(fn)
return out
|