| from dataclasses import dataclass, field |
| import time |
| from typing import Optional |
| import warnings |
| import torch |
| from tqdm import tqdm |
| from pathlib import Path |
| from core.opt import MeshOptimizer |
| from core.remesh import calc_edge_length, calc_edges, calc_vertex_normals |
| from util.func import ( |
| laplacian, |
| load_obj, |
| make_sphere, |
| make_star_cameras, |
| normalize_vertices, |
| save_images, |
| to_numpy, |
| ) |
| from util.render import NormalsRenderer |
| from util.snapshot import Snapshot, snapshot |
| import numpy as np |
|
|
| try: |
| from pyremesh import remesh_botsch |
| except: |
| remesh_botsch = None |
|
|
| |
| warnings.filterwarnings( |
| "ignore", |
| message="torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument.", |
| ) |
|
|
|
|
| @dataclass |
| class OptimizeSettings: |
| |
| target_fname: Path = None |
| target_vertices: torch.Tensor = None |
| target_faces: torch.Tensor = None |
|
|
| |
| steps: Optional[int] = 500 |
| timeout: Optional[float] = None |
|
|
| outdir: str = "out" |
| method: str = "ours" |
| image_size: int = 512 |
| sphere_size: float = 0.5 |
| sphere_level: int = 2 |
| sphere_shift: tuple[float, float, float] = None |
| cameras: tuple[int, int] = (4, 4) |
| device = "cuda" |
|
|
| |
| lr: float = 0.5 |
| laplacian_weight: float = 0.1 |
| ramp: float = 3.0 |
| betas: tuple[float, float, float] = (0.8, 0.8, 0) |
| remesh_interval: int = 1 |
| edge_len_lims: tuple[float, float] = (0.01, 0.15) |
|
|
| |
| gammas: tuple[float, float, float] = (0, 0, 0) |
| nu_ref: float = 0.3 |
| edge_len_tol: float = 0.5 |
| gain: float = 0.2 |
| local_edgelen: bool = True |
|
|
| |
| remesh_ratio: float = 0.5 |
|
|
| |
| result_interval: int = 5 |
| result_meshes: bool = False |
| result_snapshots: bool = False |
|
|
| save_images: bool = False |
|
|
|
|
| @dataclass |
| class OptimizeResult: |
| settings: OptimizeSettings |
| target_vertices: torch.Tensor = None |
| target_faces: torch.Tensor = None |
| snapshots: list[Snapshot] = field(default_factory=list) |
|
|
|
|
| def make_optimizer(settings, vertices, faces): |
| edges, _ = calc_edges(faces) |
| mean_edge_length = calc_edge_length(vertices, edges).mean().item() |
| lr = settings.lr * mean_edge_length |
| Laplacian = None |
|
|
| if settings.method == "adam": |
| vertices.requires_grad_() |
| opt = torch.optim.Adam([vertices], lr=lr, betas=settings.betas) |
| edges, _ = calc_edges(faces) |
| Laplacian = laplacian(vertices.shape[0], edges) |
| loss = (vertices * (Laplacian @ vertices)).mean() |
| elif settings.method == "ours": |
| opt = MeshOptimizer( |
| vertices, |
| faces, |
| lr=settings.lr, |
| betas=settings.betas, |
| gammas=settings.gammas, |
| nu_ref=settings.nu_ref, |
| edge_len_lims=settings.edge_len_lims, |
| edge_len_tol=settings.edge_len_tol, |
| gain=settings.gain, |
| laplacian_weight=settings.laplacian_weight, |
| ramp=settings.ramp, |
| remesh_interval=settings.remesh_interval, |
| local_edgelen=settings.local_edgelen, |
| ) |
|
|
| vertices = opt.vertices |
| else: |
| raise RuntimeError("unknown method") |
|
|
| return opt, lr, vertices, Laplacian |
|
|
|
|
| def load_target_mesh(fname, device="cuda"): |
| vertices, faces = load_obj(fname, device=device) |
| vertices = normalize_vertices(vertices) |
| return vertices, faces |
|
|
|
|
| def optimize(settings: OptimizeSettings): |
| result = OptimizeResult(settings=settings) |
| outdir = Path(settings.outdir) |
|
|
| vertices, faces = make_sphere( |
| level=settings.sphere_level, radius=settings.sphere_size, device=settings.device |
| ) |
| if settings.sphere_shift: |
| vertices += torch.tensor(settings.sphere_shift, device=settings.device) |
|
|
| mv, proj = make_star_cameras( |
| settings.cameras[0], |
| settings.cameras[1], |
| distance=10, |
| image_size=[settings.image_size, settings.image_size], |
| device=settings.device, |
| ) |
|
|
| renderer = NormalsRenderer( |
| mv, proj, image_size=[settings.image_size, settings.image_size] |
| ) |
|
|
| if settings.target_vertices is None: |
| target_vertices, target_faces = load_target_mesh(settings.target_fname) |
| else: |
| target_vertices, target_faces = settings.target_vertices, settings.target_faces |
|
|
| result.target_vertices, result.target_faces = target_vertices, target_faces |
|
|
| target_normals = calc_vertex_normals(target_vertices, target_faces) |
| target_images = renderer.render(target_vertices, target_normals, target_faces) |
|
|
| if settings.save_images: |
| save_images(target_images, outdir / "target_images") |
|
|
| opt, lr, vertices, Laplacian = make_optimizer(settings, vertices, faces) |
| start = time.time() |
| step = 1 |
| last_remesh_step = 0 |
| with tqdm( |
| desc="Optimize", |
| total=settings.steps if settings.timeout is None else settings.timeout, |
| leave=False, |
| ) as tqdm_: |
| is_last = False |
| while not is_last: |
| is_last = ( |
| step == settings.steps |
| if settings.steps |
| else time.time() - start > settings.timeout |
| ) |
|
|
| opt.zero_grad() |
|
|
| normals = calc_vertex_normals(vertices, faces) |
| images = renderer.render(vertices, normals, faces) |
| loss = (images - target_images).abs().mean() |
|
|
| if isinstance(opt, torch.optim.Adam): |
| |
| loss = ( |
| loss |
| + (vertices * (Laplacian @ vertices)).mean() |
| * settings.laplacian_weight |
| ) |
|
|
| loss.backward() |
|
|
| if isinstance(opt, torch.optim.Adam): |
| |
| ramped_lr = lr * min( |
| 1, |
| (step - last_remesh_step) * (1 - settings.betas[0]) / settings.ramp, |
| ) |
| opt.param_groups[0]["lr"] = ramped_lr |
|
|
| opt.step() |
|
|
| |
| with torch.no_grad(): |
| if ( |
| settings.result_interval and step % settings.result_interval == 1 |
| ) or is_last: |
| if settings.method == "ours": |
| s = snapshot(opt) |
| else: |
| s = Snapshot( |
| step=step, |
| time=time.time() - start, |
| vertices=vertices.clone().requires_grad_(False), |
| faces=faces.clone(), |
| ) |
| result.snapshots.append(s) |
|
|
| |
| if ( |
| settings.remesh_interval is not None |
| and (step % settings.remesh_interval) == settings.remesh_interval - 1 |
| and not is_last |
| ): |
|
|
| if isinstance(opt, MeshOptimizer): |
| vertices, faces = opt.remesh() |
| else: |
| with torch.no_grad(): |
| edges, _ = calc_edges(faces) |
| mean_edge_length = ( |
| calc_edge_length(vertices, edges).mean().item() |
| ) |
| target_edgelen = mean_edge_length * settings.remesh_ratio |
| target_edgelen = max(target_edgelen, settings.edge_len_lims[0]) |
| v = to_numpy(vertices).astype(np.double) |
| f = to_numpy(faces).astype(np.int32) |
| v, f = remesh_botsch(v, f, 5, target_edgelen, True) |
| vertices = torch.tensor( |
| v, dtype=torch.float, device=vertices.device |
| ).contiguous() |
| faces = torch.tensor( |
| f, dtype=torch.long, device=vertices.device |
| ).contiguous() |
| opt, lr, vertices, Laplacian = make_optimizer( |
| settings, vertices, faces |
| ) |
| last_remesh_step = step |
|
|
| if vertices.shape[0] == 0: |
| is_last = True |
|
|
| if settings.save_images: |
| save_images(images, outdir / "images") |
|
|
| step += 1 |
| if settings.steps is not None: |
| tqdm_.update(1) |
| else: |
| tqdm_.update( |
| min(settings.timeout, round(time.time() - start, 3)) - tqdm_.n |
| ) |
|
|
| return result |
|
|
