optimize.py

"""
Two-stage optimization:
  1. SDF learning from point cloud
  2. Vertex generation + Delaunay meshing

All pure PyTorch, no compiled extensions.
"""
import os
import math
import time
import torch
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm

from .sdfnet import SDFNetwork
from .vgnet import VGNetwork
from . import losses as loss_utils
from . import meshing as mesh_utils
from .io_utils import (
    load_pointcloud,
    normalize_pointcloud,
    denormalize_pointcloud,
    estimate_normals,
    fps_sample,
    build_sigma_knn,
    save_mesh_ply,
    save_mesh_obj,
)


class Runner:
    def __init__(self,
                 pointcloud_path,
                 out_dir='./output',
                 device='cpu',
                 sdf_iters=20_000,
                 vg_iters=8_000,
                 sdf_lr=1e-3,
                 vg_lr=1e-3,
                 sdf_batch=5_000,
                 vg_batch=3_400,
                 vertices_size=3_400,
                 update_size=5,
                 update_ratio=1.2,
                 k_samples=21,
                 multires=8,
                 queries_size=1_000_000,
                 surface_queries=200_000,
                 project_sdf_level=0.0,
                 save_freq=2_000,
                 loss_weights_sdf=None,
                 loss_weights_vg=None,
                 ):
        self.device = torch.device(device)
        self.out_dir = out_dir
        os.makedirs(out_dir, exist_ok=True)

        # Load & normalize point cloud
        print("Loading point cloud...")
        raw_pts = load_pointcloud(pointcloud_path)
        self.raw_pts = raw_pts
        self.points, self.loc, self.scale = normalize_pointcloud(raw_pts)
        print(f"  Points: {len(self.points)}  Scale: {self.scale:.4f}")

        # Preprocess: build sigma + query samples for SDF training
        print("Preprocessing queries for SDF training...")
        self._preprocess_sdf_queries(queries_size)

        # Config
        self.sdf_iters = sdf_iters
        self.vg_iters = vg_iters
        self.sdf_lr = sdf_lr
        self.vg_lr = vg_lr
        self.sdf_batch = sdf_batch
        self.vg_batch = vg_batch
        self.vertices_size = vertices_size
        self.update_size = update_size
        self.update_ratio = update_ratio
        self.k_samples = k_samples
        self.project_sdf_level = project_sdf_level
        self.save_freq = save_freq

        self.loss_weights_sdf = loss_weights_sdf or [1.0, 0.1, 0.001, 0.0]
        self.loss_weights_vg = loss_weights_vg or [100.0, 1.0, 1.0, 1.0, 100.0]

        # Networks
        self.sdf_net = SDFNetwork(
            d_in=3, d_out=1, d_hidden=256, n_layers=8,
            skip_in=(4,), multires=multires,
            bias=0.5, scale=1.0,
            geometric_init=True, weight_norm=True,
        ).to(self.device)
        self.vg_net = VGNetwork(
            d_in=3, d_out=3, d_hidden=256, n_layers=8,
            skip_in=(4,), multires=multires,
            scale=1.0, geometric_init=True, weight_norm=True,
        ).to(self.device)

        self.sdf_optimizer = torch.optim.Adam(self.sdf_net.parameters(), lr=self.sdf_lr)
        self.vg_optimizer = torch.optim.Adam(self.vg_net.parameters(), lr=self.vg_lr)

        self.iter_step = 0

    # ------------------------------------------------------------------
    # SDF preprocessing
    # ------------------------------------------------------------------
    def _preprocess_sdf_queries(self, queries_size):
        pts = self.points
        point_num = len(pts)
        point_num_gt = (point_num // 60) * 60
        if point_num_gt == 0:
            point_num_gt = point_num
        query_each = max(queries_size // point_num_gt, 1)

        # subsample to ~1/60
        if point_num > point_num_gt:
            idx = np.random.choice(point_num, point_num_gt, replace=False)
        else:
            idx = np.arange(point_num)
        subsample = pts[idx]

        sigmas = build_sigma_knn(subsample, k=min(51, len(subsample)))

        sample = []
        sample_near = []
        scale = 0.25 * np.sqrt(max(point_num_gt, 1) / 20000.0)
        for _ in range(query_each):
            tt = subsample + scale * sigmas[:, None] * np.random.normal(0.0, 1.0, size=subsample.shape)
            sample.append(tt)
            sample_near.append(subsample)

        sample = np.concatenate(sample, axis=0).astype(np.float32)
        sample_near = np.concatenate(sample_near, axis=0).astype(np.float32)
        n_uniform = max(sample.shape[0] // 10, 1)
        sample_uniform = 1.1 * (np.random.rand(n_uniform, 3).astype(np.float32) - 0.5)
        sample_uniform_near = subsample[np.random.choice(len(subsample), n_uniform, replace=True)]

        self.sample = torch.from_numpy(sample).to(self.device)
        self.sample_near = torch.from_numpy(sample_near).to(self.device)
        self.sample_uniform = torch.from_numpy(sample_uniform).to(self.device)
        self.sample_uniform_near = torch.from_numpy(sample_uniform_near).to(self.device)
        self.point_gt = torch.from_numpy(subsample).to(self.device)
        self.surface_queries_size = min(200_000, len(subsample))

        # bbox
        self.bbox_min = subsample.min(axis=0) - 0.05
        self.bbox_max = subsample.max(axis=0) + 0.05

    # ------------------------------------------------------------------
    # SDF stage
    # ------------------------------------------------------------------
    def train_sdf(self):
        print("\n=== Stage 1: SDF Learning ===")
        self.sdf_net.train()
        pbar = tqdm(range(self.sdf_iters), desc="SDF")
        for iter_i in pbar:
            self.update_lr(self.sdf_optimizer, iter_i, self.sdf_iters, self.sdf_lr, warm_up_end=1000)

            # Sample batch
            n_near = self.sdf_batch
            idx_near = np.random.choice(len(self.sample), n_near, replace=False)
            idx_uniform = np.random.choice(len(self.sample_uniform), max(n_near // 2, 1), replace=False)

            sample_near = self.sample[idx_near]
            points_near = self.sample_near[idx_near]
            sample_uniform = self.sample_uniform[idx_uniform]
            points_uniform = self.sample_uniform_near[idx_uniform]

            samples = torch.cat([sample_near, sample_uniform], dim=0)
            gradients_samples, sdf_samples = self.sdf_net.gradient(samples)
            gradients_samples_norm = F.normalize(gradients_samples, dim=-1)
            samples_moved = samples - gradients_samples_norm * sdf_samples

            # Gradient consistency
            move_pos = samples_moved.detach()
            grad_moved, _ = self.sdf_net.gradient(move_pos)
            grad_moved_norm = F.normalize(grad_moved, dim=-1)
            loss_grad_consis = (1.0 - F.cosine_similarity(grad_moved_norm, gradients_samples_norm, dim=-1)).mean()

            points = torch.cat([points_near, points_uniform], dim=0)
            sdf_points = self.sdf_net.sdf(points)

            loss_pull = torch.linalg.norm((points - samples_moved), ord=2, dim=-1).mean()
            loss_sdf = torch.abs(sdf_points).mean()
            loss_inter = torch.exp(-100.0 * torch.abs(sdf_samples)).mean()
            loss_normal = torch.zeros(1, device=self.device)
            loss_eik = loss_utils.eikonal_loss(gradients_samples)
            loss_div = loss_utils.div_loss(samples, gradients_samples)

            w = self.loss_weights_sdf
            loss = (w[0] * loss_pull +
                    w[1] * loss_sdf +
                    w[2] * loss_grad_consis +
                    w[3] * loss_inter +
                    0.01 * loss_normal +
                    0.005 * loss_eik +
                    0.001 * loss_div)

            self.sdf_optimizer.zero_grad()
            loss.backward()
            self.sdf_optimizer.step()

            if (iter_i + 1) % 500 == 0:
                pbar.set_postfix(loss=f"{loss.item():.4f}")
            if (iter_i + 1) % self.save_freq == 0:
                self.save_sdf_checkpoint(iter_i + 1)

        print("SDF training complete.")
        self.save_sdf_checkpoint('final')

    def update_lr(self, optimizer, iter_step, max_iter, init_lr, warm_up_end=1000):
        if iter_step < warm_up_end:
            lr = (iter_step / warm_up_end) * init_lr
        else:
            lr = 0.5 * (math.cos((iter_step - warm_up_end) / (max_iter - warm_up_end) * math.pi) + 1) * init_lr
        for g in optimizer.param_groups:
            g['lr'] = lr

    def save_sdf_checkpoint(self, tag):
        ckpt = {
            'iter_step': self.iter_step,
            'sdf_network': self.sdf_net.state_dict(),
        }
        os.makedirs(os.path.join(self.out_dir, 'sdf_checkpoints'), exist_ok=True)
        torch.save(ckpt, os.path.join(self.out_dir, 'sdf_checkpoints', f'sdf_{tag}.pth'))

    def load_sdf_checkpoint(self, path):
        ckpt = torch.load(path, map_location=self.device)
        self.sdf_net.load_state_dict(ckpt['sdf_network'])
        self.iter_step = ckpt.get('iter_step', 0)

    # ------------------------------------------------------------------
    # VG stage helpers
    # ------------------------------------------------------------------
    @torch.no_grad()
    def get_surface_queries(self, noisy_pts=False):
        """Project point_gt onto the learned SDF surface."""
        sdf_level = self.project_sdf_level
        queries = self.point_gt.clone()
        if noisy_pts or sdf_level != 0.0:
            queries = self.project_queries(queries, sdf_level)

        n = len(queries)
        target = min(self.surface_queries_size, n + len(self.sample))
        if target > n:
            pad_size = target - n
            # Use FPS on projected samples
            pad_queries = self.sample.clone()
            pad_queries = self.project_queries(pad_queries, sdf_level)
            idx = fps_sample(pad_queries.cpu().numpy(), pad_size)
            pad_queries = pad_queries[idx]
            queries = torch.cat([queries, pad_queries], dim=0)
        return queries.detach()

    @torch.no_grad()
    def project_queries(self, queries, sdf_level):
        batch_size = 100_000
        out = []
        for i in range(0, len(queries), batch_size):
            batch = queries[i:i + batch_size]
            for _ in range(10):
                grad, sdf = self.sdf_net.gradient(batch)
                grad = F.normalize(grad, dim=-1)
                batch = batch - grad * (sdf - sdf_level)
            out.append(batch)
        return torch.cat(out, dim=0)

    # ------------------------------------------------------------------
    # VG stage
    # ------------------------------------------------------------------
    def train_vg(self, vertices_size=None):
        if vertices_size is None:
            vertices_size = self.vertices_size
        print(f"\n=== Stage 2: Vertex Generation ({vertices_size} vertices) ===")
        self.vg_net.train()
        self.sdf_net.eval()

        # Build target surface queries
        print("Projecting surface queries...")
        point_gt = self.get_surface_queries()
        print(f"  Surface queries: {len(point_gt)}")

        # Sample initial vertices via FPS
        sample_points = self.fps_select_vertices(point_gt, vertices_size)
        sample_normal, _ = self.sdf_net.gradient(sample_points)
        sample_normal = F.normalize(sample_normal.detach(), dim=-1)

        # Curvature on surface
        normal_gt, _ = self.sdf_net.gradient(point_gt)
        normal_gt = F.normalize(normal_gt.detach(), dim=-1)
        curvature_surface = loss_utils.cal_curvature_with_normal(
            point_gt, normal_gt, knn=min(16, len(point_gt) - 1)).detach()

        # Generate curriculum sizes
        batch_sizes = self.generate_list_with_ratio(vertices_size)
        print(f"  Curriculum sizes: {batch_sizes}")

        cur_size_idx = 0
        current_batch_size = batch_sizes[cur_size_idx]
        sample_points = self.fps_select_vertices(point_gt, current_batch_size)
        sample_normal, _ = self.sdf_net.gradient(sample_points)
        sample_normal = F.normalize(sample_normal.detach(), dim=-1)

        pbar = tqdm(range(self.vg_iters), desc="VG")
        size_update_freq = self.vg_iters // (self.update_size + 1)
        if size_update_freq == 0:
            size_update_freq = self.vg_iters

        nearest_clamp = self.cal_nearest_clamp(sample_points)

        for iter_i in pbar:
            generated = self.vg_net(sample_points, sample_normal)
            vertices_grad, _ = self.sdf_net.gradient(generated)

            loss = loss_utils.cal_vg_loss(
                point_gt, normal_gt, curvature_surface,
                generated, vertices_grad,
                self.loss_weights_vg, nearest_clamp)

            self.vg_optimizer.zero_grad()
            loss.backward(retain_graph=True)
            self.vg_optimizer.step()

            if (iter_i + 1) % 500 == 0:
                pbar.set_postfix(loss=f"{loss.item():.4f}")

            # Curriculum: increase vertex count
            if (iter_i + 1) % size_update_freq == 0:
                cur_size_idx += 1
                if cur_size_idx < len(batch_sizes):
                    current_batch_size = batch_sizes[cur_size_idx]
                    moved = self.move_to_surface(generated)
                    curv = loss_utils.cal_curvature_with_normal(
                        moved, F.normalize(vertices_grad.detach(), dim=-1),
                        knn=min(16, len(moved) - 1))
                    sample_points = self.upsample(curv, moved, point_gt, current_batch_size)
                    sample_points = sample_points.detach()
                    sn, _ = self.sdf_net.gradient(sample_points)
                    sample_normal = F.normalize(sn.detach(), dim=-1)
                    nearest_clamp = self.cal_nearest_clamp(sample_points)

        # Final projection to surface
        final_vertices = self.move_to_surface(generated).detach().cpu().numpy()
        print(f"  Generated {len(final_vertices)} vertices.")
        return final_vertices

    def generate_list_with_ratio(self, final_size):
        """Build curriculum vertex counts."""
        sizes = [int(final_size / (self.update_ratio ** (self.update_size - i)))
                 for i in range(self.update_size)]
        sizes.append(final_size)
        # Ensure monotonic
        for i in range(1, len(sizes)):
            sizes[i] = max(sizes[i], sizes[i - 1] + 1)
        return sizes

    def fps_select_vertices(self, point_gt, batch_size):
        idx = fps_sample(point_gt.cpu().numpy(), min(batch_size, len(point_gt)))
        return point_gt[idx].detach()

    def cal_nearest_clamp(self, sample_pts):
        pts_np = sample_pts.detach().cpu().numpy()
        from scipy.spatial import KDTree
        tree = KDTree(pts_np)
        _, idx = tree.query(pts_np, k=2)
        idx = torch.from_numpy(idx[:, 1]).long().to(sample_pts.device)
        neigh = sample_pts[idx]
        dist = torch.linalg.norm(neigh - sample_pts, ord=2, dim=-1) ** 2
        return dist.mean().item()

    def move_to_surface(self, generated, step=10):
        for _ in range(step):
            grad, sdf = self.sdf_net.gradient(generated)
            grad = F.normalize(grad.detach(), dim=-1)
            generated = generated - grad * (sdf.detach() - self.project_sdf_level)
        return generated.detach()

    def upsample(self, curvature, pts, point_gt, sample_size):
        """Upsample to target size by adding high-curvature neighbors."""
        if len(pts) >= sample_size:
            return pts[:sample_size]
        up = sample_size - len(pts)
        topk = min(up, len(pts))
        _, top_idx = torch.topk(curvature.view(-1), k=topk, largest=True)
        best = pts[top_idx]
        from scipy.spatial import KDTree
        tree = KDTree(point_gt.cpu().numpy())
        _, idx = tree.query(best.cpu().numpy(), k=1)
        idx = torch.from_numpy(idx).long().to(pts.device)
        added = point_gt[idx]
        return torch.cat([pts, added], dim=0)

    # ------------------------------------------------------------------
    # Meshing
    # ------------------------------------------------------------------
    def generate_mesh(self, vertices, save_path=None):
        print("\n=== Meshing ===")
        v, f = mesh_utils.delaunay_meshing(
            vertices, self.sdf_net,
            sdf_threshold=self.project_sdf_level,
            k_samples=self.k_samples,
            device=self.device)

        if len(f) > 0:
            v, f = mesh_utils.add_mid_vertices(v, f)

        # Denormalize
        v = denormalize_pointcloud(v, self.loc, self.scale)

        if save_path:
            if save_path.endswith('.obj'):
                save_mesh_obj(save_path, v, f)
            else:
                save_mesh_ply(save_path, v, f)
            print(f"Saved mesh to {save_path}")
        return v, f

    # ------------------------------------------------------------------
    # End-to-end
    # ------------------------------------------------------------------
    def run(self, mesh_path=None):
        self.train_sdf()
        vertices = self.train_vg()
        v, f = self.generate_mesh(vertices, save_path=mesh_path)
        return v, f