| from typing import Dict, List, Optional, Tuple, Union |
|
|
| import matplotlib.pyplot as plt |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from pytorch3d.loss import ( |
| chamfer_distance, |
| mesh_edge_loss, |
| mesh_laplacian_smoothing, |
| mesh_normal_consistency, |
| ) |
| from pytorch3d.structures import Meshes |
|
|
| from engine.MVSRender.camera_utils import MiniCam, OrbitCamera, orbit_camera |
|
|
|
|
| def compute_chamfer_loss( |
| query_verts: torch.Tensor, gt_pts: torch.Tensor |
| ) -> torch.Tensor: |
| """Compute chamfer distance between two point sets.""" |
| return chamfer_distance(query_verts[None], gt_pts[None])[0] |
|
|
|
|
| def update_mesh_shape_prior_losses( |
| mesh: Union[Meshes, "Mesh"], loss: Dict[str, torch.Tensor] |
| ) -> None: |
| """Update mesh regularization losses.""" |
| if isinstance(mesh, Meshes): |
| loss.update( |
| { |
| "edge": mesh_edge_loss(mesh), |
| "normal": mesh_normal_consistency(mesh), |
| "laplacian": mesh_laplacian_smoothing(mesh, method="uniform"), |
| } |
| ) |
| else: |
| loss.update( |
| {"normal": mesh.normal_consistency(), "laplacian": mesh.laplacian()} |
| ) |
|
|
|
|
| def plot_losses(losses: Dict[str, Dict[str, List[float]]]) -> None: |
| """Plot training losses over iterations.""" |
| plt.figure(figsize=(13, 5)) |
| ax = plt.gca() |
| for k, l in losses.items(): |
| ax.plot(l["values"], label=f"{k} loss") |
| ax.legend(fontsize=16) |
| ax.set_xlabel("Iteration", fontsize=16) |
| ax.set_ylabel("Loss", fontsize=16) |
| ax.set_title("Loss vs iterations", fontsize=16) |
|
|
|
|
| def inverse_sigmoid(x: Union[float, torch.Tensor]) -> torch.Tensor: |
| """Inverse sigmoid function.""" |
| x = torch.tensor(x).float() if isinstance(x, float) else x |
| return torch.log(x / (1 - x)) |
|
|
|
|
| def image_grid( |
| images: np.ndarray, |
| rows: Optional[int] = None, |
| cols: Optional[int] = None, |
| fill: bool = True, |
| show_axes: bool = False, |
| rgb: bool = True, |
| normal: bool = False, |
| ) -> None: |
| """Plot grid of images.""" |
| if (rows is None) != (cols is None): |
| raise ValueError("Specify either both rows and cols or neither") |
|
|
| rows = rows or len(images) |
| cols = cols or 1 |
|
|
| fig, axes = plt.subplots( |
| rows, |
| cols, |
| gridspec_kw={"wspace": 0.0, "hspace": 0.0} if fill else {}, |
| figsize=(15, 9), |
| ) |
| fig.subplots_adjust(left=0, bottom=0, right=1, top=1) |
|
|
| for ax, im in zip(np.array(axes).ravel(), images): |
| channel_slice = slice(3) if rgb else slice(3, 4) |
| channel_slice = slice(4, 7) if normal else channel_slice |
| ax.imshow(im[..., channel_slice]) |
| ax.set_axis_off() if not show_axes else None |
|
|
|
|
| def visualize_prediction( |
| predicted: torch.Tensor, |
| target: torch.Tensor, |
| vis_normal: bool = False, |
| gt_normal: Optional[torch.Tensor] = None, |
| pred_normal: Optional[torch.Tensor] = None, |
| title: str = "", |
| ) -> None: |
| """Visualize prediction vs target comparison.""" |
| figsize = (20, 10) if not vis_normal else (20, 20) |
| plt.figure(figsize=figsize) |
|
|
| if not vis_normal: |
| plt.subplot(1, 2, 1).imshow(predicted.detach().cpu().numpy()) |
| plt.subplot(1, 2, 2).imshow(target.detach().cpu().numpy()) |
| else: |
| plt.subplot(2, 2, 1).imshow(predicted.detach().cpu().numpy()) |
| plt.subplot(2, 2, 2).imshow(target.detach().cpu().numpy()) |
| plt.subplot(2, 2, 3).imshow(pred_normal.detach().cpu().numpy()) |
| plt.subplot(2, 2, 4).imshow(gt_normal.detach().cpu().numpy()) |
|
|
| plt.suptitle(title) |
| plt.tight_layout() |
|
|
|
|
| def camera_traj(cam, ref_size=1024, views=30, radius=2.0): |
|
|
| azimuth_bins = 360 // views |
|
|
| cameras = [] |
|
|
| for rotate_i in range(30): |
| pose = orbit_camera( |
| 0, |
| 0 + rotate_i * azimuth_bins, |
| radius, |
| ) |
|
|
| cur_cam = MiniCam( |
| pose, |
| ref_size, |
| ref_size, |
| cam.fovy, |
| cam.fovx, |
| cam.near, |
| cam.far, |
| ) |
|
|
| cameras.append(cur_cam) |
|
|
| return cameras |
|
|
