|
|
| """
|
|
|
|
|
| 模式
|
| ----
|
| - superpoint：整场景点云，按超点 id 上色
|
| - segment：整场景点云，按语义 GT 上色
|
| - superpoint_per_class：每个语义类单独一张图，仅该类点，按超点 id 上色
|
|
|
|
|
| ----
|
| cd /path/to/PAMI2026
|
| python superpoint_visualize_s3dis.py --room Area_1/office_1 --mode superpoint \\
|
| --out_png outputs/superpoint_vis/office_1_mitsuba_superpoint.png
|
|
|
| python superpoint_visualize_s3dis.py --room Area_1/office_1 --mode superpoint_per_class \\
|
| --out_dir outputs/superpoint_vis/Area_1_office_1_per_class_sp_hq --film_size 2560 --spp 256
|
| """
|
| from __future__ import annotations
|
|
|
| import argparse
|
| import math
|
| import os
|
| import sys
|
|
|
| import mitsuba as mi
|
| import numpy as np
|
|
|
| _HERE = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
| def _pick_variant() -> None:
|
| for v in ("scalar_rgb", "llvm_ad_rgb", "cuda_ad_rgb"):
|
| try:
|
| mi.set_variant(v)
|
| return
|
| 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 center_normalize(pcl: np.ndarray) -> np.ndarray:
|
| centroid = np.mean(pcl, axis=0)
|
| pcl_centered = pcl - centroid
|
| max_dist = np.max(pcl_centered)
|
| pcl_normalized = pcl_centered / (2 * max_dist)
|
| pcl_normalized[:, 2] += -0.5 - np.min(pcl_normalized[:, 2])
|
| return pcl_normalized.astype(np.float32)
|
|
|
|
|
| class MiScene(object):
|
| def __init__(self):
|
| self.scene = {
|
| "type": "scene",
|
| "integrator": {"type": "path", "max_depth": -1},
|
| }
|
|
|
| def dict(self):
|
| return self.scene
|
|
|
| def add(self, key, object):
|
| self.scene[key] = object.dict()
|
|
|
|
|
| class MiSensor(object):
|
| def __init__(
|
| self,
|
| origin: list,
|
| target: list,
|
| fov: int = 25,
|
| sample_count: int = 256,
|
| film_width: int = 540,
|
| film_height: int = 540,
|
| ):
|
| self.origin = origin
|
| self.target = target
|
| self.fov = fov
|
| self.sample_count = sample_count
|
| self.film_width = film_width
|
| self.film_height = film_height
|
|
|
| def dict(self):
|
| return {
|
| "type": "perspective",
|
| "fov": self.fov,
|
| "near_clip": 0.1,
|
| "far_clip": 100.0,
|
| "to_world": mi.Transform4f().look_at(
|
| origin=self.origin, target=self.target, up=[0, 0, 1]
|
| ),
|
| "sampler": {"type": "ldsampler", "sample_count": self.sample_count},
|
| "film": {
|
| "type": "hdrfilm",
|
| "width": self.film_width,
|
| "height": self.film_height,
|
| "rfilter": {"type": "gaussian"},
|
| },
|
| }
|
|
|
|
|
| class MiFloor(object):
|
| def __init__(self, width: int = 10, height: int = 10, color: list | None = None):
|
| self.width = width
|
| self.height = height
|
| self.color = color if color is not None else [1, 1, 1]
|
|
|
| def dict(self):
|
| return {
|
| "type": "rectangle",
|
| "to_world": mi.Transform4f()
|
| .scale([self.width, self.height, 1])
|
| .translate([0, 0, -0.5]),
|
| "bsdf": {
|
| "type": "roughplastic",
|
| "distribution": "ggx",
|
| "alpha": 0.05,
|
| "int_ior": 1.46,
|
| "diffuse_reflectance": {"type": "rgb", "value": self.color},
|
| },
|
| }
|
|
|
|
|
| class MiSoftlight(object):
|
| def __init__(
|
| self,
|
| origin: list,
|
| target: list,
|
| intensity: float = 12.0,
|
| radiance: list | None = None,
|
| width: int = 10,
|
| height: int = 10,
|
| ):
|
| self.origin = origin
|
| self.target = target
|
| self.width = width
|
| self.height = height
|
| if radiance is not None:
|
| self.radiance = [float(x) for x in radiance]
|
| else:
|
| v = float(intensity)
|
| self.radiance = [v, v, v]
|
|
|
| def dict(self):
|
| return {
|
| "type": "rectangle",
|
| "to_world": mi.Transform4f()
|
| .look_at(origin=self.origin, target=self.target, up=[0, 0, 1])
|
| .scale([self.width, self.height, 1]),
|
| "emitter": {
|
| "type": "area",
|
| "radiance": {"type": "rgb", "value": self.radiance},
|
| },
|
| }
|
|
|
|
|
| class MiSphere(object):
|
| def __init__(self, center: list, radius: float, color: list, plastic: bool = False):
|
| self.center = center
|
| self.radius = radius
|
| self.color = color
|
| self.plastic = plastic
|
|
|
| def dict(self):
|
| if self.plastic:
|
| bsdf = {
|
| "type": "plastic",
|
| "diffuse_reflectance": {"type": "rgb", "value": self.color},
|
| "specular_reflectance": {"type": "rgb", "value": [0.25, 0.25, 0.25]},
|
| "int_ior": 1.5,
|
| }
|
| else:
|
| bsdf = {
|
| "type": "diffuse",
|
| "reflectance": {"type": "rgb", "value": self.color},
|
| }
|
| return {
|
| "type": "sphere",
|
| "center": self.center,
|
| "radius": self.radius,
|
| "bsdf": bsdf,
|
| }
|
|
|
|
|
| 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("跳过裁剪", 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("裁剪失败", 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(
|
| 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,
|
| ) -> bool:
|
| 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()
|
|
|
| default_root = os.path.normpath(
|
| os.path.join(_HERE, "..", "_work_biptv3", "pointcept_framework", "data", "s3dis_official")
|
| )
|
|
|
| ap = argparse.ArgumentParser(description="S3DIS 超点superpoint/语义segment可视化")
|
| 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="输出路径",
|
| )
|
| ap.add_argument(
|
| "--out_dir",
|
| type=str,
|
| default=None,
|
| help="输出目录默认 outputs/superpoint_vis/<room>_per_class_sp/",
|
| )
|
| ap.add_argument(
|
| "--mode",
|
| choices=("superpoint", "segment", "superpoint_per_class"),
|
| default="superpoint",
|
| )
|
| ap.add_argument(
|
| "--label_npy",
|
| type=str,
|
| default=None,
|
| help="optional override for room/superpoint.npy in superpoint modes",
|
| )
|
| 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)
|
| label_path = args.label_npy or os.path.join(room, "superpoint.npy")
|
| superpoint = np.load(label_path).reshape(-1).astype(np.int64)
|
| print("LABEL_NPY", os.path.abspath(label_path))
|
| 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(
|
| 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":
|
| label_path = args.label_npy or os.path.join(room, "superpoint.npy")
|
| labels = np.load(label_path).reshape(-1).astype(np.int64)
|
| print("LABEL_NPY", os.path.abspath(label_path))
|
| 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(
|
| 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()
|
|
|
