| import os |
| import re |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
|
|
| def _safe_name(name): |
| return re.sub(r"[^A-Za-z0-9_.-]+", "_", str(name)) |
|
|
|
|
| def _prepare_uv(point_uv): |
| uv = point_uv.detach().cpu() |
| if uv.dim() == 4 and uv.shape[1] == 1: |
| uv = uv[:, 0] |
| if uv.dim() != 3 or uv.shape[-1] != 2: |
| raise ValueError(f"Expected point_uv shape [V, 1, N, 2] or [V, N, 2], got {tuple(point_uv.shape)}") |
| return uv |
|
|
|
|
| def _prepare_depth(point_z, num_views): |
| if point_z is None: |
| return None |
|
|
| depth = point_z.detach().cpu().squeeze() |
| if depth.dim() != 2: |
| return None |
| if depth.shape[0] == num_views: |
| return depth |
| if depth.shape[1] == num_views: |
| return depth.t() |
| return None |
|
|
|
|
| def save_projection_debug(input_img, point_uv, point_z=None, out_dir="train_imgs/projection_debug", prefix="debug", view_names=None, max_points=30000): |
| os.makedirs(out_dir, exist_ok=True) |
|
|
| images = input_img.detach().clamp(0.0, 1.0).cpu() |
| uv = _prepare_uv(point_uv) |
| depth = _prepare_depth(point_z, uv.shape[0]) |
|
|
| num_views = min(images.shape[0], uv.shape[0]) |
| summary = [] |
|
|
| for view_idx in range(num_views): |
| image = images[view_idx] |
| if image.shape[0] == 1: |
| image = image.repeat(3, 1, 1) |
| image_np = (image[:3].permute(1, 2, 0).numpy() * 255.0).astype(np.uint8) |
| height, width = image_np.shape[:2] |
|
|
| view_uv = uv[view_idx] |
| finite = torch.isfinite(view_uv).all(dim=-1) |
| in_frame = ( |
| finite |
| & (view_uv[:, 0] >= -1.0) |
| & (view_uv[:, 0] <= 1.0) |
| & (view_uv[:, 1] >= -1.0) |
| & (view_uv[:, 1] <= 1.0) |
| ) |
| valid_indices = torch.nonzero(in_frame, as_tuple=False).flatten() |
| total_points = view_uv.shape[0] |
| valid_points = valid_indices.numel() |
|
|
| if valid_points > max_points: |
| sample_ids = torch.linspace(0, valid_points - 1, max_points).long() |
| valid_indices = valid_indices[sample_ids] |
|
|
| sampled_uv = view_uv[valid_indices] |
| x = torch.round((sampled_uv[:, 0] + 1.0) * 0.5 * (width - 1)).long().numpy() |
| y = torch.round((sampled_uv[:, 1] + 1.0) * 0.5 * (height - 1)).long().numpy() |
|
|
| overlay = image_np.copy() |
| if depth is not None: |
| depth_values = depth[view_idx][valid_indices].clamp(0.0, 1.0).numpy() |
| colors = np.stack( |
| [ |
| (255.0 * depth_values).astype(np.uint8), |
| np.full_like(depth_values, 64.0, dtype=np.float32).astype(np.uint8), |
| (255.0 * (1.0 - depth_values)).astype(np.uint8), |
| ], |
| axis=1, |
| ) |
| else: |
| colors = np.tile(np.array([[255, 32, 32]], dtype=np.uint8), (len(x), 1)) |
|
|
| for dy in (-1, 0, 1): |
| yy = np.clip(y + dy, 0, height - 1) |
| for dx in (-1, 0, 1): |
| xx = np.clip(x + dx, 0, width - 1) |
| overlay[yy, xx] = colors |
|
|
| blended = (0.65 * image_np + 0.35 * overlay).astype(np.uint8) |
| view_name = view_names[view_idx] if view_names and view_idx < len(view_names) else f"view{view_idx}" |
| filename = f"{prefix}_{view_idx:02d}_{_safe_name(view_name)}_projection.png" |
| Image.fromarray(blended).save(os.path.join(out_dir, filename)) |
| summary.append(f"{view_idx} {view_name}: {valid_points}/{total_points} points inside image") |
|
|
| with open(os.path.join(out_dir, f"{prefix}_projection_summary.txt"), "w") as f: |
| f.write("\n".join(summary)) |
|
|
