mirror of
https://github.com/unanmed/ginka-generator.git
synced 2026-08-14 18:12:28 +08:00
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
from collections import deque
|
|
import numpy as np
|
|
import torch
|
|
|
|
DIST_MAX_BUCKET = 12 # 距离分桶上限
|
|
DIST_VOCAB = DIST_MAX_BUCKET + 1 # 0-12 共 13 个桶
|
|
|
|
def compute_distance_field(map_matrix: np.ndarray) -> np.ndarray:
|
|
# map_matrix: [H, W] 整数矩阵
|
|
# 返回: [H * W] 分桶后的整数距离，值域 [0, DIST_MAX_BUCKET]
|
|
wall_mask = (map_matrix == 1)
|
|
h, w = map_matrix.shape
|
|
dist = np.full((h, w), h + w, dtype=np.int32)
|
|
dist[wall_mask] = 0
|
|
|
|
queue = deque()
|
|
for i in range(h):
|
|
for j in range(w):
|
|
if wall_mask[i, j]:
|
|
queue.append((i, j))
|
|
|
|
directions = [(0, 1), (0, -1), (1, 0), (-1, 0)]
|
|
while queue:
|
|
i, j = queue.popleft()
|
|
d = dist[i, j]
|
|
for di, dj in directions:
|
|
ni, nj = i + di, j + dj
|
|
if 0 <= ni < h and 0 <= nj < w and dist[ni, nj] > d + 1:
|
|
dist[ni, nj] = d + 1
|
|
queue.append((ni, nj))
|
|
|
|
dist = np.clip(dist, 0, DIST_MAX_BUCKET)
|
|
return dist.flatten().astype(np.int64)
|
|
|
|
def compute_distance_field_tensor(map_tensor: torch.Tensor) -> torch.Tensor:
|
|
# map_tensor: [B, H, W] 或 [1, H, W]
|
|
batch = []
|
|
for i in range(map_tensor.shape[0]):
|
|
d = compute_distance_field(map_tensor[i].cpu().numpy())
|
|
batch.append(torch.from_numpy(d))
|
|
return torch.stack(batch, dim=0).to(map_tensor.device)
|