Initial commit

.gitignore

@@ -0,0 +1,4 @@
+*venv
+flagged
+*examples
+__pycache__

app.py

@@ -0,0 +1,192 @@
+import gradio as gr
+
+
+import numpy as np
+import cv2
+from tqdm import tqdm
+
+import torch
+from pytorch3d.io.obj_io import load_obj
+import tempfile
+import main_mcc
+import mcc_model
+import util.misc as misc
+from engine_mcc import prepare_data
+from plyfile import PlyData, PlyElement
+
+def run_inference(model, samples, device, temperature, args):
+    model.eval()
+
+    seen_xyz, valid_seen_xyz, unseen_xyz, unseen_rgb, labels, seen_images = prepare_data(
+        samples, device, is_train=False, args=args, is_viz=True
+    )
+    pred_occupy = []
+    pred_colors = []
+
+    max_n_unseen_fwd = 2000
+
+    model.cached_enc_feat = None
+    num_passes = int(np.ceil(unseen_xyz.shape[1] / max_n_unseen_fwd))
+    for p_idx in range(num_passes):
+        p_start = p_idx     * max_n_unseen_fwd
+        p_end = (p_idx + 1) * max_n_unseen_fwd
+        cur_unseen_xyz = unseen_xyz[:, p_start:p_end]
+        cur_unseen_rgb = unseen_rgb[:, p_start:p_end].zero_()
+        cur_labels = labels[:, p_start:p_end].zero_()
+
+        with torch.no_grad():
+            _, pred = model(
+                seen_images=seen_images,
+                seen_xyz=seen_xyz,
+                unseen_xyz=cur_unseen_xyz,
+                unseen_rgb=cur_unseen_rgb,
+                unseen_occupy=cur_labels,
+                cache_enc=True,
+                valid_seen_xyz=valid_seen_xyz,
+            )
+        if device == "cuda":
+            pred_occupy.append(pred[..., 0].cuda())
+        else:
+            pred_occupy.append(pred[..., 0].cpu())
+        if args.regress_color:
+            pred_colors.append(pred[..., 1:].reshape((-1, 3)))
+        else:
+            pred_colors.append(
+                (
+                    torch.nn.Softmax(dim=2)(
+                        pred[..., 1:].reshape((-1, 3, 256)) / temperature
+                    ) * torch.linspace(0, 1, 256, device=pred.device)
+                ).sum(axis=2)
+            )
+    
+    pred_occupy = torch.cat(pred_occupy, dim=1)
+    pred_occupy = torch.nn.Sigmoid()(pred_occupy)
+    return torch.cat(pred_colors, dim=0).cpu().numpy(), pred_occupy.cpu().numpy(), unseen_xyz.cpu().numpy()
+
+def pad_image(im, value):
+    if im.shape[0] > im.shape[1]:
+        diff = im.shape[0] - im.shape[1]
+        return torch.cat([im, (torch.zeros((im.shape[0], diff, im.shape[2])) + value)], dim=1)
+    else:
+        diff = im.shape[1] - im.shape[0]
+        return torch.cat([im, (torch.zeros((diff, im.shape[1], im.shape[2])) + value)], dim=0)
+
+
+def normalize(seen_xyz):
+    seen_xyz = seen_xyz / (seen_xyz[torch.isfinite(seen_xyz.sum(dim=-1))].var(dim=0) ** 0.5).mean()
+    seen_xyz = seen_xyz - seen_xyz[torch.isfinite(seen_xyz.sum(dim=-1))].mean(axis=0)
+    return seen_xyz
+
+def infer(
+          image,
+          point_cloud,
+          seg,
+          granularity,
+          temperature,
+          ):
+    
+    score_thresholds = [0.1, 0.2, 0.3, 0.4, 0.5]
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    parser = main_mcc.get_args_parser()
+    parser.set_defaults(eval=True)
+
+    args = parser.parse_args()
+    
+    model = mcc_model.get_mcc_model(
+        occupancy_weight=1.0,
+        rgb_weight=0.01,
+        args=args,
+    )
+    
+    if device == "cuda":
+        model = model.cuda()
+
+    misc.load_model(args=args, model_without_ddp=model, optimizer=None, loss_scaler=None)
+
+    rgb = image
+    obj = load_obj(point_cloud.name)
+
+    seen_rgb = (torch.tensor(rgb).float() / 255)[..., [2, 1, 0]]
+    H, W = seen_rgb.shape[:2]
+    seen_rgb = torch.nn.functional.interpolate(
+        seen_rgb.permute(2, 0, 1)[None],
+        size=[H, W],
+        mode="bilinear",
+        align_corners=False,
+    )[0].permute(1, 2, 0)
+
+    seen_xyz = obj[0].reshape(H, W, 3)
+    seg = cv2.imread(seg.name, cv2.IMREAD_UNCHANGED)
+    mask = torch.tensor(cv2.resize(seg, (W, H))).bool()
+    seen_xyz[~mask] = float('inf')
+
+    seen_xyz = normalize(seen_xyz)
+
+    bottom, right = mask.nonzero().max(dim=0)[0]
+    top, left = mask.nonzero().min(dim=0)[0]
+
+    bottom = bottom + 40
+    right = right + 40
+    top = max(top - 40, 0)
+    left = max(left - 40, 0)
+
+    seen_xyz = seen_xyz[top:bottom+1, left:right+1]
+    seen_rgb = seen_rgb[top:bottom+1, left:right+1]
+
+    seen_xyz = pad_image(seen_xyz, float('inf'))
+    seen_rgb = pad_image(seen_rgb, 0)
+
+    seen_rgb = torch.nn.functional.interpolate(
+        seen_rgb.permute(2, 0, 1)[None],
+        size=[800, 800],
+        mode="bilinear",
+        align_corners=False,
+    )
+
+    seen_xyz = torch.nn.functional.interpolate(
+        seen_xyz.permute(2, 0, 1)[None],
+        size=[112, 112],
+        mode="bilinear",
+        align_corners=False,
+    ).permute(0, 2, 3, 1)
+
+    samples = [
+        [seen_xyz, seen_rgb],
+        [torch.zeros((20000, 3)), torch.zeros((20000, 3))],
+    ]
+
+    pred_colors, pred_occupy, unseen_xyz = run_inference(model, samples, device, temperature, args)
+    _masks = pred_occupy > 0.1    
+    unseen_xyz = unseen_xyz[_masks]
+    pred_colors = pred_colors[None, ...][_masks] * 255
+    
+    # Prepare data for PlyElement
+    vertex = np.core.records.fromarrays(np.hstack((unseen_xyz, pred_colors)).transpose(), 
+                                               names='x, y, z, red, green, blue', 
+                                               formats='f8, f8, f8, u1, u1, u1')
+    
+
+    # Create PlyElement
+    element = PlyElement.describe(vertex, 'vertex')
+    
+    # Save point cloud data to a temporary file
+    with tempfile.NamedTemporaryFile(suffix=".ply", delete=False) as f:
+        PlyData([element], text=True).write(f)
+        temp_file_name = f.name
+
+    return temp_file_name
+
+
+demo = gr.Interface(fn=infer, 
+                    inputs=[gr.Image(label="Input Image"),
+                            gr.File(label="Pointcloud File"),
+                            gr.File(label="Segmentation File"),
+                            gr.Slider(minimum=0.05, maximum=0.5, step=0.05, value=0.2, label="Granularity"),
+                            gr.Slider(minimum=0, maximum=1.0, step=0.1, value=0.1, label="Temperature")
+                            ], 
+                    outputs=[gr.outputs.File(label="Point Cloud Json")],
+                    examples=[["demo/quest2.jpg", "demo/quest2.obj", "demo/quest2_seg.png", 0.2, 0.1]],
+                    cache_examples=True)
+demo.launch(server_name="0.0.0.0", server_port=7860)

engine_mcc.py

@@ -0,0 +1,587 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# MAE: https://github.com/facebookresearch/mae
+# --------------------------------------------------------
+import math
+from typing import Iterable
+import os
+import matplotlib.pyplot as plt
+import random
+import torch
+import numpy as np
+import time
+import base64
+from io import BytesIO
+
+import util.misc as misc
+import util.lr_sched as lr_sched
+
+from pytorch3d.structures import Pointclouds
+from pytorch3d.vis.plotly_vis import plot_scene
+from pytorch3d.transforms import RotateAxisAngle
+from pytorch3d.io import IO
+
+
+def evaluate_points(predicted_xyz, gt_xyz, dist_thres):
+    if predicted_xyz.shape[0] == 0:
+        return 0.0, 0.0, 0.0
+    slice_size = 1000
+    precision = 0.0
+    for i in range(int(np.ceil(predicted_xyz.shape[0] / slice_size))):
+        start = slice_size * i
+        end   = slice_size * (i + 1)
+        dist = ((predicted_xyz[start:end, None] - gt_xyz[None]) ** 2.0).sum(axis=-1) ** 0.5
+        precision += ((dist < dist_thres).sum(axis=1) > 0).sum()
+    precision /= predicted_xyz.shape[0]
+
+    recall = 0.0
+    for i in range(int(np.ceil(predicted_xyz.shape[0] / slice_size))):
+        start = slice_size * i
+        end   = slice_size * (i + 1)
+        dist = ((predicted_xyz[:, None] - gt_xyz[None, start:end]) ** 2.0).sum(axis=-1) ** 0.5
+        recall += ((dist < dist_thres).sum(axis=0) > 0).sum()
+    recall /= gt_xyz.shape[0]
+    return precision, recall, get_f1(precision, recall)
+
+def aug_xyz(seen_xyz, unseen_xyz, args, is_train):
+    degree_x = 0
+    degree_y = 0
+    degree_z = 0
+    if is_train:
+        r_delta = args.random_scale_delta
+        scale = torch.tensor([
+            random.uniform(1.0 - r_delta, 1.0 + r_delta),
+            random.uniform(1.0 - r_delta, 1.0 + r_delta),
+            random.uniform(1.0 - r_delta, 1.0 + r_delta),
+        ], device=seen_xyz.device)
+
+        if args.use_hypersim:
+            shift = 0
+        else:
+            degree_x = random.randrange(-args.random_rotate_degree, args.random_rotate_degree + 1)
+            degree_y = random.randrange(-args.random_rotate_degree, args.random_rotate_degree + 1)
+            degree_z = random.randrange(-args.random_rotate_degree, args.random_rotate_degree + 1)
+
+            r_shift = args.random_shift
+            shift = torch.tensor([[[
+                random.uniform(-r_shift, r_shift),
+                random.uniform(-r_shift, r_shift),
+                random.uniform(-r_shift, r_shift),
+            ]]], device=seen_xyz.device)
+        seen_xyz = seen_xyz * scale + shift
+        unseen_xyz = unseen_xyz * scale + shift
+
+    B, H, W, _ = seen_xyz.shape
+    return [
+        rotate(seen_xyz.reshape((B, -1, 3)), degree_x, degree_y, degree_z).reshape((B, H, W, 3)),
+        rotate(unseen_xyz, degree_x, degree_y, degree_z),
+    ]
+
+
+def rotate(sample, degree_x, degree_y, degree_z):
+    for degree, axis in [(degree_x, "X"), (degree_y, "Y"), (degree_z, "Z")]:
+        if degree != 0:
+            sample = RotateAxisAngle(degree, axis=axis).to(sample.device).transform_points(sample)
+    return sample
+
+
+def get_grid(B, device, co3d_world_size, granularity):
+    N = int(np.ceil(2 * co3d_world_size / granularity))
+    grid_unseen_xyz = torch.zeros((N, N, N, 3), device=device)
+    for i in range(N):
+        grid_unseen_xyz[i, :, :, 0] = i
+    for j in range(N):
+        grid_unseen_xyz[:, j, :, 1] = j
+    for k in range(N):
+        grid_unseen_xyz[:, :, k, 2] = k
+    grid_unseen_xyz -= (N / 2.0)
+    grid_unseen_xyz /= (N / 2.0) / co3d_world_size
+    grid_unseen_xyz = grid_unseen_xyz.reshape((1, -1, 3)).repeat(B, 1, 1)
+    return grid_unseen_xyz
+
+
+def run_viz(model, data_loader, device, args, epoch):
+    epoch_start_time = time.time()
+    model.eval()
+    os.system(f'mkdir {args.job_dir}/viz')
+
+    print('Visualization data_loader length:', len(data_loader))
+    dataset = data_loader.dataset
+    for sample_idx, samples in enumerate(data_loader):
+        if sample_idx >= args.max_n_viz_obj:
+            break
+        seen_xyz, valid_seen_xyz, unseen_xyz, unseen_rgb, labels, seen_images = prepare_data(samples, device, is_train=False, args=args, is_viz=True)
+
+        pred_occupy = []
+        pred_colors = []
+        (model.module if hasattr(model, "module") else model).clear_cache()
+
+        # don't forward all at once to avoid oom
+        max_n_queries_fwd = 2000
+
+        total_n_passes = int(np.ceil(unseen_xyz.shape[1] / max_n_queries_fwd))
+        for p_idx in range(total_n_passes):
+            p_start = p_idx     * max_n_queries_fwd
+            p_end = (p_idx + 1) * max_n_queries_fwd
+            cur_unseen_xyz = unseen_xyz[:, p_start:p_end]
+            cur_unseen_rgb = unseen_rgb[:, p_start:p_end].zero_()
+            cur_labels = labels[:, p_start:p_end].zero_()
+
+            with torch.no_grad():
+                _, pred, = model(
+                    seen_images=seen_images,
+                    seen_xyz=seen_xyz,
+                    unseen_xyz=cur_unseen_xyz,
+                    unseen_rgb=cur_unseen_rgb,
+                    unseen_occupy=cur_labels,
+                    cache_enc=args.run_viz,
+                    valid_seen_xyz=valid_seen_xyz,
+                )
+
+            cur_occupy_out = pred[..., 0]
+
+            if args.regress_color:
+                cur_color_out = pred[..., 1:].reshape((-1, 3))
+            else:
+                cur_color_out = pred[..., 1:].reshape((-1, 3, 256)).max(dim=2)[1] / 255.0
+            pred_occupy.append(cur_occupy_out)
+            pred_colors.append(cur_color_out)
+
+        rank = misc.get_rank()
+        prefix = f'{args.job_dir}/viz/' + dataset.dataset_split + f'_ep{epoch}_rank{rank}_i{sample_idx}'
+
+        img = (seen_images[0].permute(1, 2, 0) * 255).cpu().numpy().copy().astype(np.uint8)
+
+        gt_xyz = samples[1][0].to(device).reshape(-1, 3)
+        gt_rgb = samples[1][1].to(device).reshape(-1, 3)
+        mesh_xyz = samples[2].to(device).reshape(-1, 3) if args.use_hypersim else None
+
+        with open(prefix + '.html', 'a') as f:
+            generate_html(
+                img,
+                seen_xyz, seen_images,
+                torch.cat(pred_occupy, dim=1),
+                torch.cat(pred_colors, dim=0),
+                unseen_xyz,
+                f,
+                gt_xyz=gt_xyz,
+                gt_rgb=gt_rgb,
+                mesh_xyz=mesh_xyz,
+            )
+    print("Visualization epoch time:", time.time() - epoch_start_time)
+
+
+def get_f1(precision, recall):
+    if (precision + recall) == 0:
+        return 0.0
+    return 2.0 * precision * recall / (precision + recall)
+
+
+def generate_plot(img, seen_xyz, seen_rgb, pred_occ, pred_rgb, unseen_xyz,
+        gt_xyz=None, gt_rgb=None, mesh_xyz=None, score_thresholds=[0.1, 0.3, 0.5, 0.7, 0.9],
+        pointcloud_marker_size=2,
+    ):
+    # if img is not None:
+    #     fig = plt.figure()
+    #     plt.imshow(img)
+    #     tmpfile = BytesIO()
+    #     fig.savefig(tmpfile, format='jpg')
+    #     encoded = base64.b64encode(tmpfile.getvalue()).decode('utf-8')
+
+    #     html = '<img src=\'data:image/png;base64,{}\'>'.format(encoded)
+    #     f.write(html)
+    #     plt.close()
+
+    clouds = {"MCC Output": {}}
+    # Seen
+    if seen_xyz is not None:
+        seen_xyz = seen_xyz.reshape((-1, 3)).cpu()
+        seen_rgb = torch.nn.functional.interpolate(seen_rgb, (112, 112)).permute(0, 2, 3, 1).reshape((-1, 3)).cpu()
+        good_seen = seen_xyz[:, 0] != -100
+
+        seen_pc = Pointclouds(
+            points=seen_xyz[good_seen][None],
+            features=seen_rgb[good_seen][None],
+        )
+        clouds["MCC Output"]["seen"] = seen_pc
+
+    # GT points
+    if gt_xyz is not None:
+        subset_gt = random.sample(range(gt_xyz.shape[0]), 10000)
+        gt_pc = Pointclouds(
+            points=gt_xyz[subset_gt][None],
+            features=gt_rgb[subset_gt][None],
+        )
+        clouds["MCC Output"]["GT points"] = gt_pc
+
+    # GT meshes
+    if mesh_xyz is not None:
+        subset_mesh = random.sample(range(mesh_xyz.shape[0]), 10000)
+        mesh_pc = Pointclouds(
+            points=mesh_xyz[subset_mesh][None],
+        )
+        clouds["MCC Output"]["GT mesh"] = mesh_pc
+
+    pred_occ = torch.nn.Sigmoid()(pred_occ).cpu()
+    for t in score_thresholds:
+        pos = pred_occ > t
+
+        points = unseen_xyz[pos].reshape((-1, 3))
+        features = pred_rgb[None][pos].reshape((-1, 3))
+        good_points = points[:, 0] != -100
+
+        if good_points.sum() == 0:
+            continue
+
+        pc = Pointclouds(
+            points=points[good_points][None].cpu(),
+            features=features[good_points][None].cpu(),
+        )
+
+        clouds["MCC Output"][f"pred_{t}"] = pc
+        IO().save_pointcloud(pc, "output_pointcloud.ply")
+
+    plt.figure()
+    try:
+        fig = plot_scene(clouds, pointcloud_marker_size=pointcloud_marker_size, pointcloud_max_points=20000 * 2)
+        fig.update_layout(height=1000, width=1000)
+        return fig
+    except Exception as e:
+        print('writing failed', e)
+    try:
+        plt.close()
+    except:
+        pass
+
+
+def generate_html(img, seen_xyz, seen_rgb, pred_occ, pred_rgb, unseen_xyz, f,
+        gt_xyz=None, gt_rgb=None, mesh_xyz=None, score_thresholds=[0.1, 0.3, 0.5, 0.7, 0.9],
+        pointcloud_marker_size=2,
+    ):
+    if img is not None:
+        fig = plt.figure()
+        plt.imshow(img)
+        tmpfile = BytesIO()
+        fig.savefig(tmpfile, format='jpg')
+        encoded = base64.b64encode(tmpfile.getvalue()).decode('utf-8')
+
+        html = '<img src=\'data:image/png;base64,{}\'>'.format(encoded)
+        f.write(html)
+        plt.close()
+
+    clouds = {"MCC Output": {}}
+    # Seen
+    if seen_xyz is not None:
+        seen_xyz = seen_xyz.reshape((-1, 3)).cpu()
+        seen_rgb = torch.nn.functional.interpolate(seen_rgb, (112, 112)).permute(0, 2, 3, 1).reshape((-1, 3)).cpu()
+        good_seen = seen_xyz[:, 0] != -100
+
+        seen_pc = Pointclouds(
+            points=seen_xyz[good_seen][None],
+            features=seen_rgb[good_seen][None],
+        )
+        clouds["MCC Output"]["seen"] = seen_pc
+
+    # GT points
+    if gt_xyz is not None:
+        subset_gt = random.sample(range(gt_xyz.shape[0]), 10000)
+        gt_pc = Pointclouds(
+            points=gt_xyz[subset_gt][None],
+            features=gt_rgb[subset_gt][None],
+        )
+        clouds["MCC Output"]["GT points"] = gt_pc
+
+    # GT meshes
+    if mesh_xyz is not None:
+        subset_mesh = random.sample(range(mesh_xyz.shape[0]), 10000)
+        mesh_pc = Pointclouds(
+            points=mesh_xyz[subset_mesh][None],
+        )
+        clouds["MCC Output"]["GT mesh"] = mesh_pc
+
+    pred_occ = torch.nn.Sigmoid()(pred_occ).cpu()
+    for t in score_thresholds:
+        pos = pred_occ > t
+
+        points = unseen_xyz[pos].reshape((-1, 3))
+        features = pred_rgb[None][pos].reshape((-1, 3))
+        good_points = points[:, 0] != -100
+
+        if good_points.sum() == 0:
+            continue
+
+        pc = Pointclouds(
+            points=points[good_points][None].cpu(),
+            features=features[good_points][None].cpu(),
+        )
+
+        clouds["MCC Output"][f"pred_{t}"] = pc
+
+    plt.figure()
+    try:
+        fig = plot_scene(clouds, pointcloud_marker_size=pointcloud_marker_size, pointcloud_max_points=20000 * 2)
+        fig.update_layout(height=1000, width=1000)
+        html_string = fig.to_html(full_html=False, include_plotlyjs="cnd")
+        f.write(html_string)
+        return fig, plt
+    except Exception as e:
+        print('writing failed', e)
+    try:
+        plt.close()
+    except:
+        pass
+
+
+def train_one_epoch(model: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler,
+                    args=None):
+    epoch_start_time = time.time()
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+
+    accum_iter = args.accum_iter
+
+    optimizer.zero_grad()
+
+    print('Training data_loader length:', len(data_loader))
+    for data_iter_step, samples in enumerate(data_loader):
+        # we use a per iteration (instead of per epoch) lr scheduler
+        if data_iter_step % accum_iter == 0:
+            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+        seen_xyz, valid_seen_xyz, unseen_xyz, unseen_rgb, labels, seen_images = prepare_data(samples, device, is_train=True, args=args)
+
+        with torch.cuda.amp.autocast():
+            loss, _ = model(
+                seen_images=seen_images,
+                seen_xyz=seen_xyz,
+                unseen_xyz=unseen_xyz,
+                unseen_rgb=unseen_rgb,
+                unseen_occupy=labels,
+                valid_seen_xyz=valid_seen_xyz,
+            )
+
+        loss_value = loss.item()
+        if not math.isfinite(loss_value):
+            print("Warning: Loss is {}".format(loss_value))
+            loss *= 0.0
+            loss_value = 100.0
+
+        loss /= accum_iter
+        loss_scaler(loss, optimizer, parameters=model.parameters(),
+                    clip_grad=args.clip_grad,
+                    update_grad=(data_iter_step + 1) % accum_iter == 0,
+                    verbose=(data_iter_step % 100) == 0)
+
+        if (data_iter_step + 1) % accum_iter == 0:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+        metric_logger.update(loss=loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        if data_iter_step == 30:
+            os.system('nvidia-smi')
+            os.system('free -g')
+        if args.debug and data_iter_step == 5:
+            break
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    print("Training epoch time:", time.time() - epoch_start_time)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+def eval_one_epoch(
+        model: torch.nn.Module,
+        data_loader: Iterable,
+        device: torch.device,
+        args=None
+    ):
+    epoch_start_time = time.time()
+    model.train(False)
+
+    metric_logger = misc.MetricLogger(delimiter="  ")
+
+    print('Eval len(data_loader):', len(data_loader))
+
+    for data_iter_step, samples in enumerate(data_loader):
+        seen_xyz, valid_seen_xyz, unseen_xyz, unseen_rgb, labels, seen_images = prepare_data(samples, device, is_train=False, args=args)
+
+        # don't forward all at once to avoid oom
+        max_n_queries_fwd = 5000
+        all_loss, all_preds = [], []
+        for p_idx in range(int(np.ceil(unseen_xyz.shape[1] / max_n_queries_fwd))):
+            p_start = p_idx     * max_n_queries_fwd
+            p_end = (p_idx + 1) * max_n_queries_fwd
+            cur_unseen_xyz = unseen_xyz[:, p_start:p_end]
+            cur_unseen_rgb = unseen_rgb[:, p_start:p_end]
+            cur_labels = labels[:, p_start:p_end]
+
+            with torch.no_grad():
+                loss, pred = model(
+                    seen_images=seen_images,
+                    seen_xyz=seen_xyz,
+                    unseen_xyz=cur_unseen_xyz,
+                    unseen_rgb=cur_unseen_rgb,
+                    unseen_occupy=cur_labels,
+                    valid_seen_xyz=valid_seen_xyz,
+                )
+            all_loss.append(loss)
+            all_preds.append(pred)
+
+        loss = sum(all_loss) / len(all_loss)
+        pred = torch.cat(all_preds, dim=1)
+
+        B = pred.shape[0]
+
+        gt_xyz = samples[1][0].to(device).reshape((B, -1, 3))
+        if args.use_hypersim:
+            mesh_xyz = samples[2].to(device).reshape((B, -1, 3))
+
+        s_thres = args.eval_score_threshold
+        d_thres = args.eval_dist_threshold
+
+        for b_idx in range(B):
+            geometry_metrics = {}
+            predicted_idx = torch.nn.Sigmoid()(pred[b_idx, :, 0]) > s_thres
+            predicted_xyz = unseen_xyz[b_idx, predicted_idx]
+
+            precision, recall, f1 = evaluate_points(predicted_xyz, gt_xyz[b_idx], d_thres)
+            geometry_metrics[f'd{d_thres}_s{s_thres}_point_pr'] = precision
+            geometry_metrics[f'd{d_thres}_s{s_thres}_point_rc'] = recall
+            geometry_metrics[f'd{d_thres}_s{s_thres}_point_f1'] = f1
+
+            if args.use_hypersim:
+                precision, recall, f1 = evaluate_points(predicted_xyz, mesh_xyz[b_idx], d_thres)
+                geometry_metrics[f'd{d_thres}_s{s_thres}_mesh_pr'] = precision
+                geometry_metrics[f'd{d_thres}_s{s_thres}_mesh_rc'] = recall
+                geometry_metrics[f'd{d_thres}_s{s_thres}_mesh_f1'] = f1
+
+            metric_logger.update(**geometry_metrics)
+
+        loss_value = loss.item()
+
+        torch.cuda.synchronize()
+        metric_logger.update(loss=loss_value)
+
+        if args.debug and data_iter_step == 5:
+            break
+
+    metric_logger.synchronize_between_processes()
+    print("Validation averaged stats:", metric_logger)
+    print("Val epoch time:", time.time() - epoch_start_time)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+def sample_uniform_semisphere(B, N, semisphere_size, device):
+    for _ in range(100):
+        points = torch.empty(B * N * 3, 3, device=device).uniform_(-semisphere_size, semisphere_size)
+        points[..., 2] = points[..., 2].abs()
+        dist = (points ** 2.0).sum(axis=-1) ** 0.5
+        if (dist < semisphere_size).sum() >= B * N:
+            return points[dist < semisphere_size][:B * N].reshape((B, N, 3))
+        else:
+            print('resampling sphere')
+
+
+def get_grid_semisphere(B, granularity, semisphere_size, device):
+    n_grid_pts = int(semisphere_size / granularity) * 2 + 1
+    grid_unseen_xyz = torch.zeros((n_grid_pts, n_grid_pts, n_grid_pts // 2 + 1, 3), device=device)
+    for i in range(n_grid_pts):
+        grid_unseen_xyz[i, :, :, 0] = i
+        grid_unseen_xyz[:, i, :, 1] = i
+    for i in range(n_grid_pts // 2 + 1):
+        grid_unseen_xyz[:, :, i, 2] = i
+    grid_unseen_xyz[..., :2] -= (n_grid_pts // 2.0)
+    grid_unseen_xyz *= granularity
+    dist = (grid_unseen_xyz ** 2.0).sum(axis=-1) ** 0.5
+    grid_unseen_xyz = grid_unseen_xyz[dist <= semisphere_size]
+    return grid_unseen_xyz[None].repeat(B, 1, 1)
+
+
+def get_min_dist(a, b, slice_size=1000):
+    all_min, all_idx = [], []
+    for i in range(int(np.ceil(a.shape[1] / slice_size))):
+        start = slice_size * i
+        end   = slice_size * (i + 1)
+        # B, n_queries, n_gt
+        dist = ((a[:, start:end] - b) ** 2.0).sum(axis=-1) ** 0.5
+        # B, n_queries
+        cur_min, cur_idx = dist.min(axis=2)
+        all_min.append(cur_min)
+        all_idx.append(cur_idx)
+    return torch.cat(all_min, dim=1), torch.cat(all_idx, dim=1)
+
+
+def construct_uniform_semisphere(gt_xyz, gt_rgb, semisphere_size, n_queries, dist_threshold, is_train, granularity):
+    B = gt_xyz.shape[0]
+    device = gt_xyz.device
+    if is_train:
+        unseen_xyz = sample_uniform_semisphere(B, n_queries, semisphere_size, device)
+    else:
+        unseen_xyz = get_grid_semisphere(B, granularity, semisphere_size, device)
+    dist, idx_to_gt = get_min_dist(unseen_xyz[:, :, None], gt_xyz[:, None])
+    labels = dist < dist_threshold
+    unseen_rgb = torch.zeros_like(unseen_xyz)
+    unseen_rgb[labels] = torch.gather(gt_rgb, 1, idx_to_gt.unsqueeze(-1).repeat(1, 1, 3))[labels]
+    return unseen_xyz, unseen_rgb, labels.float()
+
+
+def construct_uniform_grid(gt_xyz, gt_rgb, co3d_world_size, n_queries, dist_threshold, is_train, granularity):
+    B = gt_xyz.shape[0]
+    device = gt_xyz.device
+    if is_train:
+        unseen_xyz = torch.empty((B, n_queries, 3), device=device).uniform_(-co3d_world_size, co3d_world_size)
+    else:
+        unseen_xyz = get_grid(B, device, co3d_world_size, granularity)
+    dist, idx_to_gt = get_min_dist(unseen_xyz[:, :, None], gt_xyz[:, None])
+    labels = dist < dist_threshold
+    unseen_rgb = torch.zeros_like(unseen_xyz)
+    unseen_rgb[labels] = torch.gather(gt_rgb, 1, idx_to_gt.unsqueeze(-1).repeat(1, 1, 3))[labels]
+    return unseen_xyz, unseen_rgb, labels.float()
+
+
+def prepare_data(samples, device, is_train, args, is_viz=False):
+    # Seen
+    seen_xyz, seen_rgb = samples[0][0].to(device), samples[0][1].to(device)
+    valid_seen_xyz = torch.isfinite(seen_xyz.sum(axis=-1))
+    seen_xyz[~valid_seen_xyz] = -100
+    B = seen_xyz.shape[0]
+    # Gt
+    gt_xyz, gt_rgb = samples[1][0].to(device).reshape(B, -1, 3), samples[1][1].to(device).reshape(B, -1, 3)
+
+    sampling_func = construct_uniform_semisphere if args.use_hypersim else construct_uniform_grid
+    unseen_xyz, unseen_rgb, labels = sampling_func(
+        gt_xyz, gt_rgb,
+        args.semisphere_size if args.use_hypersim else args.co3d_world_size,
+        args.n_queries,
+        args.train_dist_threshold,
+        is_train,
+        args.viz_granularity if is_viz else args.eval_granularity,
+    )
+
+    if is_train:
+        seen_xyz, unseen_xyz = aug_xyz(seen_xyz, unseen_xyz, args, is_train=is_train)
+
+        # Random Flip
+        if random.random() < 0.5:
+            seen_xyz[..., 0] *= -1
+            unseen_xyz[..., 0] *= -1
+            seen_xyz = torch.flip(seen_xyz, [2])
+            valid_seen_xyz = torch.flip(valid_seen_xyz, [2])
+            seen_rgb = torch.flip(seen_rgb, [3])
+
+    return seen_xyz, valid_seen_xyz, unseen_xyz, unseen_rgb, labels, seen_rgb

main_mcc.py

@@ -0,0 +1,322 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# MAE: https://github.com/facebookresearch/mae
+# --------------------------------------------------------
+import argparse
+import datetime
+import json
+import numpy as np
+import os
+import time
+from pathlib import Path
+
+import torch
+import torch.backends.cudnn as cudnn
+import timm.optim.optim_factory as optim_factory
+
+import util.misc as misc
+import mcc_model
+from util.misc import NativeScalerWithGradNormCount as NativeScaler
+from util.hypersim_dataset import HyperSimDataset, hypersim_collate_fn
+from util.co3d_dataset import CO3DV2Dataset, co3dv2_collate_fn
+from engine_mcc import train_one_epoch, run_viz, eval_one_epoch
+from util.co3d_utils import get_all_dataset_maps
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('MCC', add_help=False)
+
+    # Model
+    parser.add_argument('--input_size', default=224, type=int,
+                        help='Images input size')
+    parser.add_argument('--occupancy_weight', default=1.0, type=float,
+                        help='A constant to weight the occupancy loss')
+    parser.add_argument('--rgb_weight', default=0.01, type=float,
+                        help='A constant to weight the color prediction loss')
+    parser.add_argument('--n_queries', default=550, type=int,
+                        help='Number of queries used in decoder.')
+    parser.add_argument('--drop_path', default=0.1, type=float,
+                        help='drop_path probability')
+    parser.add_argument('--regress_color', action='store_true',
+                        help='If true, regress color with MSE. Otherwise, 256-way classification for each channel.')
+
+    # Training
+    parser.add_argument('--batch_size', default=16, type=int,
+                        help='Batch size per GPU for training (effective batch size is batch_size * accum_iter * # gpus')
+    parser.add_argument('--eval_batch_size', default=2, type=int,
+                        help='Batch size per GPU for evaluation (effective batch size is batch_size * accum_iter * # gpus')
+    parser.add_argument('--epochs', default=100, type=int)
+    parser.add_argument('--accum_iter', default=1, type=int,
+                        help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
+    parser.add_argument('--weight_decay', type=float, default=0.05,
+                        help='Weight decay (default: 0.05)')
+    parser.add_argument('--lr', type=float, default=None, metavar='LR',
+                        help='Learning rate (absolute lr)')
+    parser.add_argument('--blr', type=float, default=1e-4, metavar='LR',
+                        help='Base learning rate: absolute_lr = base_lr * total_batch_size / 512')
+    parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
+                        help='Lower lr bound for cyclic schedulers that hit 0')
+    parser.add_argument('--warmup_epochs', type=int, default=5, metavar='N',
+                        help='Epochs to warmup LR')
+    parser.add_argument('--clip_grad', type=float, default=1.0,
+                        help='Clip gradient at the specified norm')
+
+    # Job
+    parser.add_argument('--job_dir', default='',
+                        help='Path to where to save, empty for no saving')
+    parser.add_argument('--output_dir', default='./output_dir',
+                        help='Path to where to save, empty for no saving')
+    parser.add_argument('--device', default='cuda',
+                        help='Device to use for training / testing')
+    parser.add_argument('--seed', default=0, type=int,
+                        help='Random seed.')
+    parser.add_argument('--resume', default='weights/co3dv2_all_categories.pth',
+                        help='Resume from checkpoint')
+
+    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
+                        help='Start epoch')
+    parser.add_argument('--num_workers', default=4, type=int,
+                        help='Number of workers for training data loader')
+    parser.add_argument('--num_eval_workers', default=4, type=int,
+                        help='Number of workers for evaluation data loader')
+    parser.add_argument('--pin_mem', action='store_true',
+                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
+    parser.set_defaults(pin_mem=True)
+
+    # Distributed training
+    parser.add_argument('--world_size', default=1, type=int,
+                        help='Number of distributed processes')
+    parser.add_argument('--local_rank', default=-1, type=int)
+    parser.add_argument('--dist_on_itp', action='store_true')
+    parser.add_argument('--dist_url', default='env://',
+                        help='Url used to set up distributed training')
+
+    # Experiments
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--run_viz', action='store_true',
+                        help='Specify to run only the visualization/inference given a trained model.')
+    parser.add_argument('--max_n_viz_obj', default=64, type=int,
+                        help='Max number of objects to visualize during training.')
+
+    # Data
+    parser.add_argument('--train_epoch_len_multiplier', default=32, type=int,
+                        help='# examples per training epoch is # objects * train_epoch_len_multiplier')
+    parser.add_argument('--eval_epoch_len_multiplier', default=1, type=int,
+                        help='# examples per eval epoch is # objects * eval_epoch_len_multiplier')
+
+    # CO3D
+    parser.add_argument('--co3d_path', type=str, default='co3d_data',
+                        help='Path to CO3D v2 data.')
+    parser.add_argument('--holdout_categories', action='store_true',
+                        help='If true, hold out 10 categories and train on only the remaining 41 categories.')
+    parser.add_argument('--co3d_world_size', default=3.0, type=float,
+                        help='The world space we consider is \in [-co3d_world_size, co3d_world_size] in each dimension.')
+
+    # Hypersim
+    parser.add_argument('--use_hypersim', action='store_true',
+                        help='If true, use hypersim, else, co3d.')
+    parser.add_argument('--hypersim_path', default="hypersim_data", type=str,
+                        help="Path to Hypersim data.")
+
+    # Data aug
+    parser.add_argument('--random_scale_delta', default=0.2, type=float,
+                        help='Random scaling each example by a scaler \in [1 - random_scale_delta, 1 + random_scale_delta].')
+    parser.add_argument('--random_shift', default=1.0, type=float,
+                        help='Random shifting an example in each axis by an amount \in [-random_shift, random_shift]')
+    parser.add_argument('--random_rotate_degree', default=180, type=int,
+                        help='Random rotation degrees.')
+
+    # Smapling, evaluation, and coordinate system
+    parser.add_argument('--shrink_threshold', default=10.0, type=float,
+                        help='Any points with distance beyond this value will be shrunk.')
+    parser.add_argument('--semisphere_size', default=6.0, type=float,
+                        help='The Hypersim task predicts points in a semisphere in front of the camera.'
+                             'This value specifies the size of the semisphere.')
+    parser.add_argument('--eval_granularity', default=0.1, type=float,
+                        help='Granularity of the evaluation points.')
+    parser.add_argument('--viz_granularity', default=0.1, type=float,
+                        help='Granularity of points in visaulizatoin.')
+
+    parser.add_argument('--eval_score_threshold', default=0.1, type=float,
+                        help='Score threshold for evaluation.')
+    parser.add_argument('--eval_dist_threshold', default=0.1, type=float,
+                        help='Points closer than this amount to a groud-truth is considered correct.')
+    parser.add_argument('--train_dist_threshold', default=0.1, type=float,
+                        help='Points closer than this amount is considered positive in training.')
+    return parser
+
+
+def build_loader(args, num_tasks, global_rank, is_train, dataset_type, collate_fn, dataset_maps):
+    '''Build data loader'''
+    dataset = dataset_type(args, is_train=is_train, dataset_maps=dataset_maps)
+
+    sampler_train = torch.utils.data.DistributedSampler(
+        dataset, num_replicas=num_tasks, rank=global_rank, shuffle=is_train
+    )
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=args.batch_size if is_train else args.eval_batch_size,
+        sampler=sampler_train,
+        num_workers=args.num_workers if is_train else args.num_eval_workers,
+        pin_memory=args.pin_mem,
+        collate_fn=collate_fn,
+    )
+    return data_loader
+
+
+def main(args):
+    misc.init_distributed_mode(args)
+
+    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
+    print("{}".format(args).replace(', ', ',\n'))
+
+    device = torch.device(args.device)
+
+    # fix the seed for reproducibility
+    seed = args.seed + misc.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    cudnn.benchmark = True
+    num_tasks = misc.get_world_size()
+    global_rank = misc.get_rank()
+
+    # define the model
+    model = mcc_model.get_mcc_model(
+        rgb_weight=args.rgb_weight,
+        occupancy_weight=args.occupancy_weight,
+        args=args,
+    )
+
+    model.to(device)
+
+    model_without_ddp = model
+    print("Model = %s" % str(model_without_ddp))
+
+    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
+    if args.lr is None:  # only base_lr is specified
+        args.lr = args.blr * eff_batch_size / 512
+
+    print("base lr: %.2e" % (args.blr))
+    print("actual lr: %.2e" % args.lr)
+
+    print("accumulate grad iterations: %d" % args.accum_iter)
+    print("effective batch size: %d" % eff_batch_size)
+
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
+        model_without_ddp = model.module
+
+    # following timm: set wd as 0 for bias and norm layers
+    param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
+    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    print(optimizer)
+    loss_scaler = NativeScaler()
+
+    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)
+
+    if args.use_hypersim:
+        dataset_type = HyperSimDataset
+        collate_fn = hypersim_collate_fn
+        dataset_maps = None
+    else:
+        dataset_type = CO3DV2Dataset
+        collate_fn = co3dv2_collate_fn
+        dataset_maps = get_all_dataset_maps(
+            args.co3d_path, args.holdout_categories,
+        )
+
+    dataset_viz = dataset_type(args, is_train=False, is_viz=True, dataset_maps=dataset_maps)
+    sampler_viz = torch.utils.data.DistributedSampler(
+        dataset_viz, num_replicas=num_tasks, rank=global_rank, shuffle=False
+    )
+
+    data_loader_viz = torch.utils.data.DataLoader(
+        dataset_viz, batch_size=1,
+        sampler=sampler_viz,
+        num_workers=args.num_eval_workers,
+        pin_memory=args.pin_mem,
+        collate_fn=collate_fn,
+    )
+
+    if args.run_viz:
+        run_viz(
+            model, data_loader_viz,
+            device, args=args, epoch=0,
+        )
+        exit()
+
+    data_loader_train, data_loader_val = [
+        build_loader(
+            args, num_tasks, global_rank,
+            is_train=is_train,
+            dataset_type=dataset_type, collate_fn=collate_fn, dataset_maps=dataset_maps
+        ) for is_train in [True, False]
+    ]
+
+    print(f"Start training for {args.epochs} epochs")
+    start_time = time.time()
+    for epoch in range(args.start_epoch, args.epochs):
+        print(f'Epoch {epoch}:')
+        if args.distributed:
+            data_loader_train.sampler.set_epoch(epoch)
+        train_stats = train_one_epoch(
+            model, data_loader_train,
+            optimizer, device, epoch, loss_scaler,
+            args=args,
+        )
+
+        val_stats = {}
+        if (epoch % 5 == 4 or epoch + 1 == args.epochs) or args.debug:
+            val_stats = eval_one_epoch(
+                model, data_loader_val,
+                device, args=args,
+            )
+
+        if ((epoch % 10 == 9 or epoch + 1 == args.epochs) or args.debug):
+            run_viz(
+                model, data_loader_viz,
+                device, args=args, epoch=epoch,
+            )
+
+        if args.output_dir and (epoch % 10 == 9 or epoch + 1 == args.epochs):
+            misc.save_model(
+                args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
+                loss_scaler=loss_scaler, epoch=epoch)
+
+        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+                     **{f'val_{k}': v for k, v in val_stats.items()},
+                     'epoch': epoch,}
+
+        if args.output_dir and misc.is_main_process():
+            with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
+                f.write(json.dumps(log_stats) + "\n")
+
+    run_viz(
+        model, data_loader_viz,
+        device, args=args, epoch=-1,
+    )
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Training time {}'.format(total_time_str))
+
+
+if __name__ == '__main__':
+
+    args = get_args_parser()
+    args = args.parse_args()
+
+    if args.output_dir:
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+
+    main(args)
+

mcc_model.py

@@ -0,0 +1,386 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# DeiT: https://github.com/facebookresearch/deit
+# MAE: https://github.com/facebookresearch/mae
+# --------------------------------------------------------
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from timm.models.vision_transformer import PatchEmbed, Block, Mlp, DropPath
+
+from util.pos_embed import get_2d_sincos_pos_embed
+
+class MCCDecoderAttention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., args=None):
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.args = args
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, unseen_size):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+
+        mask = torch.zeros((1, 1, N, N), device=attn.device)
+        mask[:, :, :, -unseen_size:] = float('-inf')
+        for i in range(unseen_size):
+            mask[:, :, -(i + 1), -(i + 1)] = 0
+        attn = attn + mask
+        attn = attn.softmax(dim=-1)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+class MCCDecoderBlock(nn.Module):
+
+    def __init__(
+            self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., init_values=None,
+            drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, args=None):
+        super().__init__()
+        self.args = args
+        self.norm1 = norm_layer(dim)
+        self.attn = MCCDecoderAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, args=args)
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x, unseen_size):
+        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x), unseen_size)))
+        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+        return x
+
+
+class XYZPosEmbed(nn.Module):
+    """ Masked Autoencoder with VisionTransformer backbone
+    """
+    def __init__(self, embed_dim):
+        super().__init__()
+        self.embed_dim = embed_dim
+
+        self.two_d_pos_embed = nn.Parameter(
+            torch.zeros(1, 64 + 1, embed_dim), requires_grad=False)  # fixed sin-cos embedding
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.win_size = 8
+
+        self.pos_embed = nn.Linear(3, embed_dim)
+
+        self.blocks = nn.ModuleList([
+            Block(embed_dim, num_heads=12, mlp_ratio=2.0, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
+            for _ in range(1)
+        ])
+
+        self.invalid_xyz_token = nn.Parameter(torch.zeros(embed_dim,))
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        torch.nn.init.normal_(self.cls_token, std=.02)
+
+        two_d_pos_embed = get_2d_sincos_pos_embed(self.two_d_pos_embed.shape[-1], 8, cls_token=True)
+        self.two_d_pos_embed.data.copy_(torch.from_numpy(two_d_pos_embed).float().unsqueeze(0))
+
+        torch.nn.init.normal_(self.invalid_xyz_token, std=.02)
+
+    def forward(self, seen_xyz, valid_seen_xyz):
+        emb = self.pos_embed(seen_xyz)
+
+        emb[~valid_seen_xyz] = 0.0
+        emb[~valid_seen_xyz] += self.invalid_xyz_token
+
+        B, H, W, C = emb.shape
+        emb = emb.view(B, H // self.win_size, self.win_size, W // self.win_size, self.win_size, C)
+        emb = emb.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.win_size * self.win_size, C)
+
+        emb = emb + self.two_d_pos_embed[:, 1:, :]
+        cls_token = self.cls_token + self.two_d_pos_embed[:, :1, :]
+
+        cls_tokens = cls_token.expand(emb.shape[0], -1, -1)
+        emb = torch.cat((cls_tokens, emb), dim=1)
+        for _, blk in enumerate(self.blocks):
+            emb = blk(emb)
+        return emb[:, 0].view(B, (H // self.win_size) * (W // self.win_size), -1)
+
+
+class DecodeXYZPosEmbed(nn.Module):
+    """ Masked Autoencoder with VisionTransformer backbone
+    """
+    def __init__(self, embed_dim):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.pos_embed = nn.Linear(3, embed_dim)
+
+    def forward(self, unseen_xyz):
+        return self.pos_embed(unseen_xyz)
+
+
+class MCC(nn.Module):
+    """ Masked Autoencoder with VisionTransformer backbone
+    """
+    def __init__(self,
+                 img_size=224, patch_size=16, in_chans=3,
+                 embed_dim=1024, depth=24, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4., norm_layer=nn.LayerNorm,
+                 rgb_weight=1.0, occupancy_weight=1.0, args=None):
+        super().__init__()
+
+        self.rgb_weight = rgb_weight
+        self.occupancy_weight = occupancy_weight
+        self.args = args
+
+        # --------------------------------------------------------------------------
+        # encoder specifics
+        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.cls_token_xyz = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)  # fixed sin-cos embedding
+
+        self.xyz_pos_embed = XYZPosEmbed(embed_dim)
+
+        self.blocks = nn.ModuleList([
+            Block(
+                embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer,
+                drop_path=args.drop_path
+            ) for i in range(depth)])
+
+        self.blocks_xyz = nn.ModuleList([
+            Block(
+                embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer,
+                drop_path=args.drop_path
+            ) for i in range(depth)])
+
+        self.norm = norm_layer(embed_dim)
+        self.norm_xyz = norm_layer(embed_dim)
+        self.cached_enc_feat = None
+
+        # --------------------------------------------------------------------------
+        # decoder specifics
+        self.decoder_embed = nn.Linear(
+            embed_dim * 2,
+            decoder_embed_dim,
+            bias=True
+        )
+
+        self.decoder_xyz_pos_embed = DecodeXYZPosEmbed(decoder_embed_dim)
+        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False)  # fixed sin-cos embedding
+
+        self.decoder_blocks = nn.ModuleList([
+            MCCDecoderBlock(
+                decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer,
+                drop_path=args.drop_path,
+                args=args,
+            ) for i in range(decoder_depth)])
+
+        self.decoder_norm = norm_layer(decoder_embed_dim)
+        if self.args.regress_color:
+            self.decoder_pred = nn.Linear(decoder_embed_dim, 3 + 1, bias=True) # decoder to patch
+        else:
+            self.decoder_pred = nn.Linear(decoder_embed_dim, 256 * 3 + 1, bias=True) # decoder to patch
+
+        self.loss_occupy = nn.BCEWithLogitsLoss()
+        if self.args.regress_color:
+            self.loss_rgb = nn.MSELoss()
+        else:
+            self.loss_rgb = nn.CrossEntropyLoss()
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
+        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))
+
+        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.patch_embed.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
+        torch.nn.init.normal_(self.cls_token, std=.02)
+        torch.nn.init.normal_(self.cls_token_xyz, std=.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+
+    def forward_encoder(self, x, seen_xyz, valid_seen_xyz):
+
+        # get tokens
+        x = self.patch_embed(x)
+        x = x + self.pos_embed[:, 1:, :]
+        y = self.xyz_pos_embed(seen_xyz, valid_seen_xyz)
+
+        ##### forward E_XYZ #####
+        # append cls token
+        cls_token_xyz = self.cls_token_xyz
+        cls_tokens_xyz = cls_token_xyz.expand(y.shape[0], -1, -1)
+
+        y = torch.cat((cls_tokens_xyz, y), dim=1)
+        # apply Transformer blocks
+        for blk in self.blocks_xyz:
+            y = blk(y)
+        y = self.norm_xyz(y)
+
+        ##### forward E_RGB #####
+        # append cls token
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+
+        x = torch.cat((cls_tokens, x), dim=1)
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        # combine encodings
+        x = torch.cat([x, y], dim=2)
+        return x
+
+    def forward_decoder(self, x, unseen_xyz):
+        # embed tokens
+        x = self.decoder_embed(x)
+        x = x + self.decoder_pos_embed
+
+        # 3D pos embed
+        unseen_xyz = self.decoder_xyz_pos_embed(unseen_xyz)
+        x = torch.cat([x, unseen_xyz], dim=1)
+
+        # apply Transformer blocks
+        for blk in self.decoder_blocks:
+            x = blk(x, unseen_xyz.shape[1])
+
+        x = self.decoder_norm(x)
+
+        # predictor projection
+        pred = self.decoder_pred(x)
+        # remove cls & seen token
+        pred = pred[:, -unseen_xyz.shape[1]:, :]
+
+        return pred
+
+    def forward_loss(self, pred, unseen_occupy, unseen_rgb):
+        loss = self.loss_occupy(
+            pred[:, :, :1].reshape((-1, 1)),
+            unseen_occupy.reshape((-1, 1)).float()
+        ) * self.occupancy_weight
+
+        if unseen_occupy.sum() > 0:
+            if self.args.regress_color:
+                pred_rgb = pred[:, :, 1:][unseen_occupy.bool()]
+                gt_rgb = unseen_rgb[unseen_occupy.bool()]
+            else:
+                pred_rgb = pred[:, :, 1:][unseen_occupy.bool()].reshape((-1, 256))
+                gt_rgb = torch.round(unseen_rgb[unseen_occupy.bool()] * 255).long().reshape((-1,))
+
+            rgb_loss = self.loss_rgb(pred_rgb, gt_rgb) * self.rgb_weight
+            loss = loss + rgb_loss
+        return loss
+
+
+    def clear_cache(self):
+        self.cached_enc_feat = None
+
+    def forward(self, seen_images, seen_xyz, unseen_xyz, unseen_rgb, unseen_occupy, valid_seen_xyz,
+                cache_enc=False):
+
+        unseen_xyz = shrink_points_beyond_threshold(unseen_xyz, self.args.shrink_threshold)
+
+        if self.cached_enc_feat is None:
+            seen_images = preprocess_img(seen_images)
+            seen_xyz = shrink_points_beyond_threshold(seen_xyz, self.args.shrink_threshold)
+            latent = self.forward_encoder(seen_images, seen_xyz, valid_seen_xyz)
+
+        if cache_enc:
+            if self.cached_enc_feat is None:
+                self.cached_enc_feat = latent
+            else:
+                latent = self.cached_enc_feat
+
+        pred = self.forward_decoder(latent, unseen_xyz)
+        loss = self.forward_loss(pred, unseen_occupy, unseen_rgb)
+        return loss, pred
+
+
+def get_mcc_model(**kwargs):
+    return MCC(
+        embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs
+    )
+
+
+def shrink_points_beyond_threshold(xyz, threshold):
+    xyz = xyz.clone().detach()
+    dist = (xyz ** 2.0).sum(axis=-1) ** 0.5
+    affected = (dist > threshold) * torch.isfinite(dist)
+    xyz[affected] = xyz[affected] * (
+        threshold * (2.0 - threshold / dist[affected]) / dist[affected]
+    )[..., None]
+    return xyz
+
+
+def preprocess_img(x):
+    if x.shape[2] != 224:
+        assert x.shape[2] == 800
+        x = F.interpolate(
+            x,
+            scale_factor=224./800.,
+            mode="bilinear",
+        )
+    resnet_mean = torch.tensor([0.485, 0.456, 0.406], device=x.device).reshape((1, 3, 1, 1))
+    resnet_std = torch.tensor([0.229, 0.224, 0.225], device=x.device).reshape((1, 3, 1, 1))
+    imgs_normed = (x - resnet_mean) / resnet_std
+    return imgs_normed
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_values=1e-5, inplace=False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x):
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma

packages.txt

@@ -0,0 +1,3 @@
+ffmpeg
+libsm6
+libxext6

pre-requirements.txt

@@ -0,0 +1,2 @@
+torch==1.13.0
+torchvision==0.14.0

requirements.txt

@@ -0,0 +1,11 @@
+h5py
+omegaconf
+submitit
+timm==0.4.5
+opencv-python
+matplotlib
+plotly
+gradio
+gradio_client==0.2.7
+plyfile
+git+https://github.com/facebookresearch/pytorch3d.git

util/co3d_dataset.py

@@ -0,0 +1,122 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import random
+from typing import cast
+
+import torch
+from pytorch3d.implicitron.dataset.dataset_base import FrameData
+
+import util.co3d_utils as co3d_utils
+
+
+def co3dv2_collate_fn(batch):
+    assert len(batch[0]) == 4
+    return (
+        FrameData.collate([x[0] for x in batch]),
+        FrameData.collate([x[1] for x in batch]),
+        [x[2] for x in batch],
+        [x[3] for x in batch],
+    )
+
+
+def pad_point_cloud(pc, N):
+    cur_N = pc._points_list[0].shape[0]
+    if cur_N == N:
+        return pc
+
+    assert cur_N > 0
+
+    n_pad = N - cur_N
+    indices = random.choices(list(range(cur_N)), k=n_pad)
+    pc._features_list[0] = torch.cat([pc._features_list[0], pc._features_list[0][indices]], dim=0)
+    pc._points_list[0] = torch.cat([pc._points_list[0], pc._points_list[0][indices]], dim=0)
+    return pc
+
+
+class CO3DV2Dataset(torch.utils.data.Dataset):
+    def __init__(self, args, is_train, is_viz=False, dataset_maps=None):
+
+        self.args = args
+        self.is_train = is_train
+        self.is_viz = is_viz
+
+        self.dataset_split = 'train' if is_train else 'val'
+        self.all_datasets = dataset_maps[0 if is_train else 1]
+        print(len(self.all_datasets), 'categories loaded')
+
+        self.all_example_names = self.get_all_example_names()
+        print('containing', len(self.all_example_names), 'examples')
+
+    def get_all_example_names(self):
+        all_example_names = []
+        for category in self.all_datasets.keys():
+            for sequence_name in self.all_datasets[category].seq_name2idx.keys():
+                all_example_names.append((category, sequence_name))
+        return all_example_names
+
+    def __getitem__(self, index):
+        for retry in range(1000):
+            try:
+                if retry > 9:
+                    index = random.choice(range(len(self)))
+                    print('retry', retry, 'new index:', index)
+                gap = 1 if self.is_train else len(self.all_example_names) // len(self)
+                assert gap >= 1
+                category, sequence_name = self.all_example_names[(index * gap) % len(self.all_example_names)]
+
+                cat_dataset = self.all_datasets[category]
+
+                frame_data = cat_dataset.__getitem__(
+                    random.choice(cat_dataset.seq_name2idx[sequence_name])
+                    if self.is_train
+                    else cat_dataset.seq_name2idx[sequence_name][
+                        hash(sequence_name) % len(cat_dataset.seq_name2idx[sequence_name])
+                    ]
+                )
+                test_frame = None
+                seen_idx = None
+
+                frame_data = cat_dataset.frame_data_type.collate([frame_data])
+                mask = (
+                    (cast(torch.Tensor, frame_data.fg_probability) > 0.5).float()
+                    if frame_data.fg_probability is not None
+                    else None
+                )
+                seen_rgb = frame_data.image_rgb.clone().detach()
+
+                # 112, 112, 3
+                seen_xyz = co3d_utils.get_rgbd_points(
+                    112, 112,
+                    frame_data.camera,
+                    frame_data.depth_map,
+                    mask,
+                )
+
+                full_point_cloud = co3d_utils._load_pointcloud(f'{self.args.co3d_path}/{category}/{sequence_name}/pointcloud.ply', max_points=20000)
+                full_point_cloud = pad_point_cloud(full_point_cloud, 20000)
+                break
+            except Exception as e:
+                print(category, sequence_name, 'sampling failed', retry, e)
+
+        seen_rgb = seen_rgb.squeeze(0)
+        full_rgb = full_point_cloud._features_list[0]
+
+        return (
+            (seen_xyz, seen_rgb),
+            (full_point_cloud._points_list[0], full_rgb),
+            test_frame,
+            (category, sequence_name, seen_idx),
+        )
+
+    def __len__(self) -> int:
+        n_objs = sum([len(cat_dataset.seq_name2idx.keys()) for cat_dataset in self.all_datasets.values()])
+        if self.is_train:
+            return int(n_objs * self.args.train_epoch_len_multiplier)
+        elif self.is_viz:
+            return n_objs
+        else:
+            return int(n_objs * self.args.eval_epoch_len_multiplier)

util/co3d_utils.py

@@ -0,0 +1,133 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import glob
+from omegaconf import DictConfig
+from typing import Optional
+
+import torch
+
+from pytorch3d.implicitron.dataset.dataset_map_provider import DatasetMap
+from pytorch3d.implicitron.dataset.json_index_dataset_map_provider_v2 import (
+    JsonIndexDatasetMapProviderV2
+)
+from pytorch3d.implicitron.tools.config import expand_args_fields
+from pytorch3d.io import IO
+from pytorch3d.renderer import (
+    NDCMultinomialRaysampler,
+    ray_bundle_to_ray_points,
+)
+from pytorch3d.renderer.cameras import CamerasBase
+from pytorch3d.structures import Pointclouds
+
+
+HOLDOUT_CATEGORIES = set([
+    'apple',
+    'baseballglove',
+    'cup',
+    'ball',
+    'toyplane',
+    'handbag',
+    'book',
+    'carrot',
+    'suitcase',
+    'bowl',
+])
+
+def get_dataset_map(
+    dataset_root: str,
+    category: str,
+    subset_name: str,
+) -> DatasetMap:
+    """
+    Obtain the dataset map that contains the train/val/test dataset objects.
+    """
+    expand_args_fields(JsonIndexDatasetMapProviderV2)
+    dataset_map_provider = JsonIndexDatasetMapProviderV2(
+        category=category,
+        subset_name=subset_name,
+        dataset_root=dataset_root,
+        test_on_train=False,
+        only_test_set=False,
+        load_eval_batches=True,
+        dataset_JsonIndexDataset_args=DictConfig({"remove_empty_masks": False, "load_point_clouds": False}),
+    )
+    return dataset_map_provider.get_dataset_map()
+
+
+def _load_pointcloud(pcl_path, max_points):
+    pcl = IO().load_pointcloud(pcl_path)
+    if max_points > 0:
+        pcl = pcl.subsample(max_points)
+
+    return pcl
+
+
+def get_all_dataset_maps(co3d_path, holdout_categories):
+    all_categories = [c.split('/')[-1] for c in list(glob.glob(co3d_path + '/*')) if not c.endswith('.json')]
+    all_categories = sorted(all_categories, key=lambda x: hash(x))
+
+    # Obtain the CO3Dv2 dataset map
+    train_dataset_maps = {}
+    val_dataset_maps = {}
+    for category in all_categories:
+
+        print(f'Loading dataset map ({category})')
+        dataset_map = {
+            'train': torch.load(f'dataset_cache/{category}_train.pt'),
+            'val': torch.load(f'dataset_cache/{category}_val.pt')
+        }
+        if not holdout_categories or category not in HOLDOUT_CATEGORIES:
+            train_dataset_maps[category] = dataset_map['train']
+        if not holdout_categories or category in HOLDOUT_CATEGORIES:
+            val_dataset_maps[category] = dataset_map['val']
+
+    print('Loaded', len(train_dataset_maps), 'categores for train')
+    print('Loaded', len(val_dataset_maps), 'categores for val')
+    return train_dataset_maps, val_dataset_maps
+
+
+def get_rgbd_points(
+    imh, imw,
+    camera: CamerasBase,
+    depth_map: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    mask_thr: float = 0.5,
+) -> Pointclouds:
+    """
+    Given a batch of images, depths, masks and cameras, generate a colored
+    point cloud by unprojecting depth maps to the  and coloring with the source
+    pixel colors.
+    """
+    depth_map = torch.nn.functional.interpolate(
+        depth_map,
+        size=[imh, imw],
+        mode="bilinear",
+        align_corners=False,
+    )
+    # convert the depth maps to point clouds using the grid ray sampler
+    pts_3d = ray_bundle_to_ray_points(
+        NDCMultinomialRaysampler(
+            image_width=imw,
+            image_height=imh,
+            n_pts_per_ray=1,
+            min_depth=1.0,
+            max_depth=1.0,
+        )(camera)._replace(lengths=depth_map[:, 0, ..., None])
+    ).squeeze(3)[None]
+
+    pts_mask = depth_map > 0.0
+    if mask is not None:
+        mask = torch.nn.functional.interpolate(
+            mask,
+            size=[imh, imw],
+            mode="bilinear",
+            align_corners=False,
+        )
+        pts_mask *= mask > mask_thr
+    pts_3d[~pts_mask] = float('inf')
+    return pts_3d.squeeze(0).squeeze(0)
+

util/crop.py

@@ -0,0 +1,42 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+import torch
+
+from torchvision import transforms
+from torchvision.transforms import functional as F
+
+
+class RandomResizedCrop(transforms.RandomResizedCrop):
+    """
+    RandomResizedCrop for matching TF/TPU implementation: no for-loop is used.
+    This may lead to results different with torchvision's version.
+    Following BYOL's TF code:
+    https://github.com/deepmind/deepmind-research/blob/master/byol/utils/dataset.py#L206
+    """
+    @staticmethod
+    def get_params(img, scale, ratio):
+        width, height = F._get_image_size(img)
+        area = height * width
+
+        target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
+        log_ratio = torch.log(torch.tensor(ratio))
+        aspect_ratio = torch.exp(
+            torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
+        ).item()
+
+        w = int(round(math.sqrt(target_area * aspect_ratio)))
+        h = int(round(math.sqrt(target_area / aspect_ratio)))
+
+        w = min(w, width)
+        h = min(h, height)
+
+        i = torch.randint(0, height - h + 1, size=(1,)).item()
+        j = torch.randint(0, width - w + 1, size=(1,)).item()
+
+        return i, j, h, w

util/hypersim_dataset.py

@@ -0,0 +1,271 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import random
+import glob
+
+import torch
+from pytorch3d.implicitron.dataset.dataset_base import FrameData
+from pytorch3d.ops import sample_points_from_meshes
+
+from util.hypersim_utils import read_h5py, read_img
+
+
+def hypersim_collate_fn(batch):
+    assert len(batch[0]) == 4
+    return (
+        FrameData.collate([x[0] for x in batch]),
+        FrameData.collate([x[1] for x in batch]),
+        FrameData.collate([x[2] for x in batch]),
+        [x[2] for x in batch]
+    )
+
+
+def is_good_xyz(xyz):
+    assert len(xyz.shape) == 3
+    return (torch.isfinite(xyz.sum(axis=2))).sum() > 2000
+
+
+def get_camera_pos_file_name_from_frame_name(frame_name):
+    tmp = frame_name.split('/')
+    tmp[-3] = '_detail'
+    tmp[-2] = 'cam_' + tmp[-2].split('_')[2]
+    tmp[-1] = 'camera_keyframe_positions.hdf5'
+    return '/'.join(tmp)
+
+
+def get_camera_look_at_file_name_from_frame_name(frame_name):
+    tmp = frame_name.split('/')
+    tmp[-3] = '_detail'
+    tmp[-2] = 'cam_' + tmp[-2].split('_')[2]
+    tmp[-1] = 'camera_keyframe_look_at_positions.hdf5'
+    return '/'.join(tmp)
+
+
+def get_camera_orientation_file_name_from_frame_name(frame_name):
+    tmp = frame_name.split('/')
+    tmp[-3] = '_detail'
+    tmp[-2] = 'cam_' + tmp[-2].split('_')[2]
+    tmp[-1] = 'camera_keyframe_orientations.hdf5'
+    return '/'.join(tmp)
+
+
+def read_scale_from_frame_name(frame_name):
+    tmp = frame_name.split('/')
+    with open('/'.join(tmp[:-3] + ['_detail', 'metadata_scene.csv'])) as f:
+        for line in f:
+            items = line.split(',')
+    return float(items[1])
+
+
+def random_crop(xyz, img, is_train=True):
+    assert xyz.shape[0] == img.shape[0]
+    assert xyz.shape[1] == img.shape[1]
+
+    width, height = img.shape[0], img.shape[1]
+    w = h = min(width, height)
+    if is_train:
+        i = torch.randint(0, width - w + 1, size=(1,)).item()
+        j = torch.randint(0, height - h + 1, size=(1,)).item()
+    else:
+        i = (width - w) // 2
+        j = (height - h) // 2
+    xyz = xyz[i:i+w, j:j+h]
+    img = img[i:i+w, j:j+h]
+    xyz = torch.nn.functional.interpolate(
+        xyz[None].permute(0, 3, 1, 2), (112, 112),
+        mode='bilinear',
+    ).permute(0, 2, 3, 1)[0]
+    img = torch.nn.functional.interpolate(
+        img[None].permute(0, 3, 1, 2), (224, 224),
+        mode='bilinear',
+    ).permute(0, 2, 3, 1)[0]
+    return xyz, img
+
+
+class HyperSimDataset(torch.utils.data.Dataset):
+    def __init__(self, args, is_train, is_viz=False, **kwargs):
+
+        self.args = args
+        self.is_train = is_train
+        self.is_viz = is_viz
+
+        self.dataset_split = 'train' if is_train else 'val'
+        self.scene_names = self.load_scene_names(is_train)
+
+        if not is_train:
+            self.meshes = self.load_meshes()
+
+        self.hypersim_gt = self.load_hypersim_gt()
+
+
+    def load_hypersim_gt(self):
+        gt_filename = 'hypersim_gt_train.pt' if self.dataset_split == 'train' else 'hypersim_gt_val.pt'
+        print('loading GT file from', gt_filename)
+        gt = torch.load(gt_filename)
+        for scene_name in gt.keys():
+            good = torch.isfinite(gt[scene_name][0].sum(axis=1)) & torch.isfinite(gt[scene_name][1].sum(axis=1))
+
+            # Subsample GT to reduce memory usage.
+            if self.is_train:
+                good = good & (torch.rand(good.shape) < 0.5)
+            else:
+                good = good & (torch.rand(good.shape) < 0.1)
+            gt[scene_name] = [gt[scene_name][0][good], gt[scene_name][1][good]]
+        return gt
+
+    def load_meshes(self):
+        return torch.load('all_hypersim_val_meshes.pt')
+
+    def load_scene_names(self, is_train):
+        split = 'train' if is_train else 'test'
+        scene_names = []
+        with open(os.path.join(
+                self.args.hypersim_path,
+                'evermotion_dataset/analysis/metadata_images_split_scene_v1.csv'),'r') as f:
+            for line in f:
+                items = line.split(',')
+                if items[-1].strip() == split:
+                    scene_names.append(items[0])
+        scene_names = sorted(list(set(scene_names)))
+        print(len(scene_names), 'scenes loaded:', scene_names)
+        return scene_names
+
+    def is_corrupted_frame(self, frame):
+        return (
+            ('ai_003_001' in frame and 'cam_00' in frame)
+            or ('ai_004_009' in frame and 'cam_01' in frame)
+        )
+
+    def get_hypersim_data(self, index):
+        for retry in range(1000):
+            try:
+                if retry < 10:
+                    scene_name = self.scene_names[index % len(self.scene_names)]
+                else:
+                    scene_name = random.choice(self.scene_names)
+
+                frames = glob.glob(os.path.join(self.args.hypersim_path, scene_name, 'images/scene_cam_*_final_preview/*tonemap*'))
+                seen_frame = random.choice(frames)
+
+                if self.is_corrupted_frame(seen_frame):
+                    continue
+
+                seen_data = self.load_frame_data(seen_frame)
+                if not is_good_xyz(seen_data[0]):
+                    continue
+
+                cur_gt = self.hypersim_gt[scene_name]
+                gt_data = [cur_gt[0], cur_gt[1]]
+
+                if self.is_train:
+                    mesh_points = torch.zeros((1,))
+                else:
+                    mesh_points = sample_points_from_meshes(self.meshes[scene_name], 1000000)
+
+                # get camera positions
+                camera_positions = read_h5py(get_camera_pos_file_name_from_frame_name(seen_frame))
+                camera_position = camera_positions[int(seen_frame.split('.')[-3])]
+
+                # get camera orientations
+                cam_orientations = read_h5py(get_camera_orientation_file_name_from_frame_name(seen_frame))
+                cam_orientation = cam_orientations[int(seen_frame.split('.')[-3])]
+                cam_orientation = cam_orientation * (-1.0)
+
+                # rotate to camera direction
+                seen_data[0] = torch.matmul(seen_data[0], cam_orientation)
+                gt_data[0] = torch.matmul(gt_data[0], cam_orientation)
+
+                # shift to camera center
+                camera_position = torch.matmul(camera_position, cam_orientation)
+                seen_data[0] -= camera_position
+                gt_data[0] -= camera_position
+                # to meter
+                asset_to_meter_scale = read_scale_from_frame_name(seen_frame)
+                seen_data[0] = seen_data[0] * asset_to_meter_scale
+                gt_data[0] = gt_data[0] * asset_to_meter_scale
+
+                # get points GT
+                n_gt = 30000
+                in_front_of_cam = (gt_data[0][..., 2] > 0)
+                if in_front_of_cam.sum() < 1000:
+                    print('Warning! Not enough in front of cam.', in_front_of_cam.sum())
+                    continue
+                gt_data = [gt_data[0][in_front_of_cam], gt_data[1][in_front_of_cam]]
+
+                if in_front_of_cam.sum() < n_gt:
+                    selected = random.choices(range(gt_data[0].shape[0]), k=n_gt)
+                else:
+                    selected = random.sample(range(gt_data[0].shape[0]), n_gt)
+                gt_data = [gt_data[0][selected][None], gt_data[1][selected][None], torch.zeros((1,))]
+
+                if not self.is_train:
+                    mesh_points = torch.matmul(mesh_points, cam_orientation)
+                    mesh_points -= camera_position * asset_to_meter_scale
+                    in_front_of_cam = (mesh_points[..., 2] > 0)
+                    if in_front_of_cam.sum() < 1000:
+                        print('Warning! Not enough mesh in front of cam.', in_front_of_cam.sum())
+                        continue
+                    mesh_points = mesh_points[in_front_of_cam]
+                    if in_front_of_cam.sum() < n_gt:
+                        selected = random.choices(range(mesh_points.shape[0]), k=n_gt)
+                    else:
+                        selected = random.sample(range(mesh_points.shape[0]), n_gt)
+                    mesh_points = mesh_points[selected][None]
+                    mesh_points[..., 0] *= -1
+
+                seen_data[0][..., 0] *= -1
+                gt_data[0][..., 0] *= -1
+
+                seen_data[1] = seen_data[1].permute(2, 0, 1)
+
+                return seen_data, gt_data, mesh_points, scene_name
+            except Exception as e:
+                print(scene_name, 'loading failed', retry, e)
+
+
+    def __getitem__(self, index):
+
+        seen_data, gt_data, mesh_points, scene_name = self.get_hypersim_data(index)
+
+        # normalize the data
+        example_std = get_example_std(seen_data[0])
+        seen_data[0] = seen_data[0] / example_std
+        gt_data[0] = gt_data[0] / example_std
+        mesh_points = mesh_points / example_std
+
+        return (
+            seen_data,
+            gt_data,
+            mesh_points,
+            scene_name,
+        )
+
+    def load_frame_data(self, frame_path):
+        frame_xyz_path = frame_path.replace('final_preview/', 'geometry_hdf5/').replace('.tonemap.jpg', '.position.hdf5')
+        xyz = read_h5py(frame_xyz_path)
+        img = read_img(frame_path)
+
+        xyz, img = random_crop(
+            xyz, img,
+            is_train=self.is_train,
+        )
+        return [xyz, img]
+
+    def __len__(self) -> int:
+        if self.is_train:
+            return int(len(self.scene_names) * self.args.train_epoch_len_multiplier)
+        elif self.is_viz:
+            return len(self.scene_names)
+        else:
+            return int(len(self.scene_names) * self.args.eval_epoch_len_multiplier)
+
+
+def get_example_std(x):
+    x = x.reshape(-1, 3)
+    x = x[torch.isfinite(x.sum(dim=1))]
+    return x.std(dim=0).mean().detach()

util/hypersim_utils.py

@@ -0,0 +1,24 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import cv2
+import h5py
+
+import torch
+
+
+def read_h5py(filename):
+    with h5py.File(filename, "r") as f:
+        data = torch.tensor(f['dataset'][:], dtype=torch.float32)
+    return data
+
+
+def read_img(frame_path):
+    for retry in range(100):
+        img = cv2.imread(frame_path)
+        if img is not None:
+            return torch.tensor(img / 255.0, dtype=torch.float32)[..., [2, 1, 0]]
+        print('retry loading', retry, frame_path)

util/lars.py

@@ -0,0 +1,47 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# LARS optimizer, implementation from MoCo v3:
+# https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+
+import torch
+
+
+class LARS(torch.optim.Optimizer):
+    """
+    LARS optimizer, no rate scaling or weight decay for parameters <= 1D.
+    """
+    def __init__(self, params, lr=0, weight_decay=0, momentum=0.9, trust_coefficient=0.001):
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, trust_coefficient=trust_coefficient)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for g in self.param_groups:
+            for p in g['params']:
+                dp = p.grad
+
+                if dp is None:
+                    continue
+
+                if p.ndim > 1: # if not normalization gamma/beta or bias
+                    dp = dp.add(p, alpha=g['weight_decay'])
+                    param_norm = torch.norm(p)
+                    update_norm = torch.norm(dp)
+                    one = torch.ones_like(param_norm)
+                    q = torch.where(param_norm > 0.,
+                                    torch.where(update_norm > 0,
+                                    (g['trust_coefficient'] * param_norm / update_norm), one),
+                                    one)
+                    dp = dp.mul(q)
+
+                param_state = self.state[p]
+                if 'mu' not in param_state:
+                    param_state['mu'] = torch.zeros_like(p)
+                mu = param_state['mu']
+                mu.mul_(g['momentum']).add_(dp)
+                p.add_(mu, alpha=-g['lr'])

util/lr_decay.py

@@ -0,0 +1,76 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# ELECTRA https://github.com/google-research/electra
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+
+import json
+
+
+def param_groups_lrd(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
+    """
+    Parameter groups for layer-wise lr decay
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
+    """
+    param_group_names = {}
+    param_groups = {}
+
+    num_layers = len(model.blocks) + 1
+
+    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))
+
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+
+        # no decay: all 1D parameters and model specific ones
+        if p.ndim == 1 or n in no_weight_decay_list:
+            g_decay = "no_decay"
+            this_decay = 0.
+        else:
+            g_decay = "decay"
+            this_decay = weight_decay
+            
+        layer_id = get_layer_id_for_vit(n, num_layers)
+        group_name = "layer_%d_%s" % (layer_id, g_decay)
+
+        if group_name not in param_group_names:
+            this_scale = layer_scales[layer_id]
+
+            param_group_names[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+            param_groups[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+
+        param_group_names[group_name]["params"].append(n)
+        param_groups[group_name]["params"].append(p)
+
+    # print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))
+
+    return list(param_groups.values())
+
+
+def get_layer_id_for_vit(name, num_layers):
+    """
+    Assign a parameter with its layer id
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
+    """
+    if name in ['cls_token', 'pos_embed']:
+        return 0
+    elif name.startswith('patch_embed'):
+        return 0
+    elif name.startswith('blocks'):
+        return int(name.split('.')[1]) + 1
+    else:
+        return num_layers

util/lr_sched.py

@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """Decay the learning rate with half-cycle cosine after warmup"""
+    if epoch < args.warmup_epochs:
+        lr = args.lr * epoch / args.warmup_epochs 
+    else:
+        lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
+            (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)))
+    for param_group in optimizer.param_groups:
+        if "lr_scale" in param_group:
+            param_group["lr"] = lr * param_group["lr_scale"]
+        else:
+            param_group["lr"] = lr
+    return lr

util/misc.py

@@ -0,0 +1,496 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# MAE: https://github.com/facebookresearch/mae
+# --------------------------------------------------------
+
+import builtins
+import datetime
+import os
+import time
+from collections import defaultdict, deque
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch._six import inf
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        log_msg = [
+            header,
+            '[{0' + space_fmt + '}/{1}]',
+            'eta: {eta}',
+            '{meters}',
+            'time: {time}',
+            'data: {data}'
+        ]
+        if torch.cuda.is_available():
+            log_msg.append('max mem: {memory:.0f}')
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print('{} Total time: {} ({:.4f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    builtin_print = builtins.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        force = force or (get_world_size() > 8)
+        if is_master or force:
+            now = datetime.datetime.now().time()
+            builtin_print('[{}] '.format(now), end='')  # print with time stamp
+            builtin_print(*args, **kwargs)
+
+    builtins.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if args.dist_on_itp:
+        args.rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
+        args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
+        args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+        args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
+        os.environ['LOCAL_RANK'] = str(args.gpu)
+        os.environ['RANK'] = str(args.rank)
+        os.environ['WORLD_SIZE'] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print('Not using distributed mode')
+        setup_for_distributed(is_master=True)  # hack
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}, gpu {}'.format(
+        args.rank, args.dist_url, args.gpu), flush=True)
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True, verbose=False):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        if verbose:
+            print('norm:', norm, 'clip:', clip_grad)
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.)
+    device = parameters[0].grad.device
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
+    return total_norm
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler):
+    output_dir = Path(args.output_dir)
+    epoch_name = f'{epoch:05d}'
+    if loss_scaler is not None:
+        checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_without_ddp.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state)
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler):
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume, map_location='cpu')
+        print("Resume checkpoint %s" % args.resume)
+        print(model_without_ddp.load_state_dict(checkpoint['model'], strict=False))
+        if 'optimizer' in checkpoint and 'epoch' in checkpoint and not (hasattr(args, 'eval') and args.eval):
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if 'scaler' in checkpoint:
+                print(loss_scaler.load_state_dict(checkpoint['scaler']))
+            print("With optim & sched!")
+            print("start epoch:", args.start_epoch)
+
+
+def all_reduce_mean(x):
+    world_size = get_world_size()
+    if world_size > 1:
+        x_reduce = torch.tensor(x).cuda()
+        dist.all_reduce(x_reduce)
+        x_reduce /= world_size
+        return x_reduce.item()
+    else:
+        return x
+
+
+import torch.distributed as dist
+
+def get_world_size():
+    """
+    Get the size of the world.
+    """
+    if not dist.is_available():
+        return 1
+    if not dist.is_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+# def all_gather_unaligned(data):
+#     """
+#     Run all_gather on arbitrary picklable data (not necessarily tensors).
+#     Args:
+#         data: any picklable object
+#         group: a torch process group. By default, will use a group which
+#             contains all ranks on gloo backend.
+#     Returns:
+#         list[data]: list of data gathered from each rank
+#     """
+#     print('world', get_world_size())
+#     if get_world_size() == 1:
+#         return [data]
+
+#     # receiving Tensor from all ranks
+#     tensor_list = [
+#         torch.zeros_like(data) for _ in range(get_world_size())
+#     ]
+#     dist.all_gather(tensor_list, data)
+#     for tl in tensor_list:
+#         print(tl)
+#         print(tl.shape)
+#     return tensor_list
+
+import pickle
+def _serialize_to_tensor(data, group):
+    """
+    Seriialize the tensor to ByteTensor. Note that only `gloo` and `nccl`
+        backend is supported.
+    Args:
+        data (data): data to be serialized.
+        group (group): pytorch dist group.
+    Returns:
+        tensor (ByteTensor): tensor that serialized.
+    """
+
+    backend = dist.get_backend(group)
+    assert backend in ["gloo", "nccl"]
+    device = torch.device("cpu" if backend == "gloo" else "cuda")
+
+    buffer = pickle.dumps(data)
+    if len(buffer) > 1024 ** 3:
+        print(
+            "Rank {} trying to all-gather {:.2f} GB of data on device {}".format(
+                get_rank(), len(buffer) / (1024 ** 3), device
+            )
+        )
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to(device=device)
+    return tensor
+
+import functools
+@functools.lru_cache()
+def _get_global_gloo_group():
+    """
+    Return a process group based on gloo backend, containing all the ranks
+    The result is cached.
+    Returns:
+        (group): pytorch dist group.
+    """
+    if dist.get_backend() == "nccl":
+        return dist.new_group(backend="gloo")
+    else:
+        return dist.group.WORLD
+
+
+def _pad_to_largest_tensor(tensor, group):
+    """
+    Padding all the tensors from different GPUs to the largest ones.
+    Args:
+        tensor (tensor): tensor to pad.
+        group (group): pytorch dist group.
+    Returns:
+        list[int]: size of the tensor, on each rank
+        Tensor: padded tensor that has the max size
+    """
+    world_size = dist.get_world_size(group=group)
+    assert (
+        world_size >= 1
+    ), "comm.gather/all_gather must be called from ranks within the given group!"
+    local_size = torch.tensor(
+        [tensor.numel()], dtype=torch.int64, device=tensor.device
+    )
+    size_list = [
+        torch.zeros([1], dtype=torch.int64, device=tensor.device)
+        for _ in range(world_size)
+    ]
+    dist.all_gather(size_list, local_size, group=group)
+    size_list = [int(size.item()) for size in size_list]
+
+    max_size = max(size_list)
+
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    if local_size != max_size:
+        padding = torch.zeros(
+            (max_size - local_size,), dtype=torch.uint8, device=tensor.device
+        )
+        tensor = torch.cat((tensor, padding), dim=0)
+    return size_list, tensor
+
+def all_gather_unaligned(data, group=None):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors).
+    Args:
+        data: any picklable object
+        group: a torch process group. By default, will use a group which
+            contains all ranks on gloo backend.
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    if get_world_size() == 1:
+        return [data]
+    if group is None:
+        group = _get_global_gloo_group()
+    if dist.get_world_size(group) == 1:
+        return [data]
+
+    tensor = _serialize_to_tensor(data, group)
+
+    size_list, tensor = _pad_to_largest_tensor(tensor, group)
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    tensor_list = [
+        torch.empty((max_size,), dtype=torch.uint8, device=tensor.device)
+        for _ in size_list
+    ]
+    dist.all_gather(tensor_list, tensor, group=group)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer).to(data.device))
+
+    return data_list
+

util/pos_embed.py

@@ -0,0 +1,117 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# Position embedding utils
+# --------------------------------------------------------
+
+import numpy as np
+
+import torch
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+def get_1d_sincos_pos_embed_from_grid_torch(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = torch.arange(embed_dim // 2, device=pos.device).float()
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = torch.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = torch.sin(out) # (M, D/2)
+    emb_cos = torch.cos(out) # (M, D/2)
+
+    emb = torch.cat([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed

weights/co3dv2_all_categories.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca861bee4c2cb27acc6855da34227ce7026cf9eb275171da3c5a33976b3d86bd
+size 2423688373