Upload external/MV-Adapter/mvadapter/utils/mesh_utils/warp.py with huggingface_hub

external/MV-Adapter/mvadapter/utils/mesh_utils/warp.py

@@ -0,0 +1,285 @@
+import os
+
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import torch.nn.functional as F
+import trimesh
+from einops import rearrange
+from PIL import Image, ImageDraw
+from tqdm import tqdm
+
+
+def draw(img, lines):
+    """
+    lines: n x 4, x1,y2,x2,y2
+    """
+    if isinstance(img, torch.Tensor):
+        img = img.detach().cpu().numpy()
+
+    if isinstance(img, np.ndarray):
+        if img.dtype == np.float32 or img.dtype == np.float64:
+            img = (img * 255.0).astype(np.uint8)
+        img = Image.fromarray(img)
+
+    img = img.resize((512, 512), Image.Resampling.BICUBIC)
+    w, h = img.size
+    assert w == h
+    if isinstance(lines, torch.Tensor):
+        lines = lines.detach().cpu().numpy()
+    lines = (lines + 1) / 2 * w
+
+    img1 = ImageDraw.Draw(img)
+    for line in lines:
+        img1.line(line.tolist(), fill="red", width=1)
+    return img
+
+
+def construct_grid_mesh(n_grid):
+    vertices = []
+    vertices_movable_ids = []
+    idx = 0
+    for j in range(n_grid + 1):
+        for i in range(n_grid + 1):
+            # 3D fit in trimesh format
+            if 0 < i < n_grid and 0 < j < n_grid:
+                vertices_movable_ids.append(idx)
+            vertices.append([i / n_grid, j / n_grid, 1.0 / 2])
+            idx += 1
+
+    vertices = np.array(vertices)
+    vertices = 2 * vertices - 1
+
+    faces = []
+    for j in range(n_grid):
+        for i in range(n_grid):
+            # clockwise
+            faces.append(
+                [
+                    i + j * (n_grid + 1),
+                    i + 1 + j * (n_grid + 1),
+                    i + (j + 1) * (n_grid + 1),
+                ]
+            )
+            faces.append(
+                [
+                    i + 1 + j * (n_grid + 1),
+                    i + 1 + (j + 1) * (n_grid + 1),
+                    i + (j + 1) * (n_grid + 1),
+                ]
+            )
+    faces = np.array(faces)
+    vertices_movable_ids = np.array(vertices_movable_ids)
+
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
+    return mesh, vertices_movable_ids, vertices
+
+
+def compute_warp_field(
+    ctx,
+    src_images_tensor,
+    tgt_images_tensor,
+    n_grid,
+    optim_res,
+    optim_step_per_res,
+    lambda_reg,
+    temp_dir,
+    verbose,
+    device,
+):
+    """
+    src_images_tensor: 4 H W 4
+    tgt_images_tensor: 4 H W 4
+    return: 4 H W 4
+    """
+    lam_reg = lambda_reg
+    lam_mask = 2
+    mesh, _, vertices_np = construct_grid_mesh(n_grid)
+
+    # vertices = mesh.vertices
+    faces = mesh.faces
+    edges = mesh.edges_unique
+
+    n_degree = torch.tensor(mesh.vertex_degree)
+    vertices = torch.tensor(vertices_np, device=device, dtype=torch.float32)
+    vertices_unopt = vertices.clone()
+    faces = torch.tensor(faces, device=device)
+    edges = torch.tensor(edges, device=device)
+
+    warped_images_tensor = []
+    bs = src_images_tensor.shape[0]
+
+    for img_idx in range(bs):
+        src_image_tensor = src_images_tensor[img_idx]
+        tgt_image_tensor = tgt_images_tensor[img_idx]
+
+        if verbose and temp_dir is not None:
+            vis_dir = os.path.join(temp_dir, f"{img_idx}")
+            os.makedirs(vis_dir, exist_ok=True)
+        # prepare for nvdiff rendering
+        v_origin_homo = torch.cat(
+            [vertices_unopt, torch.ones_like(vertices_unopt[..., :1])], dim=-1
+        )
+
+        # prepare parameters
+        v_move_indices = torch.where(n_degree == 6)[0]
+        vertices_movable = vertices[v_move_indices].detach().clone()
+        vertices_movable.requires_grad = True
+        opt = torch.optim.Adam([vertices_movable], lr=0.02)
+
+        for resl in optim_res:
+            rast, _ = dr.rasterize(
+                ctx,
+                v_origin_homo[None, ...].float(),
+                faces.int(),
+                (resl, resl),
+                grad_db=True,
+            )
+            face_ids = rast[..., 3].long() - 1
+            assert torch.all(face_ids >= 0)
+            u = rast[..., 0]
+            v = rast[..., 1]
+            pixel_vertice_ids = faces[face_ids]
+            pixel_bary_coords = torch.stack([u, v, 1 - u - v], dim=-1)
+
+            src_img = F.interpolate(
+                rearrange(src_image_tensor[..., :3], "H W C -> () C H W"),
+                size=(resl, resl),
+                mode="bilinear",
+                align_corners=False,
+                antialias=True,
+            )
+
+            tgt_img = F.interpolate(
+                rearrange(tgt_image_tensor[..., :3], "H W C -> () C H W"),
+                size=(resl, resl),
+                mode="bilinear",
+                align_corners=False,
+                antialias=True,
+            )
+
+            # src_img, tgt_img: 1 C H W
+            if verbose:
+                Image.fromarray(
+                    (tgt_img[0].permute(1, 2, 0).detach().cpu().numpy() * 255.0).astype(
+                        np.uint8
+                    )
+                ).save(os.path.join(vis_dir, f"target_{resl:04d}.png"))
+            with tqdm(range(optim_step_per_res), disable=not verbose) as pbar:
+                for i in pbar:
+                    opt.zero_grad()
+                    vertices_all = vertices_unopt.detach().clone()
+                    vertices_all[v_move_indices] = vertices_movable
+
+                    pixel_vertices = vertices_all[
+                        pixel_vertice_ids
+                    ]  # 1 x 512 x 512 x 3 x 3
+                    pixel_coords_inter = torch.sum(
+                        pixel_vertices * pixel_bary_coords[..., None], dim=-2
+                    )
+                    src_img_warped = F.grid_sample(
+                        src_img,
+                        pixel_coords_inter[..., :2],
+                        mode="bilinear",
+                        align_corners=False,
+                    )
+
+                    edge_vertices_all = vertices_all[edges]
+                    edge_vertices_unopt = vertices_unopt[edges]
+
+                    edge_len_all = torch.linalg.norm(
+                        edge_vertices_all[:, 0, :2] - edge_vertices_all[:, 1, :2],
+                        dim=-1,
+                    )
+                    edge_len_unopt = torch.linalg.norm(
+                        edge_vertices_unopt[:, 0, :2] - edge_vertices_all[:, 1, :2],
+                        dim=-1,
+                    )
+
+                    reg_loss = ((edge_len_all - edge_len_unopt) ** 2).mean()
+
+                    # mask_loss = ((src_img_warped[:, 3:4, :, :] - tgt_img[:, 3:4, :, :])**2).mean()
+                    # rgb_loss = (((src_img_warped[:, :3, :, :] - tgt_img[:, :3, :, :]) * tgt_img[:, 3:4, :, :])**2).mean()
+                    # loss = rgb_loss + lam_reg * reg_loss + lam_mask * mask_loss
+                    img_loss = ((src_img_warped - tgt_img) ** 2).mean()
+                    loss = img_loss + lam_reg * reg_loss
+
+                    loss.backward()
+                    opt.step()
+                    if verbose:
+                        print(img_loss.item(), reg_loss.item())
+
+                    if verbose:
+                        draw(
+                            src_img[0].permute(1, 2, 0),
+                            edge_vertices_all[:, :, :2].reshape(-1, 4),
+                        ).save(os.path.join(vis_dir, f"src_{resl:04d}_{i:03d}.png"))
+                        Image.fromarray(
+                            (
+                                torch.cat(
+                                    [
+                                        tgt_img,
+                                        src_img_warped,
+                                        (tgt_img - src_img_warped).abs(),
+                                    ],
+                                    dim=-1,
+                                )[0]
+                                .permute(1, 2, 0)
+                                .detach()
+                                .cpu()
+                                .numpy()
+                                * 255.0
+                            ).astype(np.uint8)
+                        ).save(os.path.join(vis_dir, f"opt_{resl:04d}_{i:03d}.png"))
+                        warped_this_resl = Image.fromarray(
+                            (
+                                src_img_warped[0]
+                                .permute(1, 2, 0)
+                                .detach()
+                                .cpu()
+                                .numpy()
+                                * 255.0
+                            ).astype(np.uint8)
+                        )
+                        warped_this_resl.save(
+                            os.path.join(vis_dir, f"warped_{resl:04d}_{i:03d}.png")
+                        )
+
+        # warp on 512 resolution
+        resl = src_image_tensor.shape[1]
+        with torch.no_grad():
+            rast, _ = dr.rasterize(
+                ctx,
+                v_origin_homo[None, ...].float(),
+                faces.int(),
+                (resl, resl),
+                grad_db=True,
+            )
+            face_ids = rast[..., 3].long() - 1
+            assert torch.all(face_ids >= 0)
+            u = rast[..., 0]
+            v = rast[..., 1]
+            pixel_vertice_ids = faces[face_ids]
+            pixel_bary_coords = torch.stack([u, v, 1 - u - v], dim=-1)
+
+            vertices_all = vertices_unopt.detach().clone()
+            vertices_all[v_move_indices] = vertices_movable
+
+            pixel_vertices = vertices_all[pixel_vertice_ids]  # 1 x 512 x 512 x 3 x 3
+            pixel_coords_inter = torch.sum(
+                pixel_vertices * pixel_bary_coords[..., None], dim=-2
+            )
+            src_img_warped = rearrange(
+                F.grid_sample(
+                    rearrange(src_image_tensor, "H W C -> () C H W"),
+                    pixel_coords_inter[..., :2],
+                    mode="bicubic",
+                    align_corners=False,
+                ),
+                "() C H W -> H W C",
+            ).clamp(0, 1)
+        warped_images_tensor.append(src_img_warped)
+
+    warped_images_tensor = torch.stack(warped_images_tensor, dim=0)  # 4 H W 3
+
+    return warped_images_tensor