"""Tile baker — extracts 128x128 RGBA8 tiles from Blender images and streams
them to Godot via shared memory + UDS signaling.
Flow:
1. Texture paint or image edit marks a dirty rect
2. TileBaker identifies affected tiles at the base mip level
3. Extracts RGBA8 pixel data, writes to SHM ring buffer
4. Sends "tile_ready:tex_id,tx,ty,mip,offset,size" over UDS
5. Generates parent mip tiles (2x2 downsample) and streams those too
6. Godot-side BlenderBridge reads the tile data and calls VirtualTexture.upload_tile()
"""
import struct
import numpy as np
from .shm_protocol import ShmRegion
try:
import bpy
except ImportError:
bpy = None
class TileBaker:
"""Bakes image tiles and streams them via SHM ring buffer."""
def __init__(
self,
shm: ShmRegion,
tile_size: int = 128,
max_mips: int = 8,
):
self._shm = shm
self._tile_size = tile_size
self._max_mips = max_mips
# Ring buffer region in SHM (top half)
self._ring_base = 0
self._ring_size = 0
self._ring_offset = 0
# Texture name → integer ID mapping (set by Godot after registration)
self._texture_ids: dict[str, int] = {}
# Pending tile_ready messages (caller drains via get_pending_messages)
self._pending_messages: list[str] = []
# ── Configuration ────────────────────────────────────────────────────
def init_ring_buffer(self, base_offset: int, ring_size: int):
"""Set up the tile ring buffer region in SHM.
Typically base_offset = shm_size // 2, ring_size = shm_size // 2.
This gives the tile ring the top half of shared memory.
"""
self._ring_base = base_offset
self._ring_size = ring_size
self._ring_offset = 0
def register_texture(self, name: str, texture_id: int):
"""Register a Blender image name → SVT texture ID mapping."""
self._texture_ids[name] = texture_id
def unregister_texture(self, name: str):
"""Remove a texture mapping."""
self._texture_ids.pop(name, None)
# ── Main baking entry points ─────────────────────────────────────────
def bake_dirty_region(
self,
image_name: str,
dirty_x: int,
dirty_y: int,
dirty_w: int,
dirty_h: int,
):
"""Bake all tiles that overlap a dirty rectangle.
Args:
image_name: Name of the Blender image
dirty_x, dirty_y: Top-left of dirty region in pixels
dirty_w, dirty_h: Size of dirty region in pixels
"""
tex_id = self._texture_ids.get(image_name)
if tex_id is None:
return
if bpy is None:
return
image = bpy.data.images.get(image_name)
if image is None:
return
pixels = self._get_image_pixels(image)
if pixels is None:
return
img_w, img_h = image.size[0], image.size[1]
self._bake_tiles_in_rect(pixels, img_w, img_h, tex_id, dirty_x, dirty_y, dirty_w, dirty_h)
def bake_full_image(self, image_name: str):
"""Bake all tiles for an entire image (initial sync or full invalidation)."""
tex_id = self._texture_ids.get(image_name)
if tex_id is None:
return
if bpy is None:
return
image = bpy.data.images.get(image_name)
if image is None:
return
pixels = self._get_image_pixels(image)
if pixels is None:
return
img_w, img_h = image.size[0], image.size[1]
self._bake_tiles_in_rect(pixels, img_w, img_h, tex_id, 0, 0, img_w, img_h)
def bake_from_array(
self,
pixels: np.ndarray,
img_w: int,
img_h: int,
tex_id: int,
dirty_x: int,
dirty_y: int,
dirty_w: int,
dirty_h: int,
):
"""Bake tiles from a raw pixel array (for testing without Blender).
Args:
pixels: RGBA float32 array shaped (img_h, img_w, 4)
img_w, img_h: Image dimensions
tex_id: Virtual texture ID
dirty_x, dirty_y, dirty_w, dirty_h: Dirty region
"""
self._bake_tiles_in_rect(pixels, img_w, img_h, tex_id, dirty_x, dirty_y, dirty_w, dirty_h)
# ── Pending message access ───────────────────────────────────────────
def get_pending_messages(self) -> list[str]:
"""Drain and return all pending tile_ready messages."""
msgs = self._pending_messages
self._pending_messages = []
return msgs
# ── Internal: pixel extraction ───────────────────────────────────────
def _get_image_pixels(self, image) -> np.ndarray | None:
"""Extract RGBA float32 pixels from a Blender image as (H, W, 4) array."""
w, h = image.size[0], image.size[1]
if w == 0 or h == 0:
return None
pixel_count = w * h * 4
flat = np.empty(pixel_count, dtype=np.float32)
image.pixels.foreach_get(flat)
return flat.reshape(h, w, 4)
# ── Internal: tile baking ────────────────────────────────────────────
def _bake_tiles_in_rect(
self,
pixels: np.ndarray,
img_w: int,
img_h: int,
tex_id: int,
dirty_x: int,
dirty_y: int,
dirty_w: int,
dirty_h: int,
):
"""Identify and bake all tiles overlapping the dirty rectangle."""
ts = self._tile_size
# Clamp dirty rect to image bounds
dirty_x = max(0, dirty_x)
dirty_y = max(0, dirty_y)
dirty_w = min(dirty_w, img_w - dirty_x)
dirty_h = min(dirty_h, img_h - dirty_y)
if dirty_w <= 0 or dirty_h <= 0:
return
# Find affected tile range
tx0 = dirty_x // ts
ty0 = dirty_y // ts
tx1 = (dirty_x + dirty_w - 1) // ts
ty1 = (dirty_y + dirty_h - 1) // ts
# Bake each affected tile at mip 0
baked_tiles = []
for ty in range(ty0, ty1 + 1):
for tx in range(tx0, tx1 + 1):
tile_data = self._extract_tile(pixels, img_w, img_h, tx, ty, ts)
offset = self._write_to_ring(tile_data)
if offset >= 0:
msg = f"tile_ready:{tex_id},{tx},{ty},0,{offset},{len(tile_data)}"
self._pending_messages.append(msg)
baked_tiles.append((tx, ty, tile_data))
# Generate parent mip tiles
self._generate_mips(baked_tiles, tex_id, img_w, img_h)
def _extract_tile(
self,
pixels: np.ndarray,
img_w: int,
img_h: int,
tx: int,
ty: int,
tile_size: int,
) -> bytes:
"""Extract a single tile as RGBA8 bytes.
Args:
pixels: (H, W, 4) float32 array
img_w, img_h: Image dimensions
tx, ty: Tile coordinates
tile_size: Tile edge size in pixels
Returns:
tile_size * tile_size * 4 bytes of RGBA8 data
"""
x0 = tx * tile_size
y0 = ty * tile_size
x1 = min(x0 + tile_size, img_w)
y1 = min(y0 + tile_size, img_h)
tile_f32 = pixels[y0:y1, x0:x1]
# Convert float32 [0,1] → uint8 [0,255]
tile_u8 = (np.clip(tile_f32, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
# Pad edge tiles to full tile_size
if tile_u8.shape[0] < tile_size or tile_u8.shape[1] < tile_size:
padded = np.zeros((tile_size, tile_size, 4), dtype=np.uint8)
padded[: tile_u8.shape[0], : tile_u8.shape[1]] = tile_u8
tile_u8 = padded
return tile_u8.tobytes()
# ── Internal: mip generation ─────────────────────────────────────────
def _generate_mips(
self,
base_tiles: list[tuple[int, int, bytes]],
tex_id: int,
img_w: int,
img_h: int,
):
"""Generate parent mip tiles by 2x2 downsampling.
For each base tile at (tx,ty), the parent at mip+1 is (tx//2, ty//2).
We collect unique parent tile coords, downsample from the baked data,
and stream those too.
"""
ts = self._tile_size
current_tiles = {} # (tx, ty) → RGBA8 uint8 array (ts, ts, 4)
for tx, ty, data in base_tiles:
arr = np.frombuffer(data, dtype=np.uint8).reshape(ts, ts, 4)
current_tiles[(tx, ty)] = arr
for mip in range(1, self._max_mips):
# Find unique parent tile coordinates
parents: dict[tuple[int, int], list[tuple[int, int]]] = {}
for tx, ty in current_tiles:
px, py = tx // 2, ty // 2
if (px, py) not in parents:
parents[(px, py)] = []
parents[(px, py)].append((tx, ty))
if not parents:
break
next_tiles = {}
for (px, py), children in parents.items():
# Build a 2x2 arrangement of child tiles (each ts x ts)
# Parent tile = downsample of the 2x2 block
combined = np.zeros((ts * 2, ts * 2, 4), dtype=np.uint16)
for cx, cy in children:
local_x = (cx % 2) * ts
local_y = (cy % 2) * ts
child_data = current_tiles[(cx, cy)]
combined[local_y : local_y + ts, local_x : local_x + ts] = child_data.astype(np.uint16)
# Vectorized box-filter downsample: average 2x2 blocks
# Reshape to (ts, 2, ts, 2, 4) then mean over axes (1, 3)
ds = combined.reshape(ts, 2, ts, 2, 4)
parent_tile = ds.mean(axis=(1, 3)).astype(np.uint8)
tile_data = parent_tile.tobytes()
offset = self._write_to_ring(tile_data)
if offset >= 0:
msg = f"tile_ready:{tex_id},{px},{py},{mip},{offset},{len(tile_data)}"
self._pending_messages.append(msg)
next_tiles[(px, py)] = parent_tile
if not next_tiles:
break
current_tiles = next_tiles
# ── Internal: ring buffer ────────────────────────────────────────────
def _write_to_ring(self, data: bytes) -> int:
"""Write data to the SHM ring buffer. Returns offset or -1 if failed."""
if not self._shm.valid:
return -1
size = len(data)
if size > self._ring_size:
return -1
# Wrap around if necessary
if self._ring_offset + size > self._ring_size:
self._ring_offset = 0
offset = self._ring_base + self._ring_offset
self._shm.write_blob(offset, data)
self._ring_offset += size
return offset