| |
| """ |
| S3DIS 超点 / 语义 GT 可视化（Mitsuba 3）。 |
| |
| 模式 superpoint_per_class：对 13 个语义类分别只画该类的点，按超点 id 上色，各输出一张 PNG。 |
| |
| 用法见 --help。 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import math |
| import os |
| import sys |
|
|
| _HERE = os.path.dirname(os.path.abspath(__file__)) |
| if _HERE not in sys.path: |
| sys.path.insert(0, _HERE) |
|
|
|
|
| def _pick_variant(): |
| """scalar_rgb 与 mi_objects 中 Transform 序列化兼容。""" |
| import mitsuba as mi |
|
|
| for v in ("scalar_rgb", "llvm_ad_rgb", "cuda_ad_rgb"): |
| try: |
| mi.set_variant(v) |
| return v |
| except Exception: |
| continue |
| raise RuntimeError("无法设置任何 Mitsuba variant") |
|
|
|
|
| _GOLDEN = 0.618033988749895 |
|
|
|
|
| def _hsv_to_rgb(h: float, s: float, v: float): |
| import colorsys |
|
|
| r, g, b = colorsys.hsv_to_rgb(h % 1.0, s, v) |
| return r, g, b |
|
|
|
|
| def _hash_color(uid: int): |
| h = (int(uid) * _GOLDEN) % 1.0 |
| return _hsv_to_rgb(h, 0.82, 0.96) |
|
|
|
|
| S3DIS_SEGMENT_RGB = [ |
| [0.90, 0.90, 0.92], |
| [0.72, 0.52, 0.30], |
| [0.42, 0.72, 0.48], |
| [0.35, 0.55, 0.85], |
| [0.95, 0.75, 0.25], |
| [0.45, 0.85, 0.95], |
| [0.92, 0.35, 0.35], |
| [0.95, 0.55, 0.20], |
| [0.75, 0.45, 0.92], |
| [0.55, 0.40, 0.88], |
| [0.50, 0.82, 0.45], |
| [0.88, 0.50, 0.65], |
| [0.55, 0.55, 0.58], |
| ] |
|
|
| S3DIS_CLASS_NAMES = [ |
| "ceiling", |
| "floor", |
| "wall", |
| "beam", |
| "column", |
| "window", |
| "door", |
| "table", |
| "chair", |
| "sofa", |
| "bookcase", |
| "board", |
| "clutter", |
| ] |
|
|
|
|
| def _class_color(cid: int, n_cls: int = 13): |
| i = int(cid) % max(n_cls, 1) |
| return tuple(S3DIS_SEGMENT_RGB[i]) |
|
|
|
|
| def _maybe_crop_png( |
| path: str, crop_wh: tuple[int, int] | None, center: bool = False |
| ) -> str | None: |
| if not crop_wh: |
| return None |
| try: |
| import time |
| from PIL import Image |
| except ImportError: |
| print("未安装 Pillow，跳过裁剪", file=sys.stderr) |
| return None |
| w, h = crop_wh |
| im = None |
| for _ in range(40): |
| try: |
| im = Image.open(path) |
| im.load() |
| break |
| except (OSError, Exception): |
| time.sleep(0.05) |
| if im is None: |
| print("裁剪失败：无法读取刚写入的 PNG", file=sys.stderr) |
| return None |
| im = im.convert("RGB") |
| cw, ch = min(w, im.size[0]), min(h, im.size[1]) |
| if center: |
| left = max(0, (im.size[0] - cw) // 2) |
| top = max(0, (im.size[1] - ch) // 2) |
| im = im.crop((left, top, left + cw, top + ch)) |
| else: |
| im = im.crop((0, 0, cw, ch)) |
| base, ext = os.path.splitext(path) |
| out = base + "_crop" + ext |
| im.save(out) |
| return out |
|
|
|
|
| def _render_one( |
| mi, |
| points: "np.ndarray", |
| label_for_color: "np.ndarray", |
| color_fn, |
| out_png: str, |
| sphere_radius: float, |
| film_size: int, |
| spp: int, |
| fov: float, |
| use_plastic: bool, |
| light_intensity: float, |
| fill_irradiance: float, |
| ): |
| """points: (N,3) 已归一化前原始坐标；label_for_color 与 points 等长。""" |
| import numpy as np |
|
|
| from lib_render import center_normalize |
| from mi_objects import MiFloor, MiScene, MiSensor, MiSoftlight, MiSphere |
|
|
| n = points.shape[0] |
| if n == 0: |
| return False |
|
|
| pts = center_normalize(points.astype(np.float64)) |
| pts[:, 2] += sphere_radius / 2.0 |
|
|
| lds_spp = max(16, 2 ** int(round(math.log(max(int(spp), 16), 2)))) |
| lds_spp = min(lds_spp, 1024) |
|
|
| scene = MiScene() |
| scene.add( |
| "sensor", |
| MiSensor( |
| origin=[2, 2, 2], |
| target=[0, 0, 0], |
| fov=int(round(fov)), |
| sample_count=lds_spp, |
| film_width=int(film_size), |
| film_height=int(film_size), |
| ), |
| ) |
| scene.add("floor", MiFloor(width=10, height=10, color=[0.96, 0.96, 0.98])) |
| scene.add( |
| "soft_light", |
| MiSoftlight( |
| origin=[-4, 4, 20], |
| target=[0, 0, 0], |
| intensity=float(light_intensity), |
| ), |
| ) |
|
|
| for i in range(n): |
| c = color_fn(int(label_for_color[i])) |
| clr = [float(c[0]), float(c[1]), float(c[2])] |
| scene.add( |
| f"sphere{i}", |
| MiSphere(pts[i].tolist(), float(sphere_radius), clr, plastic=use_plastic), |
| ) |
|
|
| scene_dict = scene.dict() |
| fi = float(fill_irradiance) |
| scene_dict["fill_sun"] = { |
| "type": "directional", |
| "direction": [0.35, -0.55, 0.75], |
| "irradiance": {"type": "rgb", "value": [fi, fi * 0.98, fi * 0.95]}, |
| } |
|
|
| scene_mi = mi.load_dict(scene_dict) |
| img = mi.render(scene_mi, spp=int(spp)) |
| out = os.path.abspath(out_png) |
| os.makedirs(os.path.dirname(out) or ".", exist_ok=True) |
| mi.util.write_bitmap(out, img) |
| print("MITSUBA_PNG", out) |
| return True |
|
|
|
|
| def main(): |
| _pick_variant() |
| import numpy as np |
| import mitsuba as mi |
|
|
| default_root = os.path.normpath( |
| os.path.join(_HERE, "..", "_work_biptv3", "pointcept_framework", "data", "s3dis_official") |
| ) |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--data_root", type=str, default=default_root) |
| ap.add_argument("--room", type=str, required=True) |
| ap.add_argument( |
| "--out_png", |
| type=str, |
| default=None, |
| help="单张输出路径（superpoint / segment 模式必填）", |
| ) |
| ap.add_argument( |
| "--out_dir", |
| type=str, |
| default=None, |
| help="superpoint_per_class 时输出目录，默认 outputs/superpoint_vis/<room>_per_class_sp/", |
| ) |
| ap.add_argument( |
| "--mode", |
| choices=("superpoint", "segment", "superpoint_per_class"), |
| default="superpoint", |
| ) |
| ap.add_argument("--max_points", type=int, default=40000) |
| ap.add_argument("--sphere_radius", type=float, default=0.008) |
| ap.add_argument("--film_size", type=int, default=2560) |
| ap.add_argument("--spp", type=int, default=256) |
| ap.add_argument("--fov", type=float, default=25.0) |
| ap.add_argument("--crop", type=str, default=None) |
| ap.add_argument("--crop_center", action="store_true") |
| ap.add_argument("--no_plastic_spheres", action="store_true") |
| ap.add_argument("--light_intensity", type=float, default=14.0) |
| ap.add_argument("--fill_irradiance", type=float, default=3.5) |
| ap.add_argument( |
| "--min_class_points", |
| type=int, |
| default=50, |
| help="superpoint_per_class：点数少于此的类跳过", |
| ) |
| args = ap.parse_args() |
|
|
| room = os.path.join(args.data_root, args.room) |
| coord = np.load(os.path.join(room, "coord.npy")) |
| n_all = coord.shape[0] |
|
|
| use_plastic = not args.no_plastic_spheres |
| room_tag = args.room.replace("/", "_") |
|
|
| if args.mode == "superpoint_per_class": |
| segment = np.load(os.path.join(room, "segment.npy")).reshape(-1).astype(np.int64) |
| superpoint = np.load(os.path.join(room, "superpoint.npy")).reshape(-1).astype(np.int64) |
| if len(segment) != n_all or len(superpoint) != n_all: |
| raise ValueError("segment/superpoint 与 coord 长度不一致") |
|
|
| out_dir = args.out_dir |
| if not out_dir: |
| out_dir = os.path.join(_HERE, "outputs", "superpoint_vis", f"{room_tag}_per_class_sp") |
| out_dir = os.path.abspath(out_dir) |
| os.makedirs(out_dir, exist_ok=True) |
|
|
| written = [] |
| for cls in range(13): |
| mask = segment == cls |
| cnt = int(mask.sum()) |
| if cnt < args.min_class_points: |
| print(f"skip cls{cls:02d} {S3DIS_CLASS_NAMES[cls]}: n={cnt} (<{args.min_class_points})") |
| continue |
| coord_c = coord[mask] |
| sp_c = superpoint[mask] |
| if cnt > args.max_points: |
| rng = np.random.default_rng(cls) |
| idx = rng.choice(cnt, size=args.max_points, replace=False) |
| coord_c = coord_c[idx] |
| sp_c = sp_c[idx] |
|
|
| name = S3DIS_CLASS_NAMES[cls] |
| out_png = os.path.join(out_dir, f"{room_tag}_cls{cls:02d}_{name}_superpoint.png") |
| ok = _render_one( |
| mi, |
| coord_c, |
| sp_c, |
| _hash_color, |
| out_png, |
| args.sphere_radius, |
| args.film_size, |
| args.spp, |
| args.fov, |
| use_plastic, |
| args.light_intensity, |
| args.fill_irradiance, |
| ) |
| if ok: |
| written.append(out_png) |
| if args.crop: |
| wh = args.crop.lower().replace(" ", "").split("x") |
| if len(wh) == 2: |
| cw, ch = int(wh[0]), int(wh[1]) |
| cpath = _maybe_crop_png(out_png, (cw, ch), center=args.crop_center) |
| if cpath: |
| print("MITSUBA_PNG_CROP", cpath) |
|
|
| print("PER_CLASS_DONE", len(written), "files in", out_dir) |
| return |
|
|
| if not args.out_png: |
| sys.exit("非 superpoint_per_class 模式必须指定 --out_png") |
|
|
| if args.mode == "superpoint": |
| labels = np.load(os.path.join(room, "superpoint.npy")).reshape(-1).astype(np.int64) |
| color_fn = _hash_color |
| else: |
| labels = np.load(os.path.join(room, "segment.npy")).reshape(-1).astype(np.int64) |
| color_fn = _class_color |
|
|
| if len(labels) != n_all: |
| raise ValueError("标签与 coord 长度不一致") |
|
|
| coord_u = coord |
| labels_u = labels |
| if n_all > args.max_points: |
| rng = np.random.default_rng(0) |
| idx = rng.choice(n_all, size=args.max_points, replace=False) |
| coord_u = coord[idx] |
| labels_u = labels[idx] |
|
|
| _render_one( |
| mi, |
| coord_u, |
| labels_u, |
| color_fn, |
| args.out_png, |
| args.sphere_radius, |
| args.film_size, |
| args.spp, |
| args.fov, |
| use_plastic, |
| args.light_intensity, |
| args.fill_irradiance, |
| ) |
|
|
| if args.crop: |
| wh = args.crop.lower().replace(" ", "").split("x") |
| if len(wh) == 2: |
| cw, ch = int(wh[0]), int(wh[1]) |
| cropped = _maybe_crop_png(args.out_png, (cw, ch), center=args.crop_center) |
| if cropped: |
| print("MITSUBA_PNG_CROP", cropped) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
