Pixel-Aligned Recurrent Queries for Multi-View 3D Object Detection (ICCV23)
from yacs.config import CfgNode as CN

_C = CN()


# general configs
_C.SEED = 100
_C.MEMORY_GB = 230
_C.CHECKPOINT_PATH = None
_C.PRETRAINED_PATH = None # Load the pretrained model, which may have different keys from the current model
_C.NAME = "release"
_C.LOG_PATH = "/work/vig/yimingx/parq_logs"
_C.TAG = ""
_C.LOG_IMAGES = True
_C.LOG_IMAGES_FREQUENCY = 4800
_C.LOG_RANK_ZERO_ONLY = True


# lightning trainer
_C.TRAINER = CN()
_C.TRAINER.PROFILER = 'simple'
_C.TRAINER.ACCELERATOR = 'gpu'
_C.TRAINER.GPUS = 2
_C.TRAINER.NUM_NODES = 1
_C.TRAINER.ACCUMULATE_GRAD_BATCHES = 1
_C.TRAINER.MAX_EPOCHS = 100
_C.TRAINER.LOG_EVERY_N_STEPS = 100
_C.TRAINER.GRADIENT_CLIP_VAL = 1.0
_C.TRAINER.RELOAD_DATALOADERS_EVERY_N_EPOCHS = 0
_C.TRAINER.REPLACE_SAMPLER_DDP = True
_C.TRAINER.OVERFIT_BATCHES = 0.0
_C.TRAINER.AUTO_SCALE_BATCH_SIZE = 'binsearch'
_C.TRAINER.CHECK_VAL_EVERY_N_EPOCH = 1
_C.TRAINER.PRECISION = 32
_C.TRAINER.VAL_CHECK_INTERVAL = 1.0
_C.TRAINER.LIMIT_VAL_BATCHES = 1.0
_C.TRAINER.LIMIT_TRAIN_BATCHES = 1.0


# callback
_C.CALLBACK = CN()
_C.CALLBACK.MONITOR = 'val/metrics/0.5_f1' # name of the logged metric which determines when model is improving
_C.CALLBACK.SAVE_TOP_K = 3 # save k best models (determined by above metric)
_C.CALLBACK.SAVE_LAST = True # additionaly always save model from last epoch
_C.CALLBACK.VERBOSE = False
_C.CALLBACK.DIRPATH = None
_C.CALLBACK.FILENAME = None
_C.CALLBACK.AUTO_INSERT_METRIC_NAME = False
_C.CALLBACK.MODE = 'max'

# optimization
_C.OPTIMIZER = CN()
_C.OPTIMIZER.NAME = 'adamw'
_C.OPTIMIZER.LEARNING_RATE = 1e-4 # cosine-annealing with restarts; by default just a single cycle
_C.OPTIMIZER.CYCLE_MULT = 1
_C.OPTIMIZER.WARMUP_EPOCHS = 0
_C.OPTIMIZER.NUM_RESTARTS = 1
_C.OPTIMIZER.IGNORE_FROZEN_PARAMS = True # filter out all model parameters with requires_grad=False
_C.OPTIMIZER.AUTOSCALE_LR = True # Auto-scale the learning rate based on "Accurate Large Minibatch SGD ..."


# datamodule
_C.DATAMODULE = CN()
_C.DATAMODULE.DATA_PATH = './data/scannet/scans'
_C.DATAMODULE.TRAIN_ANNOTATION_PATH = './data/scannet/scan2cad_box3d_anno_view3_overlap/scannet_train_gt_roidb.pkl'
_C.DATAMODULE.VAL_ANNOTATION_PATH = './data/scannet/scan2cad_box3d_anno_view3_overlap/scannet_val_gt_roidb.pkl'
_C.DATAMODULE.BATCH_SIZE = 1
_C.DATAMODULE.NUM_WORKERS = 1
_C.DATAMODULE.NUM_FRAMES_PER_SNIPPET = 3
_C.DATAMODULE.SHUFFLE = True
_C.DATAMODULE.GRAVITY_ALIGNED = True

# model
feature_dim = 1024
_C.MODEL = CN()
# 2d backbone
_C.MODEL.BACKBONE2D = CN()
_C.MODEL.BACKBONE2D.RESNET_NAME = 'resnet50'
_C.MODEL.BACKBONE2D.LAYER = 0
_C.MODEL.BACKBONE2D.FREEZE = False

# tokenizer
_C.MODEL.TOKENIZER = CN()
_C.MODEL.TOKENIZER.OUT_CHANNELS = feature_dim
_C.MODEL.TOKENIZER.PATCH_SIZE = 1 # patch size when using tokenization
_C.MODEL.TOKENIZER.RAY_POINTS_SCALE = [-2, 2, -1.5, 0, 0.25, 4.25] # [min_x, max_x, min_y, max_y, min_z, max_z] used to normalize the points along each ray
_C.MODEL.TOKENIZER.NUM_SAMPLES = 64 # number of ray points
_C.MODEL.TOKENIZER.MIN_DEPTH = 0.25 # min depth of ray points
_C.MODEL.TOKENIZER.MAX_DEPTH = 5.25 # max depth of ray points

# PARQ decoder
_C.MODEL.DECODER = CN()
_C.MODEL.DECODER.DIM_IN = feature_dim
_C.MODEL.DECODER.NUM_QUERIES = 128
_C.MODEL.DECODER.NUM_SEMCLS = 9 # number of semantic classes, e.g. 9 on scannet, 17 on ARKitScenes
_C.MODEL.DECODER.BOX_SIZE = [1, 1, 1]
_C.MODEL.DECODER.LOSS_WEIGHT = [5.0, 5.0, 5.0, 1.0] # center, size, rotation, category
_C.MODEL.DECODER.CONF_THRESH = 0.1 # confidence threshold, box will be filtered if its score is lower than this thresh
_C.MODEL.DECODER.MEAN_SIZE_PATH = None # the path to the file of the mean size for each object in scannet
_C.MODEL.DECODER.EVAL_TYPE = 'f1' # using f1 score to evaluate the model, follow ODAM https://arxiv.org/abs/2108.10165
_C.MODEL.DECODER.ENABLE_NMS = True # if True, the model will use nms to filter out potential duplicate boxes
_C.MODEL.DECODER.SHARE_MLP_HEADS = True # detection head share weight in each iteration
_C.MODEL.DECODER.FOR_VIS = False # if True, the model will preserve all outputs from one snippet for visualization without ant filtering, e.g. filter out far-away boxes or using nms
_C.MODEL.DECODER.TRACK_SCALE = [-1.5, 1.5, -2, 1, 0, 2] # only track boxes within this scale, because the quality of predictions is usually bad due to truncation... also used for other baselines

_C.MODEL.DECODER.TRANSFORMER = CN()
_C.MODEL.DECODER.TRANSFORMER.DEC_DIM = feature_dim # decoder dimension
_C.MODEL.DECODER.TRANSFORMER.DEC_HEADS = 4 # number of decoder heads
_C.MODEL.DECODER.TRANSFORMER.DEC_FFN_DIM = 768 # decoder feedforward dimension
_C.MODEL.DECODER.TRANSFORMER.DEC_LAYERS = 8 # number of decoder layers
_C.MODEL.DECODER.TRANSFORMER.DROPOUT_RATE = 0.1 # dropout_rate=0.0 disables dropout
_C.MODEL.DECODER.TRANSFORMER.QUERIES_DIM = feature_dim # same as dec_dim but needed for other configs
_C.MODEL.DECODER.TRANSFORMER.SCALE = _C.MODEL.TOKENIZER.RAY_POINTS_SCALE
_C.MODEL.DECODER.TRANSFORMER.SHARE_WEIGHTS = True




def update_config(cfg, args):
    cfg.defrost()
    cfg.merge_from_file(args.cfg)
    cfg.merge_from_list(args.opts)

    cfg.freeze()


def check_config(cfg):
    pass
config/default.py
ymingxie/PARQ
from .scannet_dataset import *

datasets/__init__.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
import os
import pickle

import numpy as np
import torch

from PIL import Image

from pytorch_lightning.core import LightningDataModule

from utils import collate
from datasets.transforms import ScanNetBaseTransform

    

class ScanNetDataModule(LightningDataModule):
    """
    data module for scannet snapshots.
    """

    def __init__(self, cfg):
        super().__init__()

        self.train_anno_path = cfg.TRAIN_ANNOTATION_PATH
        self.val_anno_path = cfg.VAL_ANNOTATION_PATH
        self.data_path = cfg.DATA_PATH
        self.num_workers = cfg.NUM_WORKERS
        self.batch_size = cfg.BATCH_SIZE
        self.drop_last = True
        self.transform = ScanNetBaseTransform(nviews=cfg.NUM_FRAMES_PER_SNIPPET, gravity_aligned=cfg.GRAVITY_ALIGNED)
        self.shuffle = cfg.SHUFFLE
        self.num_framesets_per_snippet=cfg.NUM_FRAMES_PER_SNIPPET

    def train_dataloader(self):
        return self.getLoader(self.data_path, self.train_anno_path)

    def val_dataloader(self):
        return self.getLoader(self.data_path, self.val_anno_path)

    def test_dataloader(self):
        return self.getLoader(self.data_path, self.val_anno_path)

    def getLoader(self, data_path, gt_path):
        dataset = ScanNetDataset(
            data_path=data_path,
            gt_path=gt_path,
            transform=self.transform,
            num_framesets_per_snippet=self.num_framesets_per_snippet,
        )

        params = {
            "batch_size": self.batch_size,
            "num_workers": self.num_workers,
            "collate_fn": collate,
            "drop_last": self.drop_last,
            "shuffle": self.shuffle,
        }
        dataloader = torch.utils.data.DataLoader(dataset, **params)
        return dataloader



class ScanNetDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        data_path: str,
        gt_path: str,
        transform=None,
        shuffle=False,
        num_framesets_per_snippet=3,
    ):
        self.data_path = data_path
        self.gt_path = gt_path
        if transform is None:
            transform = ScanNetBaseTransform()
        self.transform = transform
        self.shuffle = shuffle
        self.item_list = self.get_roidb()
        self.anno_path = os.path.join(*(self.gt_path.split("/")[:-1] + ["scene_anno"]))
        self.scene_cache = {}
        self.max_cache = 100
        self.num_framesets_per_snippet = num_framesets_per_snippet
        
    def get_roidb(self):
        with open(self.gt_path, "rb") as f:
            item_list = pickle.load(f)
        print(
            "roidb loaded from {}, totally {} samples".format(
                self.gt_path, len(item_list)
            )
        )
        # item_list = item_list[:100]
        return item_list

    def read_img(self, filepath):
        img = Image.open(filepath)
        return img

    def read_anno(self, scene_name, snippet_id):
        if scene_name not in self.scene_cache.keys():
            if len(self.scene_cache) > self.max_cache:
                self.scene_cache = {}
            anno_path = os.path.join(self.anno_path, scene_name + ".pkl")
            with open(anno_path, "rb") as f:
                anno = pickle.load(f)
            self.scene_cache[scene_name] = anno
        return self.scene_cache[scene_name][snippet_id]

    def __len__(self):
        return len(self.item_list)

    def __getitem__(self, idx):
        """
        Returns a dict with following keys:
            scene_name:         str, e.g. scene0000_00, the scene name of the snippet
            rgb_img:            tensor (N,3,H,W), the RGB image sequence, N is the number of frames
            sym:                tensor (B,), symmetry type of the object, B is number of objects
            camera:             camera intrinsics 
            obb_padded:         ground truth bounding boxes
            T_world_camera:     Pose (N, 12), pose matrixes from camera to world
            T_world_pseudoCam:  Pose (N, 12), pose matrixes from pseudo camera to world
            T_camera_pseudoCam: Pose (N, 12), pose matrixes from pseudo camera to camera
            T_world_local:      Pose (N, 12), pose matrixes from local to world
        Note: pseudoCam coordinate is the gravity-aligned camera coordinate. 
        We rotate the camera coordinate to make it gravity-aligned.
        """
        item = self.item_list[idx]
        scene_name = item["scene_name"]
        snippet_id = item["snippet_id"]
        item = self.read_anno(item["scene_name"], item["snippet_id"])

        image_ids = item["image_ids"]
        T_scan_camera = item["T_scan_camera"]
        intrinsic = item["intrinsic"]
        
        num_frames = len(image_ids)
        if self.num_framesets_per_snippet != num_frames:
            replace = True if self.num_framesets_per_snippet > num_frames else False
            choose = np.random.choice(num_frames, self.num_framesets_per_snippet, replace=replace)
            choose = np.sort(choose)
            image_ids = [image_ids[choo] for choo in choose]
            T_scan_camera = [T_scan_camera[choo] for choo in choose]
            intrinsic = [intrinsic[choo] for choo in choose]
            
        annos = item["annotations"]
        bboxes = np.array(annos["bboxes"])
        T_scan_object = np.array(annos["T_scan_object"])
        label = annos["label"]
        sym = annos["sym"]
        sym_mapping = {
            "__SYM_NONE": 0,
            "__SYM_ROTATE_UP_2": 1,
            "__SYM_ROTATE_UP_4": 2,
            "__SYM_ROTATE_UP_INF": 3,
        }
        for i in range(len(sym)):
            if sym[i] in sym_mapping.keys():
                sym[i] = sym_mapping[sym[i]]
        # padding sym
        padding_size = 50
        for i in range(padding_size - len(sym)):
            sym.append(-1)

        sym = np.array(sym)

        imgs = []
        for id_ in image_ids:
            imgs.append(
                self.read_img(
                    os.path.join(
                        self.data_path,
                        scene_name,
                        "color",
                        'frame-{0:06d}.color.jpg'.format(id_),
                    )
                )
            )

        out_dict = {
            "scene_name": scene_name,
            "snippet_id": snippet_id,
            "image_ids": image_ids,
            "rgb_img": imgs,
            "bboxes": bboxes,
            "intrinsics": np.copy(intrinsic),
            "T_world_camera": T_scan_camera,
            "T_world_object": T_scan_object,
            "label": label,
            "sym": sym,
        }

        out_dict = self.transform(out_dict)
        out_dict.pop("bboxes")
        out_dict.pop("T_world_object")
        out_dict.pop("label")
        out_dict.pop("intrinsics")
        return out_dict

datasets/scannet_dataset.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
import numpy as np
import torch
from PIL import Image, ImageOps
from utils import Camera, Obb3D, Pose
from typing import Literal


UP_W = {
    "SCANNET": [0, 0, 1],
}

class GravityAligned(object):
    """
    Generate pseudoCam coordinates, which is the gravity-aligned camera coordinate
    """
    def __init__(self, dataset_type: Literal["SCANNET"]):
        self.dataset_type = dataset_type
        assert dataset_type in UP_W.keys()
        pass

    def __call__(self, batch):
        T_world_camera = batch["T_world_camera"]
        T_world_pseudoCam = self.camera_to_gravity_aligned(T_world_camera.matrix)
        T_world_pseudoCam = Pose.from_4x4mat(T_world_pseudoCam)
        batch["T_world_pseudoCam"] = T_world_pseudoCam
        batch["T_camera_pseudoCam"] = T_world_camera.inverse() @ T_world_pseudoCam

        return batch


    def camera_to_gravity_aligned(self, T_world_camera):
        up_w = UP_W[self.dataset_type]

        up_w = torch.tensor(up_w, device=T_world_camera.device).float()
        T_wv = torch.clone(T_world_camera)

        camForward = T_world_camera[..., :3, 2]
        R_wv = torch.zeros_like(T_wv[..., :3, :3], device=T_world_camera.device)
        R_wv[..., 1] = up_w
        R_wv[..., 2] = self.normalize(self.reject(camForward, up_w))
        R_wv[..., 0] = self.normalize(torch.cross(R_wv[..., 1], R_wv[..., 2], dim=-1))
        T_wv[..., :3, :3] = R_wv
        return T_wv


    def normalize(self, v):
        """normalize' a vector, in the traditional linear algebra sense."""
        norm = torch.norm(v, dim=-1, keepdim=True)
        if (norm == 0).any():
            return v
        return v / norm


    def reject(self, A, B):
        """Create a 'projection', and subract it from the original vector"""
        project = self.bdot(A, self.normalize(B)) * self.normalize(B)
        return A - project


    def bdot(self, a, b):
        return (a.unsqueeze(-2) @ b.unsqueeze(-1)).squeeze(-1)


def pad_scannet(img, intrinsics):
    """Scannet images are 1296x968 but 1296x972 is 4x3
    so we pad vertically 4 pixels to make it 4x3
    """

    w, h = img.size
    if w == 1296 and h == 968:
        img = ImageOps.expand(img, border=(0, 2))
        intrinsics[1, 2] += 2
    return img, intrinsics


class ResizeImage(object):
    """Resize everything to given size.
    Intrinsics are assumed to refer to image prior to resize.
    After resize everything (ex: depth) should have the same intrinsics
    """

    def __init__(self, size):
        self.size = size

    def __call__(self, data):
        for i, im in enumerate(data["rgb_img"]):
            # intrinsic = np.copy(data["intrinsic"][i])
            im, intrinsics = pad_scannet(im, data["intrinsics"][i])
            w, h = im.size
            im = im.resize(self.size, Image.BILINEAR)
            intrinsics[0, :] /= w / self.size[0]
            intrinsics[1, :] /= h / self.size[1]

            im = np.array(im, dtype=np.float32)
            # if load with opencv
            # im = im[..., [2, 1, 0]]
            data["rgb_img"][i] = im
            data["intrinsics"][i] = intrinsics
        return data

    def __repr__(self):
        return self.__class__.__name__ + "(size={0})".format(self.size)


class Compose(object):
    """Apply a list of transforms sequentially"""

    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, data):
        for transform in self.transforms:
            data = transform(data)
        return data


class ToTensor(object):
    """Convert to torch tensors"""

    def __call__(self, data):
        data["rgb_img"] = torch.Tensor(
            np.stack(data["rgb_img"]).transpose([0, 3, 1, 2])
        )
        data["intrinsics"] = torch.Tensor(data["intrinsics"])
        data["T_world_camera"] = torch.Tensor(data["T_world_camera"])
        if data["bboxes"] is not None:
            data["bboxes"] = torch.Tensor(data["bboxes"])
            data["T_world_object"] = torch.Tensor(data["T_world_object"])
            data["label"] = torch.Tensor(data["label"])
            data["sym"] = torch.Tensor(data["sym"])
        return data


class Convert2Objects:
    """
    Wrap the 3d object, camera parameters, and poses. See utils/wrapper.py for details
    """
    def __init__(self, nviews=9):
        self.nviews = nviews

    def __call__(self, data):
        t, c, h, w = data["rgb_img"].shape
        intrinsics = data["intrinsics"][0]
        # fx, fy, cx, cy, w, h
        camera_param = torch.stack(
            [
                torch.Tensor([w]),
                torch.Tensor([h]),
                intrinsics[0, 0].unsqueeze(0),
                intrinsics[1, 1].unsqueeze(0),
                intrinsics[0, 2].unsqueeze(0),
                intrinsics[1, 2].unsqueeze(0),
            ],
            dim=-1
        )
        camera_param = camera_param.expand(self.nviews, -1)

        data["camera"] = Camera(camera_param)

        if data["bboxes"] is not None:
            T_world_object = Pose.from_4x4mat(data["T_world_object"])._data
            obbs = Obb3D.separate_init(
                bb3_object=data["bboxes"],
                T_world_object=T_world_object,
                sem_id=data["label"],
            )
            data["obbs_padded"] = obbs.add_padding()
        else:
            data["obbs_padded"] = None

        data["T_world_camera"] = Pose.from_4x4mat(data["T_world_camera"].float())

        return data


class Normalize:
    def __init__(self):
        self.pixel_mean = torch.Tensor([103.53, 116.28, 123.675]).view(-1, 1, 1)
        self.pixel_std = torch.Tensor([1.0, 1.0, 1.0]).view(-1, 1, 1)

    def __call__(self, data):
        """Normalizes the RGB images to the input range"""
        # data["img_rgb"] = (
        #     data["img_rgb"] - self.pixel_mean.type_as(data["img_rgb"])
        # ) / self.pixel_std.type_as(data["img_rgb"])
        data["rgb_img"] = data["rgb_img"] / 255
        return data


class SnippetLocal:
    """
    Select a coordinate as the local coordinate in the snippet. We will predict 3d objects in this local coordinate
    """
    def __init__(
        self,         
        frame_selection: float = 0.5,
    ):    
        self.frame_selection = frame_selection

    def __call__(self, batch):
        with torch.no_grad():
            if "T_world_pseudoCam" in batch.keys():
                Ts_world_psuedoCam = batch["T_world_pseudoCam"]
                T = Ts_world_psuedoCam.shape[0]
                t = int(T * self.frame_selection)
                batch["T_world_local"] = Ts_world_psuedoCam[t, :].clone().unsqueeze(0)
        return batch


class ScanNetBaseTransform(object):
    """
    Transforms for ScanNet dataset.
    """

    def __init__(self, nviews=9, gravity_aligned=True):
        transform = []
        transform += [
            ResizeImage((320, 240)), # resize image to 320x240
            ToTensor(),              # convert to torch tensor
            Normalize(),             # normalize the image
        ]
        transform += [
            Convert2Objects(nviews), # wrap the 3d object, camera parameters, and poses into a class. See utils/wrapper.py for details
        ]
        if gravity_aligned:
            transform += [
            GravityAligned(dataset_type='SCANNET') # generate pseudoCam coordinates, which is the gravity-aligned camera coordinate
            ]
        transform += [
            SnippetLocal(), # select a coordinate as the local coordinate in the snippet. We will predict 3d objects in this local coordinate
        ]
        self.transforms = Compose(transform)

    def __call__(self, batch):
        batch = self.transforms(batch)
        return batch

datasets/transforms.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and its affiliates. All Rights Reserved
import argparse
import faulthandler
import logging
import time

from model import PARQ
import torch
from config import cfg, update_config
from datasets.scannet_dataset import ScanNetDataModule

# Improve multiprocess debugging
faulthandler.enable(all_threads=True)
logging.basicConfig(level=logging.INFO)


def test_model(cfg):

    # create datamodule
    data_module = ScanNetDataModule(cfg.DATAMODULE)
    loader = data_module.val_dataloader()
    model = PARQ(cfg)

    if cfg.CHECKPOINT_PATH is not None:
        state_dict = torch.load(cfg.CHECKPOINT_PATH)
        state_dict = state_dict["state_dict"]
        model.load_state_dict(state_dict, strict=True)

    model.cuda()
    model.eval()
    time_all = []
    gpu_mem_usage = []
    for i, batch in enumerate(loader):
        print(i)
        for key, value in batch.items():
            try:
                batch[key] = value.cuda()
            except:
                pass
        print(batch["scene_name"])
        start_time = time.time()
        with torch.no_grad():
            loss = model.validation_step(batch, 0)
        print("inference time {}".format(time.time() - start_time))
        time_all.append(time.time() - start_time)
        average = torch.tensor(time_all).mean()
        print("average time {}".format(average.item()))
        print("loss {}".format(loss.data.cpu().item()))
        # torch.cuda.empty_cache()
        # gpu_mem_usage.append(torch.cuda.memory_reserved())
        # summary_text = f"""
        #         average time: {average.item()} 
        #         Average GPU memory usage (GB): {sum(gpu_mem_usage) / len(gpu_mem_usage) / (1024 ** 3)} 
        #         Max GPU memory usage (GB): {max(gpu_mem_usage) / (1024 ** 3)} 
        #     """
        # print(summary_text)
        del loss
    ap = model.validation_epoch_end(None)

    for key, value in ap.items():
        print(key)
        print(value)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='A PyTorch Implementation of PARQ')

    # general
    parser.add_argument('--cfg',
                        help='experiment configure file name',
                        required=True,
                        type=str)

    parser.add_argument('--CHECKPOINT_PATH',
                        help='checkpoint path',
                        type=str)

    parser.add_argument('opts',
                        help="Modify config options using the command-line",
                        default=None,
                        nargs=argparse.REMAINDER)
    
    # parse arguments and check
    args = parser.parse_args()
    update_config(cfg, args)

    print("training with pytorch lightning")
    test_model(cfg)

eval.py
ymingxie/PARQ
from .parq_lightning import *  # noqa

model/__init__.py
ymingxie/PARQ

# Copyright (c) Meta Platforms, Inc. and its affiliates. All Rights Reserved
# This file is derived from [3DETR](https://github.com/facebookresearch/3detr/blob/main/models/helpers.py).
# Originating Author: Ishan Misra

# Original header:
# Copyright (c) Facebook, Inc. and its affiliates.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from functools import partial

import torch


class BatchNormDim1Swap(torch.nn.BatchNorm1d):
    """
    Used for nn.Transformer that uses a HW x N x C rep
    """

    def forward(self, x):
        """
        x: HW x N x C
        permute to N x C x HW
        Apply BN on C
        permute back
        """
        hw, n, c = x.shape
        x = x.permute(1, 2, 0)
        x = super(BatchNormDim1Swap, self).forward(x)
        # x: n x c x hw -> hw x n x c
        x = x.permute(2, 0, 1)
        return x


NORM_DICT = {
    "bn": BatchNormDim1Swap,
    "bn1d": torch.nn.BatchNorm1d,
    "id": torch.nn.Identity,
    "ln": torch.nn.LayerNorm,
}

ACTIVATION_DICT = {
    "relu": torch.nn.ReLU,
    "gelu": torch.nn.GELU,
    "leakyrelu": partial(torch.nn.LeakyReLU, negative_slope=0.1),
}

WEIGHT_INIT_DICT = {
    "xavier_uniform": torch.nn.init.xavier_uniform_,
}


class GenericMLP(torch.nn.Module):
    def __init__(
        self,
        input_dim,
        hidden_dims,
        output_dim,
        norm_fn_name=None,
        activation="relu",
        use_conv=False,
        dropout=None,
        hidden_use_bias=False,
        output_use_bias=True,
        output_use_activation=False,
        output_use_norm=False,
        weight_init_name=None,
    ):
        super().__init__()
        activation = ACTIVATION_DICT[activation]
        norm = None
        if norm_fn_name is not None:
            norm = NORM_DICT[norm_fn_name]
        if norm_fn_name == "ln" and use_conv:
            norm = lambda x: torch.nn.GroupNorm(1, x)  # easier way to use LayerNorm

        if dropout is not None:
            if not isinstance(dropout, list):
                dropout = [dropout for _ in range(len(hidden_dims))]

        layers = []
        prev_dim = input_dim
        for idx, x in enumerate(hidden_dims):
            if use_conv:
                layer = torch.nn.Conv1d(prev_dim, x, 1, bias=hidden_use_bias)
            else:
                layer = torch.nn.Linear(prev_dim, x, bias=hidden_use_bias)
            layers.append(layer)
            if norm:
                layers.append(norm(x))
            layers.append(activation())
            if dropout is not None:
                layers.append(torch.nn.Dropout(p=dropout[idx]))
            prev_dim = x
        if use_conv:
            layer = torch.nn.Conv1d(prev_dim, output_dim, 1, bias=output_use_bias)
        else:
            layer = torch.nn.Linear(prev_dim, output_dim, bias=output_use_bias)
        layers.append(layer)

        if output_use_norm:
            layers.append(norm(output_dim))

        if output_use_activation:
            layers.append(activation())

        self.layers = torch.nn.Sequential(*layers)

        if weight_init_name is not None:
            self.do_weight_init(weight_init_name)

    def do_weight_init(self, weight_init_name):
        func = WEIGHT_INIT_DICT[weight_init_name]

        for (_, param) in self.named_parameters():
            if param.dim() > 1:  # skips batchnorm/layernorm
                func(param)

    def forward(self, x):
        output = self.layers(x)
        return output

model/generic_mlp.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
from functools import partial
from typing import List

import torch
from pytorch_lightning.utilities import rank_zero_only

from model.generic_mlp import GenericMLP
from model.transformer_parq import Transformer
from utils.parq_utils import (  # noqa
    BoxProcessor,
    compute_rotation_matrix_from_ortho6d,
    draw_detections,
    F1Calculator,
    HungarianMatcherModified,
    nms,
    rot_to_6d,
    roty,
)

from utils import Obb3D, Pose

import copy


def _get_clones(module, N):
    return torch.nn.ModuleList([copy.deepcopy(module) for i in range(N)])


class PARQDecoder(torch.nn.Module):
    """
    PARQ Decoder, including one-layer transformer decoder and mlp heads
    """

    def __init__(
        self,
        cfg
    ):
        super(PARQDecoder, self).__init__()
        self.dim_in = cfg.DIM_IN
        self.num_queries = cfg.NUM_QUERIES
        self.num_semcls = cfg.NUM_SEMCLS
        self.loss_weight = cfg.LOSS_WEIGHT
        # NOTE following the official DETR rep0, bg_cls_weight means
        # relative classification weight of the no-object class.
        self.class_weight = torch.ones(self.num_semcls + 1)  # * class_weight
        # set background class as the last indice
        self.class_weight[self.num_semcls] = 0.1
        self.for_vis = False
        self.track_scale = cfg.TRACK_SCALE

        self.share_mlp_heads = cfg.SHARE_MLP_HEADS
        self.build_mlp_heads(
            self.num_semcls, self.dim_in, cfg.MEAN_SIZE_PATH, mlp_dropout=0.3
        )
        
        self.parq_module = Transformer(
                 dec_dim=cfg.TRANSFORMER.DEC_DIM, 
                 queries_dim=cfg.TRANSFORMER.QUERIES_DIM, 
                 dec_heads=cfg.TRANSFORMER.DEC_HEADS, 
                 dec_layers=cfg.TRANSFORMER.DEC_LAYERS, 
                 dec_ffn_dim=cfg.TRANSFORMER.DEC_FFN_DIM, 
                 dropout_rate=cfg.TRANSFORMER.DROPOUT_RATE,
                 scale=cfg.TRANSFORMER.SCALE,
                 share_weights=cfg.TRANSFORMER.SHARE_WEIGHTS)
        self.parq_module.decoder.mlp_heads = self.mlp_heads
        self.parq_module.decoder.box_processor = self.box_processor
        self.refpoint = torch.nn.Embedding(self.num_queries, 3)

        # match between the prediction and gt
        self.matcher = HungarianMatcherModified(cost_class=2, cost_bbox=0.25)

        if not isinstance(cfg.EVAL_TYPE, list):
            eval_type = [cfg.EVAL_TYPE]
        self.metrics_calculator = []
        for et in eval_type:
            if et == "f1":
                self.metrics_calculator.append(
                    F1Calculator(cfg.CONF_THRESH)
                )

        self.enable_nms = cfg.ENABLE_NMS

    def build_mlp_heads(
        self, num_semcls, decoder_dim, mean_size_path=None, mlp_dropout=None,
    ):
        """
        build mlp head to regress the parameters of the boxes
        """
        mlp_func = partial(
            GenericMLP,
            norm_fn_name="ln",
            activation="relu",
            use_conv=True,
            hidden_dims=[decoder_dim, decoder_dim],
            dropout=0.0,
            input_dim=decoder_dim,
        )
        mlp_func_small = partial(
            GenericMLP,
            norm_fn_name="ln",
            activation="relu",
            use_conv=True,
            hidden_dims=[],
            dropout=mlp_dropout,
            input_dim=decoder_dim,
        )

        # Semantic class of the box
        # add 1 for background/not-an-object class
        semcls_head = mlp_func_small(output_dim=num_semcls + 1)

        # geometry of the box
        center_head = mlp_func(output_dim=3)
        size_head = mlp_func_small(output_dim=3)
        # ortho6d: https://arxiv.org/pdf/1812.07035.pdf
        rotation_head = mlp_func(output_dim=6)

        if not self.share_mlp_heads:
            semcls_head = _get_clones(semcls_head, self.model_cfg.dec_layers)
            center_head =  _get_clones(center_head, self.model_cfg.dec_layers)
            size_head = _get_clones(size_head, self.model_cfg.dec_layers)
            rotation_head = _get_clones(rotation_head, self.model_cfg.dec_layers)
            
        mlp_heads = [
            ("sem_cls_head", semcls_head),
            ("center_head", center_head),
            ("size_head", size_head),
            ("rotation_head", rotation_head),
        ]
        self.mlp_heads = torch.nn.ModuleDict(mlp_heads)
        self.box_processor = BoxProcessor(num_semcls, mean_size_path)
    
    def forward(self, intput_tokens, camera, T_camera_pseudoCam, T_world_pseudoCam, T_world_local):
        """
        aplly PARQ module, including transformer decoder and mlp heads
        input:
            input_tokens:       tensor, (B, N, C), the output of tokenization of (image features + ray positional encoding)
            camera:             camera intrinsics, (B, 6), width, height, fx, fy, cx, cy.
            T_camera_pseudoCam: pose, (B, 12), pose matrix from pseudo camera to camera
            T_world_pseudoCam:  pose, (B, 12), pose matrix from pseudo camera to world
            T_world_local:      pose, (B, 12), pose matrix from local to world
        output:
            box_prediction_list (parameters of 3d boxes): 
            [
                iteration 1 (dict):
                    pred_logits:          (B, N, num_semcls + 1), the classification logits (including background class)
                    center_unnormalized:  (B, N, 3), the predicted center of the box
                    size_unnormalized:    (B, N, 3), the predicted size of the box
                    ortho6d:              (B, N, 6), the predicted rotation of the box
                    sem_cls_prob:         (B, N, num_semcls + 1), the softmax of pred_logits
                    coord_pos:            (B, N, 3), the position of reference points, used for matcher
                iteration 2 (dict):
                    ...
                ...
            ]
        """
        meta_data = {"T_world_local": T_world_local,
                     "camera": camera,
                     "T_camera_pseudoCam": T_camera_pseudoCam,
                     "T_world_pseudoCam": T_world_pseudoCam}
        box_prediction_list= self.parq_module(intput_tokens, self.refpoint.weight, meta_data)
        return box_prediction_list

    def parse_target(
        self, bb3_target: Obb3D, T_world_local: Pose
    ):
        """
        parse the target 3d boxes to the format of the model
        """
        T_local_world = T_world_local.inverse()
        bb3_target_list = []
        sem_id = bb3_target.sem_id
        bs = sem_id.shape[0]
        for i in range(bs):
            bb3_target_i = bb3_target[i].remove_padding()
            # 3d boxes in object coords (center: 0,0,0; no rotation)
            # transform from object coord to local coord
            T_local_object = T_local_world[i] @ bb3_target_i.T_world_object
            size = bb3_target_i.bb3_size
            center_local = T_local_object.transform(
                bb3_target_i.bb3_center_object.unsqueeze(1)
            ).squeeze(1)
            bb3corners_local = T_local_object.transform(bb3_target_i.bb3corners_object)

            gt_ortho6d = rot_to_6d(T_local_object.R)

            gt_corners_world = bb3_target_i.T_world_object.transform(
                bb3_target_i.bb3corners_object
            )

            bb3_target_dict = {
                "labels": bb3_target_i.sem_id.squeeze(-1).long(),
                "center": center_local,
                "size": size,
                "T_rig_object": T_local_object.matrix.view(-1, 4, 4),
                "gt_corners": bb3corners_local,
                "gt_ortho6d": gt_ortho6d,
                "gt_corners_world": gt_corners_world,
                "T_world_object": bb3_target_i.T_world_object,
            }
            bb3_target_list.append(bb3_target_dict)
        return bb3_target_list

    def rotation_loss_with_sym(self, rot_predict, rot_target, sym_):
        """
          resolve symmetry for rotation loss
          sym_mapping = {
            "__SYM_NONE": 0,
            "__SYM_ROTATE_UP_2": 1,
            "__SYM_ROTATE_UP_4": 2,
            "__SYM_ROTATE_UP_INF": 3,
        }
        """
        rot = roty
        # --> resolve symmetry
        rot_loss_list = []
        for o in range(sym_.shape[0]):
            if sym_[o] == 1:
                m = 2
                tmp = [
                    torch.pow(
                        rot_predict[o]
                        - rot_target[o]
                        @ rot((k * 2.0 / m) * torch.pi, rot_predict.device),
                        2,
                    ).mean()
                    for k in range(m)
                ]
                rot_loss = torch.min(torch.stack(tmp))
            elif sym_[o] == 2:
                m = 4
                tmp = [
                    torch.pow(
                        rot_predict[o]
                        - rot_target[o]
                        @ rot((k * 2.0 / m) * torch.pi, rot_predict.device),
                        2,
                    ).mean()
                    for k in range(m)
                ]
                rot_loss = torch.min(torch.stack(tmp))
            elif sym_[o] == 3:
                m = 36
                tmp = [
                    torch.pow(
                        rot_predict[o]
                        - rot_target[o]
                        @ rot((k * 2.0 / m) * torch.pi, rot_predict.device),
                        2,
                    ).mean()
                    for k in range(m)
                ]
                rot_loss = torch.min(torch.stack(tmp))
            else:
                rot_loss = torch.pow(
                    rot_predict[o] - rot_target[o],
                    2,
                ).mean()
            rot_loss_list.append(rot_loss)
        rot_loss = torch.mean(torch.stack(rot_loss_list))
        return rot_loss

    def loss(
        self,
        out_dict_list,
        obbs_padded: Obb3D,
        T_world_local: Pose,
        sym: List = None,
        *argv,
    ):
        """
        input:
            out_dict_list:  predicted box3d parameters
            obbs_padded:    target box3d parameters
            T_world_local:  pose matrix from local to world
            sym:            symmetry of the object
        output:
            loss
        """
        assert obbs_padded.ndim == 3, f"{obbs_padded.shape}"

        loss_total = (
            out_dict_list[-1]["ortho6d"].sum()
            * out_dict_list[-1]["size_unnormalized"].sum()
            * out_dict_list[-1]["center_unnormalized"].sum()
            * out_dict_list[-1]["pred_logits"].sum()
            * 0
        )

        loss_dict = {"center_loss": 0, "size_loss": 0, "rot_loss": 0, "cat_loss": 0}
        valid_bs = 0

        bb3_target_list = self.parse_target(
                bb3_target=obbs_padded,
                T_world_local=T_world_local,
        )
                    
        for out_dict in out_dict_list:
            # parse target
            bs = obbs_padded.shape[0]
            indices, punish_mask = self.matcher(out_dict, bb3_target_list)
            
            # compute loss
            for i in range(bs):
                if len(indices[i][0]) != 0:
                    valid_bs += 1
                    # center loss: l1
                    center_predict = out_dict["center_unnormalized"][i][indices[i][0]]
                    center_target = bb3_target_list[i]["center"][indices[i][1]]
                    center_loss = (center_predict - center_target).abs().mean()
                    center_loss *= self.loss_weight[0]
                    loss_total += center_loss
                    loss_dict["center_loss"] += center_loss
                    # size loss: l1
                    size_predict = out_dict["size_unnormalized"][i][indices[i][0]]
                    size_target = bb3_target_list[i]["size"][indices[i][1]]
                    size_loss = (size_predict - size_target).abs().mean()
                    size_loss *= self.loss_weight[1]
                    loss_total += size_loss
                    loss_dict["size_loss"] += size_loss
                    # rotation loss: cosine similarity or rmat l2
                    # cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
                    rot_6d_predict = out_dict["ortho6d"][i][indices[i][0]]
                    rot_target = bb3_target_list[i]["T_rig_object"][:, :3, :3][
                        indices[i][1]
                    ]
                    rot_predict = compute_rotation_matrix_from_ortho6d(rot_6d_predict)
                    # loss rmat
                    if sym is not None:
                        sym_ = sym[i][indices[i][1]]
                        rot_loss = self.rotation_loss_with_sym(
                            rot_predict, rot_target, sym_
                        )
                    else:
                        rot_loss = torch.pow(rot_target - rot_predict, 2).mean()
                    rot_loss *= self.loss_weight[2]
                    loss_total += rot_loss
                    loss_dict["rot_loss"] += rot_loss
                    # category loss
                    classes_target_o = bb3_target_list[i]["labels"][indices[i][1]]
                    classes_target = torch.full(
                        out_dict["pred_logits"].shape[1:2],
                        self.num_semcls,
                        dtype=torch.int64,
                        device=out_dict["pred_logits"].device,
                    )
                    classes_target[indices[i][0]] = classes_target_o
                    if punish_mask is not None:
                        cross_entropy = torch.nn.CrossEntropyLoss(
                            self.class_weight.to(classes_target_o.device), reduce=False
                        )
                        cat_loss = cross_entropy(out_dict["pred_logits"][i], classes_target)
                    
                        cat_loss = (cat_loss * punish_mask[i]).sum() / punish_mask[i].sum()
                    else:
                        cross_entropy = torch.nn.CrossEntropyLoss(
                            self.class_weight.to(classes_target_o.device)
                        )
                        cat_loss = cross_entropy(out_dict["pred_logits"][i], classes_target)
                    cat_loss *= self.loss_weight[3]
                    loss_total += cat_loss
                    loss_dict["cat_loss"] += cat_loss
        if valid_bs != 0:
            loss_total = loss_total / valid_bs
            for key, value in loss_dict.items():
                loss_dict[key] = value / valid_bs
            
        loss_dict["total_loss"] = loss_total
        return loss_dict

    def parse_pred(self, out_dict):
        """
        reorganize the predicitions into OBB class. Also aplly some filters here, e.g. remove the predictions out the scope, nms
        """
        # only use the prediciton in the last iteration
        out_dict = out_dict[-1]
        size_predict = out_dict["size_unnormalized"]
        center_predict = out_dict["center_unnormalized"]
        logits = out_dict["sem_cls_prob"]
        scores, labels = torch.max(logits, -1)
        bs = logits.shape[0]
        rot_6d_predict = out_dict["ortho6d"]
        rot_matrix = compute_rotation_matrix_from_ortho6d(
            rot_6d_predict.contiguous().view(-1, 6)
        ).view(bs, -1, 3, 3)

        # t_matrix: T_local_object
        t_matrix = Pose.from_Rt(rot_matrix, center_predict)
        x_min, x_max = -size_predict[..., 0] / 2, size_predict[..., 0] / 2
        y_min, y_max = -size_predict[..., 1] / 2, size_predict[..., 1] / 2
        z_min, z_max = -size_predict[..., 2] / 2, size_predict[..., 2] / 2
        c3o = torch.stack(
            [x_min, x_max, y_min, y_max, z_min, z_max],
            dim=-1,
        )

        obbs_pred = Obb3D.separate_init(
            bb3_object=c3o.cpu(),
            T_world_object=t_matrix._data.cpu(),
            sem_id=labels.cpu(),
        )
        out_dict["obbs_pred"] = obbs_pred.cuda()

        # the model will preserve all outputs from one snippet for visualization without ant filtering, 
        # e.g. filter out far-away boxes or using nms
        if not self.for_vis:
            valid = (
                (center_predict[..., 0] > self.track_scale[0]).int()
                + (center_predict[..., 0] < self.track_scale[1]).int()
                + (center_predict[..., 2] > self.track_scale[4]).int()
                + (center_predict[..., 2] < self.track_scale[5]).int()
            )
            valid = valid == 4
        else:
            valid = torch.ones_like(center_predict[..., 0]).bool()
        if self.enable_nms:
            if not self.for_vis:
                pred_mask = nms(obbs_pred, scores, self.num_semcls, 0.1, "nms_3d_faster")
            else:
                pred_mask = nms(obbs_pred, scores, self.num_semcls, 0.2, "nms_3d_faster_samecls")
        out_dict["pred_mask"] = torch.tensor(pred_mask).to(valid.device) & valid

        return out_dict

    def update_metrics(
        self,
        out_dict,
        obbs_padded: Obb3D,
        T_world_local: Pose,
        scene_name: str = None,
    ):
        """
        prediction parse and accumulate for f1 evaluation
        input:
            out_dict: box parameters
            target: ["obbs_padded", "T_world_rig", "scene_name"]
        """
        assert obbs_padded.ndim == 3, f"{obbs_padded.shape}"

        # parse prediction
        out_dict = self.parse_pred(out_dict)

        # parse target
        bb3_target_list = self.parse_target(
            bb3_target=obbs_padded,
            T_world_local=T_world_local,
        )
        for key, value in out_dict.items():
            if value is not None:
                out_dict[key] = value.detach()
        # for hungarian matching
        obbs_pred = out_dict["obbs_pred"]
        out_dict["scene_name"] = scene_name
        out_dict["pred_corners_world"] = T_world_local.transform(
            obbs_pred.T_world_object.transform(obbs_pred.bb3corners_object)
        ).detach()
        for calculator in self.metrics_calculator:
            calculator.step(out_dict, bb3_target_list)

    def compute_metrics(self):
        metrics = {}
        for calculator in self.metrics_calculator:
            metrics.update(calculator.compute_metrics())
        return metrics

    def reset_metrics(self):
        for calculator in self.metrics_calculator:
            calculator.reset()

    @rank_zero_only
    def log_images(
        self,
        out_dict,
        obbs_padded: Obb3D,
        Ts_world_pseudoCam: Pose,
        Ts_world_local: Pose,
        T_camera_pseudoCam: Pose,
        rgb_img=None,
        calib_rgb=None,
        slaml_img=None,
        calib_slaml=None,
        slamr_img=None,
        calib_slamr=None,
    ):
        # only log batch 0
        imgs = [rgb_img[0]]
        calibs = [calib_rgb[0]]
        tags = ["rgb_img"]
        tags_gt = ["rgb_img_gt"]
        # scannet only has rgb
        if slaml_img is not None:
            imgs.append(slaml_img[0])
            calibs.append(calib_slaml[0])
            tags.append("slaml_img")
            tags_gt.append("slaml_img_gt")
        if slamr_img is not None:
            imgs.append(slamr_img[0])
            calibs.append(calib_slamr[0])
            tags.append("slamr_img")
            tags_gt.append("slamr_img_gt")

        out_dict = self.parse_pred(out_dict)
        obbs_pred = out_dict["obbs_pred"]
        pred_mask = out_dict["pred_mask"]
        T_local_object = obbs_pred.T_world_object
        T_world_object = Ts_world_local @ T_local_object
        T_pseudoCam_world = Ts_world_pseudoCam.inverse()
        imgs_pred = draw_detections(
            imgs,
            calibs,
            obbs_pred.bb3corners_object[0],
            T_world_object[0],
            T_pseudoCam_world[0],
            T_camera_pseudoCam[0],
            obbs_pred.sem_id.squeeze(-1)[0].int(),
            tags,
            self.num_semcls,
            mask=pred_mask[0],
        )

        if obbs_padded[0] != '':
            obbs_padded = obbs_padded.remove_padding()
            imgs_gt = draw_detections(
                imgs,
                calibs,
                obbs_padded[0].bb3corners_object,
                obbs_padded[0].T_world_object,
                T_pseudoCam_world[0],
                T_camera_pseudoCam[0],
                obbs_padded[0].sem_id.squeeze(-1).int(),
                tags_gt,
                self.num_semcls,
            )
        else:
            imgs_gt = {}

        return {**imgs_pred, **imgs_gt}

model/parq_decoder.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
import logging
import math
import os

import numpy as np
import torch
from PIL import Image
from pytorch_lightning import LightningModule
from model.resnet_fpn import ResnetFPN
from model.ray_positional_encoding import AddRayPE
from model.parq_decoder import PARQDecoder
import sys
import einops
sys.path.append("..") # Adds higher directory to python modules path.


from utils.train_utils import CosineAnnealingWarmupRestarts

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# logger.setLevel(logging.DEBUG)

from utils.vis_utils import normalize, pca_compress


class PARQ(LightningModule):
    """
    PARQ Lightning Module

    Args:
        cfg:
            configure the lightning module
    """

    def __init__(
        self,
        cfg,
    ):
        super().__init__()
        self.save_hyperparameters(cfg)
        # hparams == cfg; this supports loading from checkpoint
        self.cfg = self.hparams

        self.batch_size = self.cfg.DATAMODULE.BATCH_SIZE

        self.backbone2d = ResnetFPN(
            cfg.MODEL.BACKBONE2D.RESNET_NAME,
            cfg.MODEL.BACKBONE2D.LAYER,
            cfg.MODEL.BACKBONE2D.FREEZE)
        self.add_ray_pe = AddRayPE(
            cfg.MODEL.TOKENIZER.OUT_CHANNELS,
            cfg.MODEL.TOKENIZER.RAY_POINTS_SCALE,
            cfg.MODEL.TOKENIZER.NUM_SAMPLES,
            cfg.MODEL.TOKENIZER.MIN_DEPTH,
            cfg.MODEL.TOKENIZER.MAX_DEPTH)
        self.box3d_decoder = PARQDecoder(cfg.MODEL.DECODER)

    def prepare_data(self):
        rank = os.getenv("RANK")
        logger.info(f"prepare_data rank={rank}")

    def setup(self, stage=None):
        logger.info("setup done ")

    def forward(self, batch, batch_idx):
        # image feature extraction
        batch = self.backbone2d(batch)
        # generate ray positional encoding
        encoding = self.add_ray_pe(batch['all_features'], batch['camera_feature'], batch['T_camera_pseudoCam'], batch['T_world_pseudoCam'], batch['T_world_local'])

        # add ray positional encoding to image features
        images_feat = batch['all_features'] + encoding

        # tonkenize image features
        input_tokens = einops.rearrange(
            images_feat,
            "b (t dt) c (h dh) (w dw) -> b t h w (dt dh dw c)",
            dh=1,
            dw=1,
            dt=1,
        )
        input_tokens = einops.rearrange(input_tokens, "b t h w c -> b (t h w) c")

        # compute the outputs
        outputs = self.box3d_decoder(input_tokens, batch['camera_feature'], batch['T_camera_pseudoCam'], batch['T_world_pseudoCam'], batch['T_world_local'])
        
        # compute loss
        losses = self.box3d_decoder.loss(outputs, batch['obbs_padded'], batch['T_world_local'], batch['sym'])
        return losses, outputs
    
    def training_step(self, batch, batch_idx):
        losses, outputs = self.forward(batch, batch_idx)
        self.log_step(batch, losses, outputs, batch_idx=batch_idx, stage="train")
        return losses['total_loss']

    def validation_step(self, batch, batch_idx):
        losses, outputs = self.forward(batch, batch_idx)
        self.box3d_decoder.update_metrics(
                outputs, batch['obbs_padded'], batch['T_world_local'], batch['scene_name']
            )

        self.log_step(batch, losses, outputs, batch_idx=batch_idx, stage="val")
        return losses['total_loss']

    def on_validation_epoch_start(self) -> None:
        logger.debug("starting validation epoch; reset metrics")
        self.box3d_decoder.reset_metrics()

    def validation_epoch_end(self, outs):
        logger.debug("finishing validation epoch; evaluate metrics")
        metrics = {}
        metrics = self.box3d_decoder.compute_metrics()
        for key, value in metrics.items():
            if type(value) is np.ndarray:
                continue
            elif type(value) is Image.Image:
                if self.logger.experiment is not None:
                    self.logger.experiment.add_image(
                        "val/metrics/{}".format(key),
                        np.asarray(value).transpose(2, 0, 1),
                        self.trainer.current_epoch,
                    )
            else:
                self.log(
                    "val/metrics/{}".format(key),
                    value,
                    on_epoch=True,
                    logger=True,
                    sync_dist=True,
                    rank_zero_only=False,
                )
        logger.info(metrics)
        return metrics

    def test_step(self, batch, batch_idx=0):
        return self.forward(batch)

    def on_before_zero_grad(self, optimizer: torch.optim.Optimizer):
        current_lr = [d["lr"] for d in optimizer.param_groups][0]
        logger.debug(f"Step: {self.global_step} LR: {current_lr:.4e}")

    def configure_optimizers(self):
        logger.info(f"configure optimizer {self.cfg.OPTIMIZER.NAME}")
        effective_batch_size = (
            self.cfg.DATAMODULE.BATCH_SIZE
            * self.cfg.TRAINER.NUM_NODES
            * self.cfg.TRAINER.GPUS
            * self.cfg.TRAINER.ACCUMULATE_GRAD_BATCHES
        )

        # "Accurate Large Minibatch SGD ..."
        if self.cfg.OPTIMIZER.AUTOSCALE_LR:
            lr = self.cfg.OPTIMIZER.LEARNING_RATE * effective_batch_size / 256.0
        else:
            lr = self.cfg.OPTIMIZER.LEARNING_RATE

        # use the multi_tensor implementation of AdamW to speed up optimizer.step()
        optimizer = torch.optim._multi_tensor.AdamW(
        self.get_model_params(),
            lr=lr,
        )
        warmup_epochs = self.cfg.OPTIMIZER.WARMUP_EPOCHS
        lr_min = self.cfg.OPTIMIZER.LEARNING_RATE
        if effective_batch_size <= 256:
            lr_min = lr_min / 256.0
        cycle_mult = self.cfg.OPTIMIZER.CYCLE_MULT
        num_restarts = self.cfg.OPTIMIZER.NUM_RESTARTS
        total_epochs = self.cfg.TRAINER.MAX_EPOCHS
        cycle_fractions = [pow(cycle_mult, i) for i in range(num_restarts)]
        epochs_cycle_0 = math.ceil(total_epochs / sum(cycle_fractions))
        logger.info(f"warm up epochs {warmup_epochs}")
        logger.info(f"cycle0 {epochs_cycle_0}")
        logger.info(f"total {total_epochs}")
        logger.info(f"cycle multiplier {cycle_mult}")
        logger.info(f"lr_min {lr_min} to lr_max {lr}")
        # cosine decay with warmup epochs
        scheduler = CosineAnnealingWarmupRestarts(
            optimizer,
            epochs_cycle_0,
            cycle_mult,
            lr,
            lr_min,
            warmup_epochs,
        )
        return (
            {
                "optimizer": optimizer,
                "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"},
            },
        )

    def get_model_params(self):
        return (
            filter(lambda p: p.requires_grad, self.parameters())
            if self.cfg.OPTIMIZER.IGNORE_FROZEN_PARAMS
            else self.parameters()
        )

    def log_image(self, img: torch.Tensor, tag: str, step: int = 0):
        if (
            self.cfg.LOG_IMAGES
            and self.logger is not None
            and self.logger.experiment is not None
        ):
            logger.debug(f"image to log {img.shape}")
            self.logger.experiment.add_image(tag, img, step)

    def log_step(self, batch, losses, outputs, batch_idx=0, stage="train"):
        if len(losses) > 0:
            for key, value in losses.items():
                self.log(
                    f"{stage}/{key}",
                    value,
                    on_step=True,
                    logger=True,
                    sync_dist=True,
                )

        if batch_idx == 0 or batch_idx % self.cfg.LOG_IMAGES_FREQUENCY == 0:
            logger.debug("logging images")
            log_ims = self.get_log_images(batch, outputs)
            for tag, im in log_ims.items():
                if isinstance(im, np.ndarray):
                    self.log_image(
                        im,
                        f"{stage}/{tag}",
                        self.global_step,
                    )
                else:
                    logger.warning(
                        f"not logging {stage}/{tag} because it unknown type {type(im)}"
                    )
    
    def log_input(
        self,
        images: torch.Tensor
    ):
        log_img = {}
        T = images.shape[1]
        input_snippet = einops.rearrange(images, "b t c h w -> (b t) c h w")
        input_snippet = einops.rearrange(
            input_snippet, "(b t) c h w -> b c (t h) w", t=T
        )
        log_img["input"] = input_snippet
        return log_img

    def get_log_images(self, batch, outputs):
        def convert_image(img):
            """convert logged images to format CxHxW (1<=C<=3) range 0 to 255 np.uin8"""
            if not isinstance(img, torch.Tensor):
                return img
            if img.dim() == 5:  # BxTxCxHxW
                img = img[0, 0, ...]
            elif img.dim() == 4:  # BxCxHxW
                img = img[0, ...]
            img = img.detach()
            if img.shape[0] > 3:  # C
                img = pca_compress(img.unsqueeze(0)).squeeze(0)
            img = normalize(img)
            # go from 0-1 to 0-255
            img = (img * 255).cpu().numpy()
            img = img.astype(np.uint8)
            if img.shape[0] == 1:
                img = img[0, :, :]
            if img.ndim == 2:
                img = np.stack([img, img, img])
            return img

        log_ims = {}
        # Log outputs.
        log_outputs = self.box3d_decoder.log_images(
            outputs,
            batch['obbs_padded'],
            batch['T_world_pseudoCam'],
            batch['T_world_local'],
            batch['T_camera_pseudoCam'],
            batch['rgb_img'],
            batch['camera'],
        )
        if log_outputs is not None:
            for tag, im in log_outputs.items():
                log_ims[f"{tag}"] = convert_image(im)

        # Log inputs.
        log_inputs = self.log_input(
            batch['rgb_img']
        )
        if log_inputs is not None:
            for tag, im in log_inputs.items():
                log_ims[f"{tag}"] = convert_image(im)

        return log_ims

model/parq_lightning.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
import logging

import einops
import torch
from utils.encoding_utils import (
    grid_2d,
    ray_points_snippet,
)

from torch.nn import Module

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# logger.setLevel(logging.DEBUG)
logger.setLevel(logging.WARNING)


EPS = 1e-5


def inverse_sigmoid(x, eps=1e-3):
    x = x.clamp(min=0, max=1)
    x1 = x.clamp(min=eps)
    x2 = (1 - x).clamp(min=eps)
    return torch.log(x1/x2)


class AddRayPE(Module):
    def __init__(
        self,
        dim_out: int,
        # [min_x, max_x, min_y, max_y, min_z, max_z]
        ray_points_scale: list = [-2, 2, -1.5, 0, 0.25, 4.25],
        num_samples: int = 64,  # number of ray points
        min_depth: float = 0.25,
        max_depth: float = 5.25,
    ):
        """
        Args:
            out_channels: Output channels required from model
            ray_points_scale: [min_x, max_x, min_y, max_y, min_z, max_z] used to normalize the points along each ray
            num_samples:  number of ray points
            min_depth:    minimum depth of the ray points
            max_depth:    maximum depth of the ray points
        """
        super().__init__()

        self.dim_out = dim_out

        self.ray_points_scale = ray_points_scale
        self.num_samples = num_samples
        self.min_depth = min_depth
        self.max_depth = max_depth
        self.encoder = torch.nn.Sequential(
                torch.nn.Linear(3 * self.num_samples, dim_out),
                torch.nn.ReLU(),
                torch.nn.Linear(dim_out, dim_out),
        )

    def forward(
        self,
        images_feat: torch.Tensor,
        camera: torch.Tensor = None,
        T_camera_pseudoCam: torch.Tensor = None,
        T_world_pseudoCam: torch.Tensor = None,
        T_world_local: torch.Tensor = None,
    ):
        '''
        input:
            images_feat:        (B, T, C, H, W), image features
            camera:             (B, T, 6), camera intrinsics: with, height, fx, fy, cx, cy
            T_camera_pseudoCam: (B, T, 12), pseudo camera to camera transformation
            T_world_pseudoCam:  (B, T, 12), pseudo camera to world transformation
            T_world_local:      (B, 12), local to world transformation
        output:
            images_feat:        (B, T*H*W, C), tokenized image features with ray position encoding, patch size = 1
        '''
        B, T = images_feat.shape[0], images_feat.shape[1]

        width, height = camera.size[0, 0]
        width, height = width.round().int().item(), height.round().int().item()
        pos_2d = grid_2d(width, height, output_range=[0.0, width, 0.0, height])
        pos_2d = pos_2d.to(T_camera_pseudoCam.device)
        min_depth = torch.tensor([self.min_depth], device=T_camera_pseudoCam.device)[0]
        max_depth = torch.tensor([self.max_depth], device=T_camera_pseudoCam.device)[0]

        T_local_world = T_world_local.inverse()
        points3d = ray_points_snippet(
            pos_2d,
            camera,
            T_camera_pseudoCam,
            T_world_pseudoCam,
            T_local_world,
            self.num_samples,
            min_depth,
            max_depth,
        )
        points3d = einops.rearrange(
            points3d,
            "b t h w (n c) -> (b t) h w n c",
            b=B,
            t=T,
            h=height,
            w=width,
            n=self.num_samples,
        )
        points3d[..., 0] = (points3d[..., 0] - self.ray_points_scale[0]) / (
            self.ray_points_scale[1] - self.ray_points_scale[0]
        )
        points3d[..., 1] = (points3d[..., 1] - self.ray_points_scale[2]) / (
            self.ray_points_scale[3] - self.ray_points_scale[2]
        )
        points3d[..., 2] = (points3d[..., 2] - self.ray_points_scale[4]) / (
            self.ray_points_scale[5] - self.ray_points_scale[4]
        )
        points3d = inverse_sigmoid(points3d)
        points3d = einops.rearrange(
            points3d,
            "(b t) h w n c -> (b t) h w (n c)",
            b=B,
            t=T,
            h=height,
            w=width,
            n=self.num_samples,
        )

        encoding = self.encoder(points3d.contiguous())  # B x C x H x W
        encoding = einops.rearrange(
            encoding,
            "(b t) h w c -> b t c h w",
            b=B,
            t=T,
            h=height,
            w=width,
        )

        logger.debug(f"ray grid encoding {encoding.shape}")
        return encoding

model/ray_positional_encoding.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
import logging
from typing import Literal

import einops

import torch
from torchvision import transforms
from torchvision.models.detection.backbone_utils import resnet_fpn_backbone

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# logger.setLevel(logging.DEBUG)


class ResnetFPN(torch.nn.Module):
    """
    Extract a specified ResNet backbone with FPN on top
    resnet_name: type of resnet, e.g., resnet18, resnet34, resnet50, resnet101, resnet152
    layers:  name of layers to extract
    """

    def __init__(
        self,
        resnet_name: Literal[
            "resnet18", "resnet34", "resnet50", "resnet101", "resnet152"
        ] = "resnet18",
        layer: str = "0",
        freeze: bool = True,
    ):
        super(ResnetFPN, self).__init__()
        assert resnet_name in [
            "resnet18",
            "resnet34",
            "resnet50",
            "resnet101",
            "resnet152",
        ]
        self.resnet_fpn = resnet_fpn_backbone(
            resnet_name, pretrained=True, trainable_layers=5
        )
        self.layer = str(layer)
        self.transform = transforms.Compose(
            [
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                )
            ]
        )

        self.freeze = freeze
        if self.freeze:
            self.feature_extractor.eval()
        logger.info(f"Resnet FPN Preprocessor {resnet_name}, layers: {layer}")

    def forward(self, batch):
        B = batch["rgb_img"].shape[0]

        batch_img = batch["rgb_img"]

        if batch["rgb_img"].dim() == 5:
            batch_img = einops.rearrange(batch_img, "b t c h w -> (b t) c h w")

        # normalized mean and std
        batch_img = self.transform(batch_img)

        if self.freeze:
            with torch.no_grad():
                features = self.resnet_fpn(batch_img)
        else:
            features = self.resnet_fpn(batch_img)

        feature_list = []
        for layer in range(4):
            feature = features[str(layer)]
            feature = torch.nn.functional.interpolate(
                feature, features[self.layer].shape[-2:], mode="bilinear"
            )
            feature_list.append(feature)
        v = torch.cat(feature_list, dim=1)
            
        if batch["rgb_img"].dim() == 5:
            v = v.view(B, batch["rgb_img"].shape[1], *(v.size()[-(v.dim() - 1) :]))

        batch["all_features"] = v

        logger.debug(f"{self.layer} {v.shape}")
        camera = batch["camera"]
        scale_factor = 1 / (2 ** (int(self.layer) + 2))
        batch['camera_feature'] = camera.scale(scale_factor)
        return batch

model/resnet_fpn.py
ymingxie/PARQ

"""
Copyright (c) 2023 Yiming Xie
This file is derived from [DETR](https://github.com/facebookresearch/detr/blob/main/models/transformer.py).
Modified for [PARQ] by Yiming Xie.

Original header:
Copyright 2020 - present, Facebook, Inc

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

DETR Transformer class.
Copy-paste from torch.nn.Transformer with modifications:
    * positional encodings are passed in MHattention
    * extra LN at the end of encoder is removed
    * decoder returns a stack of activations from all decoding layers
"""
import copy
from typing import Optional, List

import torch
import torch.nn.functional as F
from torch import nn, Tensor
from torch.nn.functional import grid_sample
import math


def inverse_sigmoid(x, eps=1e-3):
    x = x.clamp(min=0, max=1)
    x1 = x.clamp(min=eps)
    x2 = (1 - x).clamp(min=eps)
    return torch.log(x1/x2)


def pos2posemb3d(pos, num_pos_feats=128, temperature=10000):
    # https://github.com/megvii-research/PETR/blob/main/projects/mmdet3d_plugin/models/dense_heads/petr_head.py#L29
    scale = 2 * math.pi
    pos = pos * scale
    dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=pos.device)
    dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats)
    pos_x = pos[..., 0, None] / dim_t
    pos_y = pos[..., 1, None] / dim_t
    pos_z = pos[..., 2, None] / dim_t
    pos_x = torch.stack(
        (pos_x[..., 0::2].sin(), pos_x[..., 1::2].cos()), dim=-1
    ).flatten(-2)
    pos_y = torch.stack(
        (pos_y[..., 0::2].sin(), pos_y[..., 1::2].cos()), dim=-1
    ).flatten(-2)
    pos_z = torch.stack(
        (pos_z[..., 0::2].sin(), pos_z[..., 1::2].cos()), dim=-1
    ).flatten(-2)
    posemb = torch.cat((pos_y, pos_x, pos_z), dim=-1)
    return posemb

        
class Transformer(nn.Module):

    def __init__(self, dec_dim=512, queries_dim=512, dec_heads=8,
                 dec_layers=6, dec_ffn_dim=2048, dropout_rate=0.1,
                 activation="relu", normalize_before=False,
                 return_intermediate_dec=True, 
                 scale=None,
                 share_weights=False):
        super().__init__()
        assert (
            queries_dim == dec_dim
        ), f"queries dim {queries_dim} needs to be equal to input enc_dim {dec_dim} for transformer encoder"
        
        decoder_layer = TransformerDecoderLayer(dec_dim, dec_heads, dec_ffn_dim,
                                                    dropout_rate, activation, normalize_before)
        decoder_norm = nn.LayerNorm(dec_dim)
        self.decoder = TransformerDecoder(decoder_layer, dec_layers, dec_dim, scale, decoder_norm,
                                              return_intermediate=return_intermediate_dec, share_weights=share_weights)
        self._reset_parameters()

        self.d_model = dec_dim
        self.nhead = dec_heads

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, input_tokens, query_tokens, meta_data, mask=None):
        """
        Decoder - PARQ module, including transformer decoder and mlp heads
        input:
            input_tokens:       tensor, (B, N, C), the output of tokenization of (image features + ray positional encoding)
            query_tokens:       tensor, (N', C), reference points, N' is the number of reference points
            meta_data:          dict, meta data of the scene, including:
                                    camera:             camera intrinsics, (B, 6), width, height, fx, fy, cx, cy.
                                    T_camera_pseudoCam: pose, (B, 12), pose matrix from pseudo camera to camera
                                    T_world_pseudoCam:  pose, (B, 12), pose matrix from pseudo camera to world
                                    T_world_local:      pose, (B, 12), pose matrix from local to world
        output:
            output_list (parameters of 3d boxes): 
            [
                iteration 1 (dict):
                    pred_logits:          (B, N, num_semcls + 1), the classification logits (including background class)
                    center_unnormalized:  (B, N, 3), the predicted center of the box
                    size_unnormalized:    (B, N, 3), the predicted size of the box
                    ortho6d:              (B, N, 6), the predicted rotation of the box
                    sem_cls_prob:         (B, N, num_semcls + 1), the softmax of pred_logits
                    coord_pos:            (B, N, 3), the position of reference points, used for matcher
                iteration 2 (dict):
                    ...
                ...
            ]
        """
        bs = input_tokens.shape[0]
        query_tokens = query_tokens.repeat(bs, 1, 1)

        output_list = self.decoder(None, input_tokens, memory_key_padding_mask=mask,
                          query_pos=query_tokens, meta_data=meta_data)
        return output_list


def project(memory_hw, query_pos, T_camera_local, camera):
    '''
    Project reference points onto multi-view image to fetch appearence features
    Bilinear interpolation is used to fetch features.
    Average pooling is used to aggregate features from different views.
    '''
    w, h = camera.size[0][0].data.cpu().numpy()
    # from local coord to camera coord
    query_pos_c = T_camera_local.transform(query_pos.unsqueeze(1))

    center_b_list = []
    valid_b_list = []
    for cam_b, pos_c_b in zip(camera, query_pos_c):
        center_im_b, center_valid_b = cam_b.project(pos_c_b)
        center_b_list.append(center_im_b)
        valid_b_list.append(center_valid_b)
    center_im = torch.stack(center_b_list)
    center_valid = torch.stack(valid_b_list)
        
    im_grid = torch.stack([2 * center_im[..., 0] / (w - 1) - 1, 2 * center_im[..., 1] / (h - 1) - 1], dim=-1)
    bs, num_view, num_query, _ = im_grid.shape
    im_grid = im_grid.view(bs * num_view, 1, -1, 2)
        
    features = grid_sample(memory_hw, im_grid, padding_mode='zeros', align_corners=True)
    features = features.view(bs, num_view, -1, num_query)
    features = features.permute(0, 1, 3, 2).contiguous()
    # average across different views
    features = features.sum(dim=1)
    mask = center_valid.sum(dim=1)
    invalid_mask = mask == 0
    mask[invalid_mask] = 1
    features /= mask.unsqueeze(-1)
    return features, center_im, center_valid


class TransformerDecoder(nn.Module):

    def __init__(self, decoder_layer, num_layers, dim_in, scale, norm=None, return_intermediate=False, share_weights=False):
        super().__init__()
        if not share_weights:
            self.layers = _get_clones(decoder_layer, num_layers)
        else:
            self.layers = _get_clones(decoder_layer, 1)
        self.num_layers = num_layers
        self.share_weights = share_weights
        self.norm = norm
        self.return_intermediate = return_intermediate
        self.position_encoder = torch.nn.Sequential(
                torch.nn.Linear(128 * 3, dim_in),
                torch.nn.ReLU(),
                torch.nn.Linear(dim_in, dim_in),
            )
        self.mlp_heads = None
        self.box_processor = None
        self.scale = scale

    def normalize(self, center_offset):
        center_offset1 = (center_offset[..., 0] - self.scale[0]) / (
                    self.scale[1] - self.scale[0]
                )
        center_offset2 = (center_offset[..., 1] - self.scale[2]) / (
                    self.scale[3] - self.scale[2]
                )
        center_offset3 = (center_offset[..., 2] - self.scale[4]) / (
                    self.scale[5] - self.scale[4]
                )
        center_offset = torch.stack([center_offset1, center_offset2, center_offset3], dim=-1)
        return center_offset
    
    def denormalize(self, center_offset):
        center_offset1 = (
                center_offset[..., 0] * (self.scale[1] - self.scale[0]) + self.scale[0]
            )
        center_offset2 = (
                center_offset[..., 1] * (self.scale[3] - self.scale[2]) + self.scale[2]
            )
        center_offset3 = (
                center_offset[..., 2] * (self.scale[5] - self.scale[4]) + self.scale[4]
            )
        center_offset = torch.stack([center_offset1, center_offset2, center_offset3], dim=-1)
        return center_offset
    
    def bbox3d_prediction(self, tokens, coord, layer_num):
        """
        Predict the paramers of the boxes via multiple MLP heads
        input:
            tokens: tensor, (B, N, C), output of transformer decoder
            coord: tensor, (B, N, 3), coordinates of the reference points
            layer_num: int, the layer number of the transformer decoder
        output:
            out_dict: parameters of 3d boxes
        """
        if tokens.dim() == 4:
            tokens_list = torch.split(tokens, 1, dim=0)
        else:
            tokens_list = [tokens.unsqueeze(0)]
        
        share_mlp_heads = True
        if isinstance(self.mlp_heads["sem_cls_head"], torch.nn.ModuleList):
            share_mlp_heads = False
            
        box_prediction_list = []
        for tokens in tokens_list:
            tokens = tokens[0]
            # tokens are B x nqueries x noutput, change to B x noutput x bqueries
            tokens = tokens.permute(0, 2, 1).contiguous()
            if share_mlp_heads:
                cls_logits = self.mlp_heads["sem_cls_head"](tokens).transpose(1, 2)
                center_offset = self.mlp_heads["center_head"](tokens).transpose(1, 2)
            else:
                cls_logits = self.mlp_heads["sem_cls_head"][layer_num](tokens).transpose(1, 2)
                center_offset = self.mlp_heads["center_head"][layer_num](tokens).transpose(1, 2)

            coord_pos = self.denormalize(coord)
            center_offset = center_offset + inverse_sigmoid(coord)
            center_offset = center_offset.sigmoid()
            center_offset = self.denormalize(center_offset)
            
            if share_mlp_heads:
                size_scale = self.mlp_heads["size_head"](tokens).transpose(1, 2)
                ortho6d = self.mlp_heads["rotation_head"](tokens).transpose(1, 2)
            else:
                size_scale = self.mlp_heads["size_head"][layer_num](tokens).transpose(1, 2)
                ortho6d = self.mlp_heads["rotation_head"][layer_num](tokens).transpose(1, 2)

            (
                center_normalized,
                center_unnormalized,
            ) = self.box_processor.compute_predicted_center(center_offset)

            # below are not used in computing loss (only for matching/mAP eval)
            # we compute them with no_grad() so that distributed training does not complain about unused variables
            with torch.no_grad():
                (
                    semcls_prob,
                    objectness_prob,
                ) = self.box_processor.compute_objectness_and_cls_prob(cls_logits)

            size_unnormalized = self.box_processor.compute_predicted_size(
                size_scale, semcls_prob
            )

            box_prediction = {
                "pred_logits": cls_logits,
                "center_unnormalized": center_unnormalized,
                "size_unnormalized": size_unnormalized,
                "ortho6d": ortho6d,
                "sem_cls_prob": semcls_prob,
                # used in matcher
                "coord_pos": coord_pos, # use input reference point to match instead of output center
            }
            box_prediction_list.append(box_prediction)
        return box_prediction_list

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None,
                meta_data = None):
        # parse meta data
        camera = meta_data["camera"]
        T_camera_pseudoCam = meta_data["T_camera_pseudoCam"]
        T_world_pseudoCam = meta_data["T_world_pseudoCam"]
        T_world_local = meta_data["T_world_local"]
        bs, num_view = T_camera_pseudoCam.shape[:2]

        T_camera_local = T_camera_pseudoCam @ (
                T_world_pseudoCam.inverse() @ T_world_local
            )
        w, h = camera.size[0][0].data.cpu().numpy()
        memory_hw = memory.view(bs * num_view, int(h), int(w), -1)
        memory_hw = memory_hw.permute(0, 3, 1, 2)
        
        memory = memory.permute(1, 0, 2)
        
        
        intermediate = []
        reference_points = query_pos.sigmoid()
        for layer_num in range(self.num_layers):
            if self.share_weights:
                layer = self.layers[0]
            else:
                layer = self.layers[layer_num]
            
            # positional encoding
            pos_feat = self.position_encoder(pos2posemb3d(reference_points))
            pos_feat = pos_feat.permute(1, 0, 2)

            # project
            pixel_aligned, center_im, center_valid = project(memory_hw, self.denormalize(reference_points), T_camera_local, camera)
            output = pixel_aligned.permute(1, 0, 2)

            output, attn = layer(output, memory, tgt_mask=tgt_mask,
                           memory_mask=memory_mask,
                           tgt_key_padding_mask=tgt_key_padding_mask,
                           memory_key_padding_mask=memory_key_padding_mask,
                           pos=pos, query_pos=pos_feat)
            output = output.permute(1, 0, 2)
            output_dict = self.bbox3d_prediction(output, reference_points, layer_num)
            reference_points = self.normalize(output_dict[0]["center_unnormalized"])
            reference_points = reference_points.detach()
            
            if self.return_intermediate:
                intermediate += output_dict

        return intermediate


class TransformerDecoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt, memory,
                     tgt_mask: Optional[Tensor] = None,
                     memory_mask: Optional[Tensor] = None,
                     tgt_key_padding_mask: Optional[Tensor] = None,
                     memory_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None,
                     query_pos: Optional[Tensor] = None):
        q = k = self.with_pos_embed(tgt, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2, attn = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt, attn

    def forward_pre(self, tgt, memory,
                    tgt_mask: Optional[Tensor] = None,
                    memory_mask: Optional[Tensor] = None,
                    tgt_key_padding_mask: Optional[Tensor] = None,
                    memory_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None,
                    query_pos: Optional[Tensor] = None):
        tgt2 = self.norm1(tgt)
        q = k = self.with_pos_embed(tgt2, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt2 = self.norm2(tgt)
        tgt2, attn = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)
        tgt = tgt + self.dropout2(tgt2)
        tgt2 = self.norm3(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout3(tgt2)
        return tgt, attn

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)


def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")

model/transformer_parq.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
import argparse
import glob
import os
import pickle

import numpy as np
import torch
from PIL import Image
import ray
import torch.multiprocessing

from processing_utils import (
    collate_fn,
    get_box3d_inside_fov,
    get_catid_to_label_name,
    get_homogeneous,
    get_label,
    get_level,
    get_point_cloud,
    get_point_cloud_inside_box3d,
    get_scannet_class_to_index_RayTran,
    name2ids,
    ScanNetDataset,
    split_list,
    view_selection,
    view_selection_overlap,
    view_selection_allframes,
    view_selection_w1,
)

torch.multiprocessing.set_sharing_strategy('file_system')

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--scannet_path",
        help="the data path of the scannet dataset",
        default="/work/vig/Datasets/ScanNet/scans_uncomp",
    )
    parser.add_argument(
        "--out_filename",
        help="the data path of the output file",
        default="/work/vig/Datasets/ScanNet/deper_scannet_anno/scan2cad_box3d_anno_view3",
    )
    parser.add_argument("--window_size", default=3, type=int)
    parser.add_argument("--min_angle", default=15, type=float)
    parser.add_argument("--min_distance", default=0.1, type=float)
    parser.add_argument("--all_frames", action="store_true", help="if load all frames (keyframes)")
    
    # ray multi processes
    parser.add_argument('--n_proc', type=int, default=8, help='#processes launched to process scenes.')
    parser.add_argument('--n_gpu', type=int, default=1, help='#number of gpus')
    parser.add_argument('--num_workers', type=int, default=8)
    parser.add_argument('--loader_num_workers', type=int, default=8)

    return parser.parse_args()


args = parse_args()


@ray.remote(num_cpus=args.num_workers + 1, num_gpus=(1 / args.n_proc))
def worker_thread(args, scenes):
    catid2name = get_catid_to_label_name(args.scannet_path)
    # read scannet validation scenes
    split_path = os.path.join(args.scannet_path, "..", "ScanNet", "Tasks", "Benchmark", "scannetv2_val.txt")
    scenes_val = open(split_path, "r").readlines()
    scenes_val = [l.strip() for l in scenes_val]

    for scene in scenes:
        if os.path.exists(
            os.path.join(args.out_filename, "image_anno_{}.pkl".format(scene))
        ):
            print(
                "{} exists, ignored".format(
                    os.path.join(args.out_filename, "image_anno_{}.pkl".format(scene))
                )
            )
            continue

        val = scene in scenes_val
        # if not val:
            # continue

        n_imgs = len(
            os.listdir(os.path.join(args.scannet_path, scene, "color"))
        )
        intrinsic_dir_depth = os.path.join(
            args.scannet_path,
            scene,
            "intrinsic",
            "intrinsic_depth.txt",
        )
        intrinsic_dir_color = os.path.join(
            args.scannet_path,
            scene,
            "intrinsic",
            "intrinsic_color.txt",
        )
        
        cam_intr_depth = np.loadtxt(intrinsic_dir_depth, delimiter=" ").astype(np.float32)

        cam_intr_color = np.loadtxt(intrinsic_dir_color, delimiter=" ").astype(np.float32)

        dataset = ScanNetDataset(n_imgs, scene, args.scannet_path)

        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=None,
            collate_fn=collate_fn,
            batch_sampler=None,
            num_workers=args.loader_num_workers,
        )

        cam_pose_all = {}
        for id, (cam_pose, _, _) in enumerate(dataloader):
            if id % 100 == 0:
                print("{}: read frame {}/{}".format(scene, str(id), str(n_imgs)))
            if (
                cam_pose[0][0] == np.inf
                or cam_pose[0][0] == -np.inf
                or cam_pose[0][0] == np.nan
            ):
                continue
            cam_pose_all.update({id: cam_pose})

        save_snippet_pkl(
            args,
            scene,
            cam_pose_all,
            cam_intr_depth,
            cam_intr_color,
            catid2name,
            val,
        )


def save_snippet_pkl(
    args, scene, cam_pose_list, cam_intr_depth, cam_intr_color, g_cad2scannet, val
):

    if scene.endswith("/"):
        scene = scene[:-1]

    # view selection
    print("segment: process scene {}".format(scene))
    if args.all_frames:
        all_ids = view_selection_allframes(args, scene, cam_pose_list)
    else:
        if val:
            if args.window_size == 1:
                all_ids = view_selection_w1(args, scene, cam_pose_list)
            else:
                all_ids = view_selection(args, scene, cam_pose_list)
        else:
            all_ids = view_selection_overlap(args, scene, cam_pose_list)

    try:
        with open(os.path.join(args.out_filename, scene + ".pkl"), "rb") as f:
            scene_anno_dict = pickle.load(f)
    except:
        # handle no oriented box for this scene
        return

    aligned_models = scene_anno_dict["aligned_models"]

    # camera coordinate
    catids = [_["catid_cad"] for _ in aligned_models]
    # catids -> cat name
    labels = get_label(g_cad2scannet, catids)
    # 20: others

    bboxes = [_["bboxes"] for _ in aligned_models]
    bbox_corners = [_["bbox_corners"] for _ in aligned_models]
    T_scan_object = [_["T_so"] for _ in aligned_models]
    sym = [_["sym"] for _ in aligned_models]
    bbox_corners = torch.from_numpy(np.array(bbox_corners)).cuda()

    roidb_scene = {
        "scene_name": scene,
        "bboxes": bboxes,
        "sym": sym,
        "T_scan_object": T_scan_object,
        "labels": labels,
        "snippets": [],
    }
    # save snippets
    for i, ids in enumerate(all_ids):
        img_shape = None

        inside_point_cloud_number_list = []
        trunc_ratio_list = []
        T_scan_camera_list = []
        cam_intr_color_list = []
        for id_ in ids:
            print("processing image {} scene {}".format(id_, scene))
            if img_shape is None:
                im_file = os.path.join(
                    args.scannet_path,
                    scene,
                    "color",
                    'frame-{0:06d}.color.jpg'.format(id_),
                )
                img = Image.open(im_file)
                img_shape = np.array(img).shape

            depth_file = os.path.join(
                args.scannet_path,
                scene,
                "depth",
                'frame-{0:06d}.depth.pgm'.format(id_),
            )
            T_scan_camera = cam_pose_list[id_]
            T_scan_camera_list.append(T_scan_camera)
            cam_intr_color_list.append(np.copy(cam_intr_color))

            if not args.all_frames:
                T_camera_scan = np.linalg.inv(T_scan_camera)

                T_camera_scan = torch.from_numpy(T_camera_scan).cuda()

                bboxes_camera = (
                    get_homogeneous(bbox_corners.view(-1, 3)) @ T_camera_scan.T
                )[:, :3].view(-1, 8, 3)

                # camera coordinate
                point_cloud = get_point_cloud(depth_file, cam_intr_depth)

                inside_point_cloud_number, cs_list = get_point_cloud_inside_box3d(
                    bboxes_camera, point_cloud
                )
                inside_point_cloud_number_list.append(inside_point_cloud_number)

                trunc_ratio, box2d_list = get_box3d_inside_fov(
                    bboxes_camera, img_shape, cam_intr_color
                )
                trunc_ratio_list.append(trunc_ratio)

        if args.all_frames:
            inside_point_cloud_number = trunc_ratio = None
        else:
            inside_point_cloud_number = torch.stack(inside_point_cloud_number_list).max(
                dim=0
            )[0]
            trunc_ratio = torch.stack(trunc_ratio_list).max(dim=0)[0]

        roidb_scene["snippets"].append(
            {
                "snippet_id": i,
                "image_ids": ids,
                "intrinsic": cam_intr_color_list,
                "T_scan_camera": T_scan_camera_list,
                "point_cloud_num_list": inside_point_cloud_number,
                "truncation_ratio_list": trunc_ratio,
            }
        )

    with open(
        os.path.join(args.out_filename, "image_anno_{}.pkl".format(scene)), "wb"
    ) as f:
        pickle.dump(roidb_scene, f)
    return


def get_roidb(args, split="train"):
    split_file = "scannetv2_train.txt" if split == "train" else "scannetv2_val.txt"
    print(split_file)
    split_path = os.path.join(args.scannet_path, "..", "ScanNet", "Tasks", "Benchmark", split_file)
    scenes = open(split_path, "r").readlines()
    scenes = [l.strip() for l in scenes]

    scannet_class_to_index = get_scannet_class_to_index_RayTran()

    roidb_filename = (
        os.path.join(args.out_filename, "scannet_train_gt_roidb.pkl")
        if "train" in split
        else os.path.join(args.out_filename, "scannet_val_gt_roidb.pkl")
    )
    scene_anno_path = os.path.join(args.out_filename, "scene_anno")
    if not os.path.exists(scene_anno_path):
        os.makedirs(scene_anno_path)
    # if os.path.exists(roidb_filename):
    #     item_list = pickle.load(open(roidb_filename, "rb"))
    #     print(
    #         "roidb loaded from {}, totally {} samples".format(
    #             roidb_filename, len(item_list)
    #         )
    #     )
    #     return item_list

    item_list = []
    roidb_filenames = glob.glob(os.path.join(args.out_filename, "image_anno*"))
    roidb_filenames.sort()
    print("total files: {}".format(len(roidb_filenames)))
    for roidb_file in roidb_filenames:
        with open(roidb_file, "rb") as f:
            roidb_scene = pickle.load(f)
        scene_name = roidb_scene["scene_name"]
        snippets = roidb_scene["snippets"]
        labels = roidb_scene["labels"]
        bboxes = roidb_scene["bboxes"]
        T_scan_object = roidb_scene["T_scan_object"]
        sym = roidb_scene["sym"]

        # cat name -> defined id (0, 1, 2, 3 ...)
        ids = name2ids(labels, scannet_class_to_index)
        print(labels)
        if scene_name not in scenes:
            continue
        print(scene_name)
        item_one_scene = {}
        for im_anno in snippets:
            point_cloud_num_list = im_anno["point_cloud_num_list"]
            trunc_ratio_list = im_anno["truncation_ratio_list"]

            num_obj = len(bboxes)
            valid_obj = []
            valid_difficulty = []
            for i in range(num_obj):
                if point_cloud_num_list is not None:
                    difficulty = get_level(
                        point_cloud_num_list[i],
                        trunc_ratio_list[i],
                    )
                    valid_difficulty.append(difficulty)
                    if difficulty >= 3:
                        continue
                valid_obj.append(i)

            if len(valid_obj) == 0:
                continue

            item_list.append(
                {
                    "scene_name": roidb_scene["scene_name"],
                    "snippet_id": im_anno["snippet_id"],
                }
            )

            item_one_scene[im_anno["snippet_id"]] = {
                "image_ids": im_anno["image_ids"],
                "T_scan_camera": im_anno["T_scan_camera"],
                "intrinsic": im_anno["intrinsic"],
                "annotations": {
                    "label": [ids[i] for i in valid_obj],
                    "bboxes": [bboxes[i] for i in valid_obj],
                    "sym": [sym[i] for i in valid_obj],
                    "T_scan_object": [T_scan_object[i] for i in valid_obj],
                    # "pc_num": [point_cloud_num_list[i] for i in valid_obj],
                    # "truncation": [trunc_ratio_list[i] for i in valid_obj],
                    # "difficulty": [valid_difficulty[i] for i in valid_obj],
                },
            }
        scene_save_path = os.path.join(scene_anno_path, scene_name + ".pkl")
        with open(scene_save_path, "wb") as fout:
            pickle.dump(item_one_scene, fout)

    with open(roidb_filename, "wb") as fout:
        pickle.dump(item_list, fout)
    print(
        "annotations written to file {}, totally {} samples".format(
            roidb_filename, len(item_list)
        )
    )
    return item_list


def main():
    all_proc = args.n_proc * args.n_gpu
    ray.init(num_cpus=all_proc * (args.num_workers + 1), num_gpus=args.n_gpu)
    
    # get scene name
    scenes = os.listdir(args.scannet_path)
    scenes.sort()

    # scenes = scenes[:1]
    # scenes = ["scene0599_02"] + scenes
    # worker_thread(args, scenes)
    print("totally {} scenes".format(len(scenes)))

    scenes = split_list(scenes, all_proc)

    ray_worker_ids = []
    for w_idx in range(all_proc):
        ray_worker_ids.append(worker_thread.remote(args, scenes[w_idx]))
        
    results = ray.get(ray_worker_ids)
    
    print("generate train data")
    get_roidb(args, split="train")

    print("generate val data")
    get_roidb(args, split="val")


if __name__ == "__main__":
    main()

scripts/scannet_preprocessing/generate_scannet_anno_snippet.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
import argparse
import json
import os
import pickle

import numpy as np

from processing_utils import get_corner_by_dims, get_homogeneous, make_M_from_tqs


def generate_anno(args, full_anno):
    full_anno_list = []
    anno_num = len(full_anno)

    for j, anno in enumerate(full_anno):
        id_scan = anno["id_scan"]
        # if os.path.exists(os.path.join(args.out_filename, id_scan + ".pkl")):
        #     continue

        n_aligned_models = anno["n_aligned_models"]
        aligned_models = anno["aligned_models"]
        anno_dict = dict(
            id_scan=id_scan,
            n_aligned_models=n_aligned_models,
            aligned_models=[],
        )
        print(
            "processing bounding box {}, {} % {}, containing {} objs".format(
                id_scan, j, anno_num, n_aligned_models
            )
        )
        # T_world_scan
        T_ws = make_M_from_tqs(
            anno["trs"]["translation"],
            anno["trs"]["rotation"],
            anno["trs"]["scale"],
        )
        # T_scan_world
        T_sw = np.linalg.inv(T_ws)

        for i, model in enumerate(aligned_models):
            print("processing bounding box  {} % {}".format(i, n_aligned_models))
            t = model["trs"]["translation"]
            q = model["trs"]["rotation"]
            s = model["trs"]["scale"]

            sym = model["sym"]

            catid_cad = model["catid_cad"]
            id_cad = model["id_cad"]

            mat_off = np.eye(4)
            mat_off[:3, 3] = model["center"]

            if min(s) < 1e-3:
                continue
            scales = model["bbox"] * np.asarray(s) * 2
            T_wo = make_M_from_tqs(t, q, np.ones_like(s))
            T_so = T_sw @ T_wo @ mat_off
            bboxes = np.stack(
                [
                    -scales[0] / 2,
                    scales[0] / 2,
                    -scales[1] / 2,
                    scales[1] / 2,
                    -scales[2] / 2,
                    scales[2] / 2,
                ],
            )
            bbox_corners = get_corner_by_dims(scales)
            bbox_corners = (get_homogeneous(bbox_corners) @ T_so.T)[:, :3]

            model_box = dict(
                id_obj=i,
                catid_cad=catid_cad,
                id_cad=id_cad,
                bboxes=bboxes,
                bbox_corners=bbox_corners,
                T_so=T_so,
                sym=sym,
            )
            anno_dict["aligned_models"].append(model_box)

        with open(os.path.join(args.out_filename, id_scan + ".pkl"), "wb") as fout:
            pickle.dump(anno_dict, fout)
        print(
            "id_scan {} saved to {}".format(
                id_scan, os.path.join(args.out_filename, id_scan)
            )
        )
        full_anno_list.append(anno_dict)

    with open(
        os.path.join(args.out_filename, "scan2cad_bbox_3d_anno.pkl"), "wb"
    ) as fout:
        pickle.dump(full_anno_list, fout)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--scan2cad_path",
        help="the data path of the scan2cad dataset",
        default="/work/vig/Datasets/ScanNet/scan2cad/full_annotations.json",
    )
    parser.add_argument(
        "--scannet_path",
        help="the data path of the scannet dataset",
        default="/work/vig/Datasets/ScanNet/scans_uncomp",
    )
    parser.add_argument(
        "--out_filename",
        help="the data path of the output file",
        default="/work/vig/Datasets/ScanNet/deper_scannet_anno/scan2cad_box3d_anno_view3",
    )

    args = parser.parse_args()

    with open(args.scan2cad_path, "r") as f:
        full_anno = json.load(f)
    generate_anno(args, full_anno)

scripts/scannet_preprocessing/parse_scan2cad.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
import copy
import csv
import os
from io import BytesIO
from typing import Union

import cv2

import numpy as np
import torch

try:
    import quaternion  # @manual
except:
    import Quaternion as quaternion  # @manual


def make_M_from_tqs(t, q, s):
    q = np.quaternion(q[0], q[1], q[2], q[3])
    T = np.eye(4)
    T[0:3, 3] = t
    R = np.eye(4)
    R[0:3, 0:3] = quaternion.as_rotation_matrix(q)
    S = np.eye(4)
    S[0:3, 0:3] = np.diag(s)

    M = T.dot(R).dot(S)
    return M


def get_homogeneous(
    pts: Union["np.ndarray", "torch.tensor"]
):
    """convert [(b), N, 3] pts to homogeneous coordinate
    Args:
        pts ([(b), N, 3] Union['np.ndarray', 'torch.tensor']): input point cloud
    Returns:
        homo_pts ([(b), N, 4] Union['np.ndarray', 'torch.tensor']): output point
            cloud
    Raises:
        ValueError: if the input tensor/array is not with the shape of [b, N, 3]
            or [N, 3]
        TypeError: if input is not either tensor or array
    """

    batch = False
    if len(pts.shape) == 3:
        pts = pts[0]
        batch = True
    elif len(pts.shape) == 2:
        pts = pts
    else:
        raise ValueError("only accept [b, n_pts, 3] or [n_pts, 3]")

    if isinstance(pts, torch.Tensor):
        ones = torch.ones_like(pts[:, 2:])
        homo_pts = torch.cat([pts, ones], axis=1)
        if batch:
            return homo_pts[None, :, :]
        else:
            return homo_pts
    elif isinstance(pts, np.ndarray):
        ones = np.ones_like(pts[:, 2:])
        homo_pts = np.concatenate([pts, ones], axis=1)
        if batch:
            return homo_pts[None, :, :]
        else:
            return homo_pts
    else:
        raise TypeError("wrong data type")


def get_corner_by_dims(dimensions):
    """get 8 corner points of 3D bbox defined by self.dimensions
    Returns:
        a np.ndarray with shape [8,3] to represent 8 corner points'
        position of the 3D bounding box.
    """

    l, h, w = dimensions[0], dimensions[1], dimensions[2]
    x_corners = [-l / 2, l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2, -l / 2]
    y_corners = [-h / 2, -h / 2, h / 2, h / 2, -h / 2, -h / 2, h / 2, h / 2]
    z_corners = [-w / 2, -w / 2, -w / 2, -w / 2, w / 2, w / 2, w / 2, w / 2]
    corner_pts = np.array([x_corners, y_corners, z_corners], dtype=np.float32).T
    return corner_pts


def get_scannet_class_to_index():
    class_dict = {
        "chair": 0,
        "table": 1,
        "cabinet": 2,
        "trashbin": 3,
        "bookshelf": 4,
        "display": 5,
        "sofa": 6,
        "bathtub": 7,
        "bed": 8,
        "file cabinet": 9,
        "bag": 10,
        "printer": 11,
        "washer": 12,
        "lamp": 13,
        "microwave": 14,
        "stove": 15,
        "basket": 16,
        "bench": 17,
        "laptop": 18,
        "computer keyboard": 19,
        "other": 20,
    }
    return class_dict


def get_scannet_class_to_index_RayTran():
    # top 8
    class_dict = {
        "chair": 0,
        "table": 1,
        "cabinet": 2,
        "trashbin": 3,
        "bookshelf": 4,
        "display": 5,
        "sofa": 6,
        "bathtub": 7,
        "other": 8,
    }
    return class_dict


def get_point_cloud(depth_dir, intrinsics_depth):
    d = cv2.imread(depth_dir, -1).astype(np.float32)
    d /= 1000
    h, w = d.shape
    xx, yy = np.meshgrid(np.arange(w), np.arange(h))
    xx = xx * d
    yy = yy * d
    pc = np.stack([xx, yy, d, np.ones_like(xx)], axis=2)
    pc = pc.reshape(-1, 4)
    T_camera_screen = np.linalg.inv(intrinsics_depth)

    T_camera_screen = torch.from_numpy(T_camera_screen).cuda()
    pc = torch.from_numpy(pc).cuda().float()

    pc = (T_camera_screen @ pc.T).T

    pos_z = torch.nonzero(pc[:, 2] > 0).squeeze(1)

    # pc = np.dot(T_camera_screen, pc.T).T
    # pos_z = np.nonzero(pc[:, 2] > 0)[0]

    point_cloud = pc[pos_z]
    return point_cloud[:, :3]


def get_catid_to_label_name(csv_filename):
    CARE_CLASSES = {
        "03211117": "display",
        "04379243": "table",
        "02808440": "bathtub",
        "02747177": "trashbin",
        "04256520": "sofa",
        "03001627": "chair",
        "02933112": "cabinet",
        "02871439": "bookshelf",
        "00000000": "other",
    }

    # catid2index = {}
    # path = os.path.join(csv_filename, "..", "scannetv2-labels.combined.tsv")
    # with path_manager.open(path) as csvfile:
    #     spamreader = csv.DictReader(csvfile, delimiter="\t")
    #     for row in spamreader:
    #         try:
    #             catid2index[row["wnsynsetid"][1:]] = (
    #                 int(row["nyu40id"]),
    #                 row["nyu40class"],
    #             )
    #         except:
    #             pass
    return CARE_CLASSES


def get_label(g_cad2scannet, catid_cads):
    labels = []
    for catid in catid_cads:
        try:
            name = g_cad2scannet[catid]
            labels.append(name)
        except KeyError:
            labels.append("other")
    return labels


def name2ids(labels, class_to_index):
    ids = []
    for label in labels:
        try:
            ids.append(class_to_index[label])
        except KeyError:
            ids.append(class_to_index["other"])
    return ids


def get_box3d_inside_fov(box3d_list, image_size, intrinsics_color):
    h, w, c = image_size
    T_screen_camera = torch.from_numpy(intrinsics_color).cuda()
    ratio_list = []
    ret_2d_box = []
    for ind, box3d in enumerate(box3d_list):
        ret_2d_box.append([])
        hom_box3d = torch.cat([box3d, torch.ones((8, 1)).cuda()], axis=1).float()
        projected_box3d = (T_screen_camera @ hom_box3d.T).T[:, :3]
        projected_box3d[:, 0] /= torch.maximum(
            projected_box3d[:, 2], torch.Tensor([1]).cuda()
        )
        projected_box3d[:, 1] /= torch.maximum(
            projected_box3d[:, 2], torch.Tensor([1]).cuda()
        )
        xmin, xmax = torch.min(projected_box3d[:, 0]), torch.max(projected_box3d[:, 0])
        ymin, ymax = torch.min(projected_box3d[:, 1]), torch.max(projected_box3d[:, 1])
        area = (xmax - xmin) * (ymax - ymin)
        cliped_xmin, cliped_xmax = torch.clip(xmin, 0, w - 1), torch.clip(
            xmax, 0, w - 1
        )
        cliped_ymin, cliped_ymax = torch.clip(ymin, 0, h - 1), torch.clip(
            ymax, 0, h - 1
        )
        ret_2d_box[ind] = [cliped_xmin, cliped_ymin, cliped_xmax, cliped_ymax]
        area_inside = (cliped_xmax - cliped_xmin) * (cliped_ymax - cliped_ymin)
        ratio = 1.0 * area_inside / torch.maximum(area, torch.Tensor([1]).cuda())
        ratio_list.append(ratio)
    return torch.cat(ratio_list), ret_2d_box


def get_point_cloud_inside_box3d(bbox3, pc):
    """
    Find point cloud inside a bounding box.
    :param bbox3:
    :param pc:
    :return:
    """
    point = copy.deepcopy(pc).unsqueeze(0)
    v45 = bbox3[:, 5] - bbox3[:, 4]
    v40 = bbox3[:, 0] - bbox3[:, 4]
    v47 = bbox3[:, 7] - bbox3[:, 4]
    point = point - bbox3[:, 4:5]
    m0 = torch.matmul(point, v45.unsqueeze(-1))
    m1 = torch.matmul(point, v40.unsqueeze(-1))
    m2 = torch.matmul(point, v47.unsqueeze(-1))
    cs = []
    for m, v in zip([m0, m1, m2], [v45, v40, v47]):
        c0 = 0 < m
        c1 = m < v.unsqueeze(1) @ v.unsqueeze(-1)
        c = c0 & c1
        cs.append(c)

    cs = cs[0] & cs[1] & cs[2]
    cs = cs.squeeze(-1)

    inside_number = cs.sum(-1)
    return inside_number, cs


def collate_fn(list_data):
    cam_pose, depth_im, _ = list_data
    # Concatenate all lists
    return cam_pose, depth_im, _


class ScanNetDataset(torch.utils.data.Dataset):
    """Pytorch Dataset for a single scene. getitem loads individual frames"""

    def __init__(self, n_imgs, scene, data_path, max_depth=3.0, id_list=None):
        """
        Args:
        """
        self.n_imgs = n_imgs
        self.scene = scene
        self.data_path = data_path
        self.max_depth = max_depth
        if id_list is None:
            self.id_list = [i for i in range(n_imgs)]
        else:
            self.id_list = id_list

    def __len__(self):
        return self.n_imgs

    def __getitem__(self, id):
        """
        Returns:
            dict of meta data and images for a single frame
        """
        id = self.id_list[id]
        path = os.path.join(
            self.data_path, self.scene, "pose", 'frame-{0:06d}.pose.txt'.format(id)
        )
        cam_pose = np.loadtxt(path, delimiter=" ")
        return cam_pose, None, None


def get_level(pc, trunc):
    difficulty_level = [
        {
            "point_cloud_num": 1000,
            "truncation_ratio": 0.85,
        },
        {
            "point_cloud_num": 500,
            "truncation_ratio": 0.70,
        },
        {
            "point_cloud_num": 100,
            "truncation_ratio": 0.50,
        },
    ]

    if (
        pc > difficulty_level[0]["point_cloud_num"]
        and trunc > difficulty_level[0]["truncation_ratio"]
    ):
        return 0
    elif (
        pc > difficulty_level[1]["point_cloud_num"]
        and trunc > difficulty_level[1]["truncation_ratio"]
    ):
        return 1
    elif (
        pc > difficulty_level[2]["point_cloud_num"]
        and trunc > difficulty_level[2]["truncation_ratio"]
    ):
        return 2
    else:
        return 3


def decode_image(value: Union[torch.ByteTensor, BytesIO]):
    """
    Decode a image from ByteTensor loaded by koski
        It supports all opencv format (e.g., 16-bit depth image) and keeps
        original image format.
    """
    if isinstance(value, torch.ByteTensor):
        value = value.numpy().view()
    else:
        value = np.asarray(bytearray(value.read()), dtype=np.uint8)
    return cv2.imdecode(value, cv2.IMREAD_UNCHANGED)


def view_selection(args, scene, cam_pose_list):
    all_ids = []
    ids = []
    count = 0
    last_pose = None
    for id in cam_pose_list.keys():
        cam_pose = cam_pose_list[id]

        if count == 0:
            ids.append(id)
            last_pose = cam_pose
            count += 1
        else:
            angle = np.arccos(
                (
                    (
                        np.linalg.inv(cam_pose[:3, :3])
                        @ last_pose[:3, :3]
                        @ np.array([0, 0, 1]).T
                    )
                    * np.array([0, 0, 1])
                ).sum()
            )
            dis = np.linalg.norm(cam_pose[:3, 3] - last_pose[:3, 3])
            if angle > (args.min_angle / 180) * np.pi or dis > args.min_distance:
                ids.append(id)
                last_pose = cam_pose
                count += 1
                if count == args.window_size:
                    all_ids.append(ids)
                    ids = []
                    count = 0
    return all_ids


def view_selection_w1(args, scene, cam_pose_list):
    all_ids = []
    ids = []
    count = 0
    last_pose = None
    last_id = 0
    for id in cam_pose_list.keys():
        cam_pose = cam_pose_list[id]
        last_id = id
        if count == 0:
            ids.append(id)
            last_pose = cam_pose
            count += 1
        else:
            angle = np.arccos(
                (
                    (
                        np.linalg.inv(cam_pose[:3, :3])
                        @ last_pose[:3, :3]
                        @ np.array([0, 0, 1]).T
                    )
                    * np.array([0, 0, 1])
                ).sum()
            )
            dis = np.linalg.norm(cam_pose[:3, 3] - last_pose[:3, 3])
            if angle > (args.min_angle / 180) * np.pi or dis > args.min_distance:
                ids.append(id)
                last_pose = cam_pose
                count += 1
    
    for id in ids:
        all_ids.append([id])
    return all_ids


def view_selection_overlap(args, scene, cam_pose_list):
    all_ids = []
    ids = []
    count = 0
    last_pose = None
    last_id = 0
    for id in cam_pose_list.keys():
        cam_pose = cam_pose_list[id]
        last_id = id
        if count == 0:
            ids.append(id)
            last_pose = cam_pose
            count += 1
        else:
            angle = np.arccos(
                (
                    (
                        np.linalg.inv(cam_pose[:3, :3])
                        @ last_pose[:3, :3]
                        @ np.array([0, 0, 1]).T
                    )
                    * np.array([0, 0, 1])
                ).sum()
            )
            dis = np.linalg.norm(cam_pose[:3, 3] - last_pose[:3, 3])
            if angle > (args.min_angle / 180) * np.pi or dis > args.min_distance:
                ids.append(id)
                last_pose = cam_pose
                count += 1

    for i in range(10):
        for j in range(count):
            if j + args.window_size <= count:
                ids_snippet = ids[j : j + args.window_size]
                if ids_snippet[-1] + i <= last_id:
                    ids_snippet = [
                        IS + i for IS in ids_snippet if IS + i in cam_pose_list.keys()
                    ]
                    if len(ids_snippet) == args.window_size:
                        all_ids.append(ids_snippet)

    # remove repeated data
    ids_iden = []
    for ids in all_ids:
        if ids not in ids_iden:
            ids_iden.append(ids)

    return ids_iden


def view_selection_allframes(args, scene, cam_pose_list):
    all_ids = []
    ids = []
    count = 0
    last_pose = None
    for id in cam_pose_list.keys():
        cam_pose = cam_pose_list[id]

        if count == 0:
            ids.append(id)
            last_pose = cam_pose
            count += 1
        else:
            angle = np.arccos(
                (
                    (
                        np.linalg.inv(cam_pose[:3, :3])
                        @ last_pose[:3, :3]
                        @ np.array([0, 0, 1]).T
                    )
                    * np.array([0, 0, 1])
                ).sum()
            )
            dis = np.linalg.norm(cam_pose[:3, 3] - last_pose[:3, 3])
            if angle > (args.min_angle / 180) * np.pi or dis > args.min_distance:
                ids.append(id)
                last_pose = cam_pose
                count += 1
                # if count == args.window_size:
                #     all_ids.append(ids)
                #     ids = []
                #     count = 0
    all_ids.append(ids)
    return all_ids


def split_list(_list, n):
    assert len(_list) >= n
    ret = [[] for _ in range(n)]
    for idx, item in enumerate(_list):
        ret[idx % n].append(item)
    return ret

scripts/scannet_preprocessing/processing_utils.py
ymingxie/PARQ
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

import logging
import os

import time
from datetime import datetime
from os import path as osp
from typing import List
import argparse

from model import PARQ

import torch

from omegaconf import DictConfig, OmegaConf
from pytorch_lightning import seed_everything, Trainer
from pytorch_lightning.callbacks.base import Callback
from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
from pytorch_lightning.strategies.ddp import DDPStrategy
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor

from datasets.scannet_dataset import ScanNetDataModule
from config import cfg, update_config


def setup_tensorboard(cfg):
    tensorboard_run = f"parq_{datetime.fromtimestamp(time.time()).strftime('%y-%m-%d-%H-%M-%S')}"

    tb_logger = TensorBoardLogger(
        save_dir=cfg.LOG_PATH,
        name=cfg.NAME,
        version="model_" + tensorboard_run,
        default_hp_metric=False,
    )
    log_path = osp.join(cfg.LOG_PATH, cfg.NAME, "model_" + tensorboard_run)
    return log_path, tensorboard_run, tb_logger


def train(cfg):
    logging.basicConfig()
    logger = logging.getLogger(__name__)
    logger.info("training started")
    # pring config
    rank = int(os.getenv("RANK", 0))
    logger = logging.getLogger("train")
    logger.setLevel(logging.INFO)
    logger.info(f"starting train function {rank}")

    seed_everything(cfg.SEED)

    # create datamodule
    data_module = ScanNetDataModule(cfg.DATAMODULE)

    # create model from scratch
    logger.info("Setting up Perceiver model")
    model = PARQ(cfg)

    if cfg.PRETRAINED_PATH is not None:
        # just reload the weights, without other parameters (e.g. optimizer, scheduler, lr, etc.)
        assert cfg.PRETRAINED_PATH is None
        logger.info("Load pretrained model.")
        checkpoint = torch.load(cfg.PRETRAINED_PATH, map_location=model.device)[
            "state_dict"
        ]
        model.load_state_dict(checkpoint, strict=False)

    logger.info("Setting up tb logger")
    log_path, tb_run, tb_logger = setup_tensorboard(cfg)

    # Init lightning callbacks
    callbacks: List[Callback] = []
    callbacks.append(
        ModelCheckpoint(
            monitor=cfg.CALLBACK.MONITOR,
            save_top_k=cfg.CALLBACK.SAVE_TOP_K,
            save_last=cfg.CALLBACK.SAVE_LAST,
            verbose=cfg.CALLBACK.VERBOSE,
            auto_insert_metric_name=cfg.CALLBACK.AUTO_INSERT_METRIC_NAME,
            mode=cfg.CALLBACK.MODE,
            dirpath=log_path,
            filename=tb_run + "-epoch{epoch:04d}-val_loss{val/total_loss:.2f}-05_f1{val/metrics/0.5_f1:.3f}",
            )
    )
    callbacks.append(
        LearningRateMonitor(
            logging_interval= 'step'
            )
    )

    if cfg.CHECKPOINT_PATH is None:
        logger.info("no model checkpoint found - starting from scratch!")
    else:
        logger.info(f"loading model checkpoint from {cfg.CHECKPOINT_PATH}")

    effective_batch_size = (
        cfg.DATAMODULE.BATCH_SIZE
        * cfg.TRAINER.NUM_NODES
        * cfg.TRAINER.GPUS
        * cfg.TRAINER.ACCUMULATE_GRAD_BATCHES
    )

    strategy = DDPStrategy(find_unused_parameters=False, static_graph=True)
    trainer = Trainer(
        strategy=strategy,
        profiler=cfg.TRAINER.PROFILER,
        accelerator=cfg.TRAINER.ACCELERATOR,
        gpus=cfg.TRAINER.GPUS,
        num_nodes=cfg.TRAINER.NUM_NODES,
        accumulate_grad_batches=cfg.TRAINER.ACCUMULATE_GRAD_BATCHES,
        max_epochs=cfg.TRAINER.MAX_EPOCHS,
        log_every_n_steps=cfg.TRAINER.LOG_EVERY_N_STEPS,
        gradient_clip_val=cfg.TRAINER.GRADIENT_CLIP_VAL,
        reload_dataloaders_every_n_epochs=cfg.TRAINER.RELOAD_DATALOADERS_EVERY_N_EPOCHS,
        replace_sampler_ddp=cfg.TRAINER.REPLACE_SAMPLER_DDP,
        overfit_batches=cfg.TRAINER.OVERFIT_BATCHES,
        auto_scale_batch_size=cfg.TRAINER.AUTO_SCALE_BATCH_SIZE,
        check_val_every_n_epoch=cfg.TRAINER.CHECK_VAL_EVERY_N_EPOCH,
        precision=cfg.TRAINER.PRECISION,
        val_check_interval=cfg.TRAINER.VAL_CHECK_INTERVAL,
        resume_from_checkpoint=cfg.CHECKPOINT_PATH,
        limit_val_batches=8192 // effective_batch_size,
        callbacks=callbacks,
        logger=tb_logger,
    )

    # Training mode
    logger.info("starting training")
    trainer.fit(model=model, datamodule=data_module)
    logger.info(f"k best models: {trainer.checkpoint_callback.best_k_models}")
    logger.info(f"best model: {trainer.checkpoint_callback.best_model_path}")
    modelPath = trainer.checkpoint_callback.best_model_path
    if modelPath == "":
        logger.info(
            "starting testing the current model since no best model is saved yet."
        )
        retTest = trainer.test(model=model, datamodule=data_module)
        retVal = trainer.validate(model=model, datamodule=data_module)
    else:
        logger.info(f"starting testing model: {modelPath}")
        retTest = trainer.test(ckpt_path=modelPath, datamodule=data_module)
        retVal = trainer.validate(ckpt_path=modelPath, datamodule=data_module)

    logger.info("done training")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='A PyTorch Implementation of PARQ')

    # general
    parser.add_argument('--cfg',
                        help='experiment configure file name',
                        required=True,
                        type=str)

    parser.add_argument('opts',
                        help="Modify config options using the command-line",
                        default=None,
                        nargs=argparse.REMAINDER)

    # parse arguments and check
    args = parser.parse_args()
    update_config(cfg, args)

    print("training with pytorch lightning")
    train(cfg)

train.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
import logging

import einops

import torch

logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.setLevel(logging.WARNING)
# logger.setLevel(logging.DEBUG)


def grid_2d(width: int, height: int, output_range=(-1.0, 1.0, -1.0, 1.0)):
    x = torch.linspace(output_range[0], output_range[1], width + 1)[:-1]
    y = torch.linspace(output_range[2], output_range[3], height + 1)[:-1]
    xx, yy = torch.meshgrid(x, y, indexing="xy")
    grid = torch.stack([xx, yy], dim=-1)
    return grid


def ray_points_snippet(
    pixels: torch.Tensor,
    camera: torch.Tensor,
    T_camera_pseudoCam: torch.Tensor,
    T_world_pseudoCam: torch.Tensor,
    T_local_world: torch.Tensor,
    num_samples: int = 32,
    min_depth: float = 0.25,
    max_depth: float = 3.25,
):
    """
    ray positional encoding: sample points along the ray as positional encoding

    input:
        pixels: grid of rays [H x W x 2]
        T_camera_pseudoCam:   -
        T_world_pseudoCam:    -
        T_local_world:        -
        num_samples: number of points to sample along the ray
        min_depth: minimum depth of the ray
        max_depth: maximum depth of the ray

    output:
        points: [B x T x H x W x (num_samples x 3)] rays points in a consistent local coordinate frame
    """
    batch_size, time_snippet = T_camera_pseudoCam.shape[0], T_camera_pseudoCam.shape[1]

    # T_camera_pseudoCam = T_camera_pseudoCam.view(batch_size * time_snippet, 4, 4)
    # intrinsics = intrinsics.view(batch_size * time_snippet, 4, 4)
    T_camera_pseudoCam = T_camera_pseudoCam.view(batch_size * time_snippet, -1)
    camera = camera.view(batch_size * time_snippet, -1)

    points = ray_points(pixels, camera, T_camera_pseudoCam, num_samples, min_depth, max_depth)

    T_local_pseudoCam = T_local_world @ T_world_pseudoCam
    T_local_pseudoCam = T_local_pseudoCam.view(batch_size * time_snippet, -1)
    points = einops.rearrange(points, "b h w n c -> b (h w n) c")
    points = T_local_pseudoCam.transform(points)

    points = einops.rearrange(
        points,
        "(b t) (h w n) c -> b t h w (n c)",
        b=batch_size,
        t=time_snippet,
        h=pixels.shape[0],
        w=pixels.shape[1],
        n=num_samples,
    )
    return points


def ray_points(pixel_grid, camera, T_camera_pseudoCam, num_samples=32, min_depth=0.25, max_depth=3.25):
    grid_height, grid_width = pixel_grid.shape[0], pixel_grid.shape[1]
    batch_size = T_camera_pseudoCam.shape[0]
    pixel_grid = pixel_grid.reshape(-1, 2)
    pixel_grid = einops.repeat(pixel_grid, "n c -> b n c", b=batch_size)
    # unproject
    rays = camera.unproject(pixel_grid)

    linear_ramp = (
        torch.linspace(0, 1, num_samples).view(1, 1, num_samples, 1).to(rays.device)
    )
    log_depth_planes_bd11 = (
        torch.log(min_depth) + torch.log(max_depth / min_depth) * linear_ramp
    )
    depth_planes_bd11 = torch.exp(log_depth_planes_bd11)
    points = rays.unsqueeze(-2) * depth_planes_bd11

    points = points.view(batch_size, -1, 3)

    assert not torch.isnan(
        points
    ).any(), f"have {torch.isnan(points).count_nonzero().item()} nans in rays"
    T_pseudoCam_camera = T_camera_pseudoCam.inverse()
    points = T_pseudoCam_camera.transform(points)

    assert not torch.isnan(points).any()
    return points.view([batch_size, grid_height, grid_width, num_samples, -1])


utils/encoding_utils.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and affiliates.
# This file is derived from [ODAM](https://github.com/likojack/ODAM/blob/main/src/utils/eval_utils.py 
# and https://github.com/likojack/ODAM/blob/main/src/utils/box_utils.py).
# Originating Author: Kejie Li

# Original header:
# MIT License Copyright (c) 2022 kejieli

from copy import deepcopy

import numpy as np

import numpy as np

import torch
from numba import jit

# from utils import load_list_from_folder, fileparts, mkdir_if_missing
from scipy.spatial import ConvexHull
from scipy.optimize import linear_sum_assignment as linear_assignment


CARE_CLASSES = {
    0: "chair",
    1: "table",
    2: "cabinet",
    3: "trash bin",
    4: "bookshelf",
    5: "display",
    6: "sofa",
    7: "bathtub",
    8: "other",
}


def match_sequence(
    total_gts,
    total_preds,
    total_tps,
    predictions,
    gts,
    threshold,
    scan_id,
):
    used_gts = []
    for gt in gts:
        total_gts[gt[0]] += 1
    for prediction in predictions:
        pred_class = prediction[0]
        pred_bbx = prediction[1]
        total_preds[pred_class] += 1
        for i, gt in enumerate(gts):
            gt_class = gt[0]
            gt_bbx = gt[1]
            if gt_class == pred_class:
                rotx_matrix = rotx(np.pi / 2)
                pred_bbx_rot = (rotx_matrix @ pred_bbx[[4, 0, 1, 5, 7, 3, 2, 6]].T).T
                gt_bbx_rot = (rotx_matrix @ gt_bbx[[4, 0, 1, 5, 7, 3, 2, 6]].T).T
                iou, _ = iou3d(gt_bbx_rot, pred_bbx_rot)
                if iou > threshold and i not in used_gts:
                    used_gts.append(i)
                    total_tps[pred_class] += 1


def rotx(t):
    """3D Rotation about the x-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])


@jit
def poly_area(x, y):
    return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))


def iou3d(corners1, corners2):
    """Compute 3D bounding box IoU.
    Input:
        corners1: numpy array (8,3), assume up direction is negative Y
        corners2: numpy array (8,3), assume up direction is negative Y
    Output:
        iou: 3D bounding box IoU
        iou_2d: bird's eye view 2D bounding box IoU
    """
    if np.isnan(corners1).any() or np.isnan(corners2).any():
        return 0, 0
    try:
        # corner points are in counter clockwise order
        rect1 = [(corners1[i, 0], corners1[i, 2]) for i in range(3, -1, -1)]
        rect2 = [(corners2[i, 0], corners2[i, 2]) for i in range(3, -1, -1)]
        area1 = poly_area(np.array(rect1)[:, 0], np.array(rect1)[:, 1])
        area2 = poly_area(np.array(rect2)[:, 0], np.array(rect2)[:, 1])
        inter, inter_area = convex_hull_intersection(rect1, rect2)
        iou_2d = inter_area / (area1 + area2 - inter_area)
        ymax = min(corners1[0, 1], corners2[0, 1])
        ymin = max(corners1[4, 1], corners2[4, 1])
        inter_vol = inter_area * max(0.0, ymax - ymin)
        vol1 = box3d_vol(corners1)
        vol2 = box3d_vol(corners2)
        iou = inter_vol / (vol1 + vol2 - inter_vol)
    except:
        print(np.isnan(corners1).any())
        print(np.isnan(corners2).any())
        return 0, 0
    return iou, iou_2d


@jit
def box3d_vol(corners):
    """corners: (8,3) no assumption on axis direction"""
    a = np.sqrt(np.sum((corners[0, :] - corners[1, :]) ** 2))
    b = np.sqrt(np.sum((corners[1, :] - corners[2, :]) ** 2))
    c = np.sqrt(np.sum((corners[0, :] - corners[4, :]) ** 2))
    return a * b * c


@jit
def convex_hull_intersection(p1, p2):
    """Compute area of two convex hull's intersection area.
    p1,p2 are a list of (x,y) tuples of hull vertices.
    return a list of (x,y) for the intersection and its volume
    """
    inter_p = polygon_clip(p1, p2)
    if inter_p is not None:
        hull_inter = ConvexHull(inter_p)
        return inter_p, hull_inter.volume
    else:
        return None, 0.0


def polygon_clip(subjectPolygon, clipPolygon):
    """Clip a polygon with another polygon.
    Args:
      subjectPolygon: a list of (x,y) 2d points, any polygon.
      clipPolygon: a list of (x,y) 2d points, has to be *convex*
    Note:
      **points have to be counter-clockwise ordered**
    Return:
      a list of (x,y) vertex point for the intersection polygon.
    """

    def inside(p):
        return (cp2[0] - cp1[0]) * (p[1] - cp1[1]) > (cp2[1] - cp1[1]) * (p[0] - cp1[0])

    def computeIntersection():
        dc = [cp1[0] - cp2[0], cp1[1] - cp2[1]]
        dp = [s[0] - e[0], s[1] - e[1]]
        n1 = cp1[0] * cp2[1] - cp1[1] * cp2[0]
        n2 = s[0] * e[1] - s[1] * e[0]
        n3 = 1.0 / (dc[0] * dp[1] - dc[1] * dp[0])
        return [(n1 * dp[0] - n2 * dc[0]) * n3, (n1 * dp[1] - n2 * dc[1]) * n3]

    outputList = subjectPolygon
    cp1 = clipPolygon[-1]

    for clipVertex in clipPolygon:
        cp2 = clipVertex
        inputList = outputList
        outputList = []
        s = inputList[-1]

        for subjectVertex in inputList:
            e = subjectVertex
            if inside(e):
                if not inside(s):
                    outputList.append(computeIntersection())
                outputList.append(e)
            elif inside(s):
                outputList.append(computeIntersection())
            s = e
        cp1 = cp2
        if len(outputList) == 0:
            return None
    return outputList


def get_f1(gts, predictions, tps):
    total_gts = 0
    total_preds = 0
    total_tps = 0
    for c in CARE_CLASSES:
        if predictions[c] == 0:
            continue
        if gts[c] == 0:
            accu = 0
        else:
            accu = tps[c] / predictions[c]
        if gts[c] == 0:
            recall = 0
        else:
            recall = tps[c] / gts[c]
        print("class {}:".format(CARE_CLASSES[c]))
        print("accuracy: {}".format(accu))
        print("recall: {}".format(recall))
        f1 = 2 * accu * recall / (accu + recall) if accu + recall != 0 else 0
        print("F1: {}".format(f1))
        total_gts += gts[c]
        total_preds += predictions[c]
        total_tps += tps[c]
    if total_preds != 0:
        accuracy = total_tps / total_preds
    else:
        accuracy = 0
    if total_gts != 0:
        recall = total_tps / total_gts
    else:
        recall = 0
    if (accuracy + recall) != 0:
        f1 = 2 * accuracy * recall / (accuracy + recall)
    else:
        f1 = 0
    print("average accuracy: {}, recall: {}, F1: {}".format(accuracy, recall, f1))
    print("------------")
    return accuracy, recall, f1


def to_box(size, center, rotation):
    x_min, x_max, y_min, y_max, z_min, z_max = size
    c3o = np.stack(
        [
            x_min,
            y_min,
            z_min,
            x_max,
            y_min,
            z_min,
            x_max,
            y_max,
            z_min,
            x_min,
            y_max,
            z_min,
            x_min,
            y_min,
            z_max,
            x_max,
            y_min,
            z_max,
            x_max,
            y_max,
            z_max,
            x_min,
            y_max,
            z_max,
        ],
        axis=-1,
    )
    c3o = c3o.reshape(*c3o.shape[:-1], 8, 3)
    c3o = c3o @ rotation.transpose(-1, -2) + center[np.newaxis]
    return c3o


class F1Calculator(object):
    """Calculating f1"""

    def __init__(
        self,
        conf_thresh,
        f1_iou_thresh=[0.25, 0.5, 0.7],
    ):
        """
        """
        self.f1_iou_thresh = f1_iou_thresh
        self.conf_thresh = conf_thresh
        self.iou_thresh = 0.1
        self.score_keep_diff_class = 0.9
        self.iou_average = 0.5
        self.num_trks_keep = 4
        self.preds = {}
        self.gts = {}
        self.pred_image_level = {}
        # iou computation assume up direction is negative Y
        self.rotx_matrix = rotx(np.pi / 2)
        self.reset()

    def step(self, outputs, gt_list):
        """
        """
        batch_pred_map_cls = self.parse_predictions(
            outputs,
            self.conf_thresh,
        )

        gts = self.make_gt_list(gt_list)

        batch_pred_map_cls = self.matching_pred(
            batch_pred_map_cls, outputs["scene_name"]
        )

        gts = self.matching_gt(gts, outputs["scene_name"])

    def matching_pred(self, detections, scene_names):
        """
        """
        for i, dets in enumerate(detections):
            scene_name = scene_names[i]
            if scene_name not in self.preds.keys():
                for trk_id in range(len(dets)):
                    dets[trk_id][-1] = trk_id
                self.preds[scene_name] = deepcopy(dets)
            else:
                trks = self.preds[scene_name]
                iou_matrix = np.zeros((len(dets), len(trks)), dtype=np.float32)

                for d, det in enumerate(dets):
                    for t, trk in enumerate(trks):
                        det_rot = (
                            self.rotx_matrix @ det[1][[4, 0, 1, 5, 7, 3, 2, 6]].T
                        ).T
                        trk_rot = (
                            self.rotx_matrix @ trk[1][[4, 0, 1, 5, 7, 3, 2, 6]].T
                        ).T
                        iou_matrix[d, t] = iou3d(det_rot, trk_rot)[
                            0
                        ]  # det: 8 x 3, trk: 8 x 3

                matched_indices = linear_assignment(
                    -iou_matrix
                )  # hougarian algorithm
                matched_indices = np.stack(matched_indices, axis=1)

                unmatched_detections = []
                for d, det in enumerate(dets):
                    if d not in matched_indices[:, 0]:
                        unmatched_detections.append(d)
                unmatched_trackers = []
                for t, trk in enumerate(trks):
                    if t not in matched_indices[:, 1]:
                        unmatched_trackers.append(t)

                # filter out matched with low IOU
                matches = []
                for m in matched_indices:
                    if iou_matrix[m[0], m[1]] < self.iou_thresh:
                        unmatched_detections.append(m[0])
                        unmatched_trackers.append(m[1])
                    else:
                        matches.append(m.reshape(1, 2))

                for match in matches:
                    dets[match[:, 0].item()][-1] = trks[match[:, 1].item()][-1]
                    if trks[match[:, 1].item()][2] < dets[match[:, 0].item()][2]:
                        trks[match[:, 1].item()] = dets[match[:, 0].item()]

                pre_trk_num = len(self.preds[scene_name])
                # add new detections
                for trk_id, um in enumerate(unmatched_detections):
                    dets[um][-1] = trk_id + pre_trk_num
                    trks.append(dets[um])
                self.preds[scene_name] = deepcopy(trks)
        return detections

    def make_gt_list(self, gt_list):
        batch_gt_map_cls = []
        for gt_dict_b in gt_list:
            gt_sem_cls_labels = gt_dict_b["labels"].data.cpu().numpy()
            gt_corners = gt_dict_b["gt_corners_world"].data.cpu().numpy()

            batch_gt_map_cls.append(
                [
                    (
                        gt_sem_cls_labels[j].item(),
                        gt_corners[j] + np.random.randn(1) * 0.001,
                        1, # score, placeholder
                    )
                    for j in range(gt_corners.shape[0])
                ]
            )
        return batch_gt_map_cls

    def parse_predictions(
        self,
        outputs,
        conf_thresh,
    ):
        """Parse predictions
        Args:
        Returns:
        """
        pred_corners = outputs["pred_corners_world"]
        sem_cls_probs = outputs["sem_cls_prob"]
        pred_mask = outputs["pred_mask"]

        pred_sem_cls_prob, pred_sem_cls = torch.max(sem_cls_probs, -1)
        scores = pred_sem_cls_prob

        pred_sem_cls = pred_sem_cls.data.cpu().numpy()
        pred_corners = pred_corners.data.cpu().numpy()
        pred_sem_cls_prob = pred_sem_cls_prob.data.cpu().numpy()
        scores = scores.data.cpu().numpy()

        bsize = pred_corners.shape[0]

        batch_pred_map_cls = (
            []
        ) 
        for i in range(bsize):
            batch_pred_map_cls.append(
                [
                    [
                        pred_sem_cls[i, j].item(),
                        pred_corners[i, j],
                        scores[i, j],
                        # trk_id, always last one
                        -1,
                    ]
                    for j in range(pred_corners.shape[1])
                    if pred_sem_cls[i, j].item() != 9
                    and scores[i, j] > conf_thresh
                    and pred_mask[i, j] == 1
                ]
            )
        return batch_pred_map_cls

    def matching_gt(self, gts, scene_names):
        """
        """
        gts_snippet = []
        for i, dets in enumerate(gts):
            scene_name = scene_names[i]
            if scene_name not in self.gts.keys():
                self.gts[scene_name] = dets
                gts_snippet.append(deepcopy(dets))
            else:
                trks = self.gts[scene_name]
                iou_matrix = np.zeros((len(dets), len(trks)), dtype=np.float32)

                for d, det in enumerate(dets):
                    for t, trk in enumerate(trks):
                        det_rot = (
                            self.rotx_matrix @ det[1][[4, 0, 1, 5, 7, 3, 2, 6]].T
                        ).T
                        trk_rot = (
                            self.rotx_matrix @ trk[1][[4, 0, 1, 5, 7, 3, 2, 6]].T
                        ).T
                        iou_matrix[d, t] = iou3d(det_rot, trk_rot)[
                            0
                        ]  # det: 8 x 3, trk: 8 x 3

                matched_indices = linear_assignment(-iou_matrix)  # hougarian algorithm
                matched_indices = np.stack(matched_indices, axis=1)

                unmatched_detections = []
                for d, det in enumerate(dets):
                    if d not in matched_indices[:, 0]:
                        unmatched_detections.append(d)
                unmatched_trackers = []
                for t, trk in enumerate(trks):
                    if t not in matched_indices[:, 1]:
                        unmatched_trackers.append(t)

                # filter out matched with low IOU
                matches = []
                for m in matched_indices:
                    if iou_matrix[m[0], m[1]] < self.iou_thresh:
                        unmatched_detections.append(m[0])
                        unmatched_trackers.append(m[1])
                    else:
                        matches.append(m.reshape(1, 2))

                for match in matches:
                    if trks[match[:, 1].item()][2] < dets[match[:, 0].item()][2]:
                        trks[match[:, 1].item()] = dets[match[:, 0].item()]

                # add new detections
                for um in unmatched_detections:
                    trks.append(dets[um])
                self.gts[scene_name] = trks
                gts_snippet.append(deepcopy(trks))
        return gts_snippet

    def compute_metrics(self):
        """Use accumulated predictions and groundtruths to compute Average Precision."""

        scenes = self.preds.keys()
        metrics = {}
        for threshold in self.f1_iou_thresh:
            total_gts = {k: 0 for k in CARE_CLASSES}
            total_preds = {k: 0 for k in CARE_CLASSES}
            total_tps = {k: 0 for k in CARE_CLASSES}

            for scene in scenes:
                preds = self.preds[scene]
                match_sequence(
                    total_gts,
                    total_preds,
                    total_tps,
                    preds,
                    self.gts[scene],
                    threshold,
                    scene,
                )
            print(total_gts)
            print(total_preds)
            print(total_tps)
            print("----------")
            accuracy, recall, f1 = get_f1(total_gts, total_preds, total_tps)
            metrics["{}_accuracy".format(threshold)] = accuracy
            metrics["{}_recall".format(threshold)] = recall
            metrics["{}_f1".format(threshold)] = f1
        return metrics

    def __str__(self):
        overall_ret = self.compute_metrics()
        return self.metrics_to_str(overall_ret)

    def metrics_to_str(self, overall_ret, per_class=True):
        mAP_strs = []
        AR_strs = []
        per_class_metrics = []
        for ap_iou_thresh in self.ap_iou_thresh:
            mAP = overall_ret[ap_iou_thresh]["mAP"] * 100
            mAP_strs.append(f"{mAP:.2f}")
            ar = overall_ret[ap_iou_thresh]["AR"] * 100
            AR_strs.append(f"{ar:.2f}")

            if per_class:
                # per-class metrics
                per_class_metrics.append("-" * 5)
                per_class_metrics.append(f"IOU Thresh={ap_iou_thresh}")
                for x in list(overall_ret[ap_iou_thresh].keys()):
                    if x == "mAP" or x == "AR":
                        pass
                    else:
                        met_str = f"{x}: {overall_ret[ap_iou_thresh][x]*100:.2f}"
                        per_class_metrics.append(met_str)

        ap_header = [f"mAP{x:.2f}" for x in self.ap_iou_thresh]
        ap_str = ", ".join(ap_header)
        ap_str += ": " + ", ".join(mAP_strs)
        ap_str += "\n"

        ar_header = [f"AR{x:.2f}" for x in self.ap_iou_thresh]
        ap_str += ", ".join(ar_header)
        ap_str += ": " + ", ".join(AR_strs)

        if per_class:
            per_class_metrics = "\n".join(per_class_metrics)
            ap_str += "\n"
            ap_str += per_class_metrics

        return ap_str

    def metrics_to_dict(self, overall_ret):
        metrics_dict = {}
        for ap_iou_thresh in self.ap_iou_thresh:
            metrics_dict[f"mAP_{ap_iou_thresh}"] = (
                overall_ret[ap_iou_thresh]["mAP"] * 100
            )
            metrics_dict[f"AR_{ap_iou_thresh}"] = overall_ret[ap_iou_thresh]["AR"] * 100
        return metrics_dict

    def reset(self):
        self.preds = {}
        self.gts = {}
        self.scan_cnt = 0

utils/f1_eval.py
ymingxie/PARQ
"""
Copyright (c) Meta Platforms, Inc. and its affiliates. All Rights Reserved
This file is derived from [DETR](https://github.com/facebookresearch/detr/blob/main/models/matcher.py).
Modified for [PARQ] by Yiming Xie.

Original header:
Copyright 2020 - present, Facebook, Inc

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

"""
Modules to compute the matching cost and solve the corresponding LSAP.
"""
import torch
from scipy.optimize import linear_sum_assignment
from torch import nn
import numpy as np
        

class HungarianMatcherModified(nn.Module):
    """This is modified from the HungarianMatcher and Aside from Hungarian matching, we also match the GT box and 
    the predictions whose corresponding reference points are in close proximity to this GT box, since 
    for two adjacent reference points which have the similar queries, they should both detect nearby objects.
    """
    def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1, ratio=0.2, max_padding=10):
        """Creates the matcher
        Params:
            cost_class: This is the relative weight of the classification error in the matching cost
            cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost
            cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost
        """
        super().__init__()
        self.cost_class = cost_class
        self.cost_bbox = cost_bbox
        self.cost_giou = cost_giou
        self.ratio = ratio
        self.max_padding = max_padding
        assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0,  "all costs cant be 0"

    def forward(self, outputs, targets):
        with torch.no_grad():
            bs, num_queries = outputs["pred_logits"].shape[:2]

            out_prob = (
                outputs["pred_logits"].softmax(-1)
            )  # [batch_size * num_queries, num_classes]
            out_bbox = outputs["coord_pos"] 

            indices = []
            punish_mask_list = []
            assert bs == len(targets)
            for batch_idx in range(bs):
                bz_boxes = out_bbox[batch_idx]
                bz_out_prob = out_prob[batch_idx]
                bz_tgt_ids = targets[batch_idx]["labels"]
                bz_tgt_center = targets[batch_idx]["center"]
                bz_tgt_corners = targets[batch_idx]["gt_corners"]
                num_insts = len(bz_tgt_ids)
                if num_insts == 0:  # empty object in key frame
                    non_valid = torch.zeros(bz_out_prob.shape[0]).to(bz_out_prob) > 0
                    indices_batchi = (non_valid, torch.arange(0, 0).to(bz_out_prob))
                    indices.append(indices_batchi)
                    continue
                
                # ---hungarian matching---
                # Compute the classification cost.
                cost_class = -bz_out_prob[:, bz_tgt_ids]
                cost_bbox = torch.cdist(bz_boxes, bz_tgt_center, p=1)
                cost = self.cost_bbox * cost_bbox + self.cost_class * cost_class # + 100.0 * (~is_in_boxes_and_center)
                indices_batchi = linear_sum_assignment(cost.cpu())
                indices_batchi = list(indices_batchi)
                # -----------------------

                # ---match the GT box and the predictions whose corresponding reference points 
                # are in close proximity to this GT box---
                pred_indices = []
                gt_indices = []
                for j, box_j in enumerate(bz_tgt_corners):
                    inside_sph = cost_bbox[..., j] < self.ratio
                    pred_ind = torch.nonzero(inside_sph).squeeze(1).data.cpu().numpy()
                    punish_mask = torch.ones_like(inside_sph).bool()
                    punish_mask[pred_ind] = False
                    if pred_ind.shape[0] > self.max_padding:
                        choose = np.random.choice(pred_ind.shape[0], self.max_padding, replace=False)
                        pred_ind = pred_ind[choose]
                    punish_mask[pred_ind] = True
                    pred_indices.append(pred_ind)
                    gt_indices.append(np.ones_like(pred_ind) * j)
                pred_indices = np.concatenate(pred_indices)
                gt_indices = np.concatenate(gt_indices)
                # ----------------------
                
                indices_batchi[0] = np.concatenate([indices_batchi[0], pred_indices])
                indices_batchi[1] = np.concatenate([indices_batchi[1], gt_indices])
                
                # remove the redundant
                _, inverse_indices = np.unique(indices_batchi[0], return_index=True)
                indices_batchi[0] = indices_batchi[0][inverse_indices]
                indices_batchi[1] = indices_batchi[1][inverse_indices]

                indices.append(indices_batchi)
                punish_mask_list.append(punish_mask)

        return indices, punish_mask_list


utils/matcher.py
ymingxie/PARQ
"""
Copyright (c) Meta Platforms, Inc. and its affiliates. All Rights Reserved
This file is derived from [VoteNet](https://github.com/facebookresearch/votenet/blob/main/utils/nms.py).
Modified for [PARQ] by Yiming Xie.

Original header:
Copyright (c) Facebook, Inc. and its affiliates.

This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""

import numpy as np
# boxes are axis aigned 2D boxes of shape (n,5) in FLOAT numbers with (x1,y1,x2,y2,score)
""" Ref: https://www.pyimagesearch.com/2015/02/16/faster-non-maximum-suppression-python/
Ref: https://github.com/vickyboy47/nms-python/blob/master/nms.py
"""


def nms(
    obbs_pred, scores, num_semcls, overlap_threshold, nms_type="nms_3d_faster_samecls"
):
    # object to local
    pred_corners = obbs_pred.T_world_object.transform(obbs_pred.bb3corners_object)
    pred_corners = pred_corners.data.cpu().numpy()
    labels = obbs_pred.sem_id.data.cpu().numpy()
    scores = scores.data.cpu().numpy()
    pred_mask = run_nms(
        pred_corners, labels, scores, num_semcls, overlap_threshold, nms_type
    )
    return pred_mask


def run_nms(
    pred_corners,
    labels,
    scores,
    num_semcls,
    overlap_threshold,
    nms_type="nms_3d_faster_samecls",
):
    bs, K = pred_corners.shape[:2]
    # ---------- NMS input: pred_with_prob in (B,K,8) -----------
    pred_mask = np.zeros((bs, K)).astype(np.bool)
    for i in range(bs):
        boxes_3d_with_prob = np.zeros((K, 8))
        boxes_3d_with_prob[:, 0] = np.min(pred_corners[i, :, :, 0], axis=-1)
        boxes_3d_with_prob[:, 1] = np.min(pred_corners[i, :, :, 1], axis=-1)
        boxes_3d_with_prob[:, 2] = np.min(pred_corners[i, :, :, 2], axis=-1)
        boxes_3d_with_prob[:, 3] = np.max(pred_corners[i, :, :, 0], axis=-1)
        boxes_3d_with_prob[:, 4] = np.max(pred_corners[i, :, :, 1], axis=-1)
        boxes_3d_with_prob[:, 5] = np.max(pred_corners[i, :, :, 2], axis=-1)
        boxes_3d_with_prob[:, 6] = scores[i]
        boxes_3d_with_prob[:, 7] = labels[
            i, :, 0
        ]  # only suppress if the two boxes are of the same class!!
        bg_box_inds = np.where(labels[i, :, 0] != num_semcls)[0]
        if nms_type == "nms_3d_faster_samecls":
            pick = nms_3d_faster_samecls(
                boxes_3d_with_prob[bg_box_inds],
                overlap_threshold,
            )
        else:
            pick = nms_3d_faster(
                boxes_3d_with_prob[bg_box_inds],
                overlap_threshold,
            )
        pred_mask[i, bg_box_inds[pick]] = 1

    return pred_mask


def nms_2d(boxes, overlap_threshold):
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    score = boxes[:, 4]
    area = (x2 - x1) * (y2 - y1)

    I = np.argsort(score)
    pick = []
    while I.size != 0:
        last = I.size
        i = I[-1]
        pick.append(i)
        suppress = [last - 1]
        for pos in range(last - 1):
            j = I[pos]
            xx1 = max(x1[i], x1[j])
            yy1 = max(y1[i], y1[j])
            xx2 = min(x2[i], x2[j])
            yy2 = min(y2[i], y2[j])
            w = xx2 - xx1
            h = yy2 - yy1
            if w > 0 and h > 0:
                o = w * h / area[j]
                print("Overlap is", o)
                if o > overlap_threshold:
                    suppress.append(pos)
        I = np.delete(I, suppress)
    return pick


def nms_2d_faster(boxes, overlap_threshold, old_type=False):
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    score = boxes[:, 4]
    area = (x2 - x1) * (y2 - y1)

    I = np.argsort(score)
    pick = []
    while I.size != 0:
        last = I.size
        i = I[-1]
        pick.append(i)

        xx1 = np.maximum(x1[i], x1[I[: last - 1]])
        yy1 = np.maximum(y1[i], y1[I[: last - 1]])
        xx2 = np.minimum(x2[i], x2[I[: last - 1]])
        yy2 = np.minimum(y2[i], y2[I[: last - 1]])

        w = np.maximum(0, xx2 - xx1)
        h = np.maximum(0, yy2 - yy1)

        if old_type:
            o = (w * h) / area[I[: last - 1]]
        else:
            inter = w * h
            o = inter / (area[i] + area[I[: last - 1]] - inter)

        I = np.delete(
            I, np.concatenate(([last - 1], np.where(o > overlap_threshold)[0]))
        )

    return pick


def nms_3d_faster(boxes, overlap_threshold, old_type=False):
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    z1 = boxes[:, 2]
    x2 = boxes[:, 3]
    y2 = boxes[:, 4]
    z2 = boxes[:, 5]
    score = boxes[:, 6]
    area = (x2 - x1) * (y2 - y1) * (z2 - z1)

    I = np.argsort(score)
    pick = []
    while I.size != 0:
        last = I.size
        i = I[-1]
        pick.append(i)

        xx1 = np.maximum(x1[i], x1[I[: last - 1]])
        yy1 = np.maximum(y1[i], y1[I[: last - 1]])
        zz1 = np.maximum(z1[i], z1[I[: last - 1]])
        xx2 = np.minimum(x2[i], x2[I[: last - 1]])
        yy2 = np.minimum(y2[i], y2[I[: last - 1]])
        zz2 = np.minimum(z2[i], z2[I[: last - 1]])

        l = np.maximum(0, xx2 - xx1)
        w = np.maximum(0, yy2 - yy1)
        h = np.maximum(0, zz2 - zz1)

        if old_type:
            o = (l * w * h) / area[I[: last - 1]]
        else:
            inter = l * w * h
            o = inter / (area[i] + area[I[: last - 1]] - inter)

        I = np.delete(
            I, np.concatenate(([last - 1], np.where(o > overlap_threshold)[0]))
        )

    return pick


def nms_3d_faster_samecls(boxes, overlap_threshold, old_type=False):
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    z1 = boxes[:, 2]
    x2 = boxes[:, 3]
    y2 = boxes[:, 4]
    z2 = boxes[:, 5]
    score = boxes[:, 6]
    cls = boxes[:, 7]
    area = (x2 - x1) * (y2 - y1) * (z2 - z1)

    I = np.argsort(score)
    pick = []
    while I.size != 0:
        last = I.size
        i = I[-1]
        pick.append(i)

        xx1 = np.maximum(x1[i], x1[I[: last - 1]])
        yy1 = np.maximum(y1[i], y1[I[: last - 1]])
        zz1 = np.maximum(z1[i], z1[I[: last - 1]])
        xx2 = np.minimum(x2[i], x2[I[: last - 1]])
        yy2 = np.minimum(y2[i], y2[I[: last - 1]])
        zz2 = np.minimum(z2[i], z2[I[: last - 1]])
        cls1 = cls[i]
        cls2 = cls[I[: last - 1]]

        l = np.maximum(0, xx2 - xx1)
        w = np.maximum(0, yy2 - yy1)
        h = np.maximum(0, zz2 - zz1)

        if old_type:
            o = (l * w * h) / area[I[: last - 1]]
        else:
            inter = l * w * h
            o = inter / (area[i] + area[I[: last - 1]] - inter)
        o = o * (cls1 == cls2)

        I = np.delete(
            I, np.concatenate(([last - 1], np.where(o > overlap_threshold)[0]))
        )

    return pick


if __name__ == "__main__":
    a = np.random.random((100, 5))
    print(nms_2d(a, 0.9))
    print(nms_2d_faster(a, 0.9))

utils/nms.py
ymingxie/PARQ
"""
Copyright (c) Meta Platforms, Inc. and its affiliates. All Rights Reserved
This file is derived from [VoteNet](https://github.com/papagina/RotationContinuity/blob/master/sanity_test/code/tools.py).
Modified for [PARQ] by Yiming Xie.

Original header:
Copyright (c) 2019 Yi_Zhou

This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""

import torch

# the representation of rotation --> ortho6d
# https://github.com/papagina/RotationContinuity/blob/758b0ce551c06372cab7022d4c0bdf331c89c696/sanity_test/code/tools.py#L47
def rot_to_6d(rot_matrix):
    # batch x 3 x 3
    return torch.concat((rot_matrix[..., 0], rot_matrix[..., 1]), dim=-1)


# batch*n
def normalize_vector(v, return_mag=False):
    batch = v.shape[0]
    v_mag = torch.sqrt(v.pow(2).sum(1))  # batch
    v_mag = torch.max(
        v_mag, torch.autograd.Variable(torch.FloatTensor([1e-8])).to(v.device)
    )
    v_mag = v_mag.view(batch, 1).expand(batch, v.shape[1])
    v = v / v_mag
    if return_mag == True:
        return v, v_mag[:, 0]
    else:
        return v


# u, v batch*n
def cross_product(u, v):
    batch = u.shape[0]
    # print (u.shape)
    # print (v.shape)
    i = u[:, 1] * v[:, 2] - u[:, 2] * v[:, 1]
    j = u[:, 2] * v[:, 0] - u[:, 0] * v[:, 2]
    k = u[:, 0] * v[:, 1] - u[:, 1] * v[:, 0]

    out = torch.cat(
        (i.view(batch, 1), j.view(batch, 1), k.view(batch, 1)), 1
    )  # batch*3

    return out


def compute_rotation_matrix_from_ortho6d(ortho6d):
    x_raw = ortho6d[:, 0:3]  # batch*3
    y_raw = ortho6d[:, 3:6]  # batch*3

    x = normalize_vector(x_raw)  # batch*3
    z = cross_product(x, y_raw)  # batch*3
    z = normalize_vector(z)  # batch*3
    y = cross_product(z, x)  # batch*3

    x = x.view(-1, 3, 1)
    y = y.view(-1, 3, 1)
    z = z.view(-1, 3, 1)
    matrix = torch.cat((x, y, z), 2)  # batch*3*3
    return matrix

utils/ortho6d_transforms.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and its affiliates. All Rights Reserved

import colorsys
import itertools
from fractions import Fraction

import cv2
import numpy as np
import torch

from .matcher import HungarianMatcherModified
from .f1_eval import F1Calculator
from .ortho6d_transforms import (
    compute_rotation_matrix_from_ortho6d,
    rot_to_6d,
)
from .nms import nms

__all__ = [
    "BoxProcessor",
    "HungarianMatcherModified",
    "prepare_for_dn",
    "dn_post_process",
    "rot_to_6d",
    "compute_rotation_matrix_from_ortho6d",
    "F1Calculator",
    "nms",
]


class BoxProcessor(object):
    """
    Class to convert 3DETR MLP head outputs into bounding boxes
    Convertion between different representation of boxes (corners, dof, etc.)
    """

    def __init__(self, num_semcls, mean_size_path):
        self.num_semcls = num_semcls

        self.num_class = 9
        self.mean_size_path = mean_size_path
        if self.mean_size_path is not None:
            self.init_mean_size()

    def init_mean_size(self):
        # TODO: only for ScanNet
        self.type2class = {
            "chair": 0,
            "table": 1,
            "cabinet": 2,
            "trash bin": 3,
            "bookshelf": 4,
            "display": 5,
            "sofa": 6,
            "bathtub": 7,
            "other": 8,
        }

        self.class2type = {self.type2class[t]: t for t in self.type2class}

        self.typelong_mean_size = {}
        with open(self.mean_size_path, "r") as f:
            for line in f.readlines():
                type_cat, size = line.split(": ")
                size = size[1:-3].split(" ")
                size_ = []
                for j, s in enumerate(size):
                    if len(s) != 0:
                        size_.append(s)
                size = [float(size_[i]) for i in [0, 1, 2]]
                self.typelong_mean_size[type_cat] = size

        self.mean_size_arr = []
        self.type_mean_size = {}
        for i in range(self.num_class):
            object_type = self.class2type[i]
            for key, value in self.typelong_mean_size.items():
                key = key.split(",")
                if object_type in key:
                    self.mean_size_arr.append(value)
                    self.type_mean_size[object_type] = value
                    break

        self.mean_size_arr.append([1, 1, 1])
        self.type_mean_size["other"] = [1, 1, 1]
        self.mean_size_arr.append([1, 1, 1])
        self.type_mean_size["non-object"] = [1, 1, 1]
        self.mean_size_arr = torch.from_numpy(np.array(self.mean_size_arr))

    def compute_predicted_center(self, center_normalized):
        center_unnormalized = center_normalized * 1
        return center_normalized, center_unnormalized

    def compute_predicted_size(self, size_scale, cls_prob):
        if self.mean_size_path is not None:
            pred_cls = cls_prob.argmax(-1)
            mean_size = self.mean_size_arr[pred_cls.data.cpu()]
            size_pred = torch.exp(size_scale) * mean_size.to(size_scale.device).float()
        return size_pred

    def compute_objectness_and_cls_prob(self, cls_logits):
        assert cls_logits.shape[-1] == self.num_semcls + 1
        cls_prob = torch.nn.functional.softmax(cls_logits, dim=-1)
        objectness_prob = 1 - cls_prob[..., -1]
        return cls_prob, objectness_prob


def get_faces():
    return [
        [0, 1, 2, 3],
        [0, 3, 7, 4],
        [0, 4, 5, 1],
        [1, 2, 6, 5],
        [2, 6, 7, 3],
        [7, 4, 5, 6],
    ]


def infinite_hues():
    yield Fraction(0)
    for k in itertools.count():
        i = 2**k  # zenos_dichotomy
        for j in range(1, i, 2):
            yield Fraction(j, i)


def hue_to_hsvs(h: Fraction):
    # tweak values to adjust scheme
    for s in [Fraction(6, 10)]:
        for v in [Fraction(6, 10), Fraction(9, 10)]:
            yield (h, s, v)


def get_colors(n):
    hues = infinite_hues()
    hsvs = itertools.chain.from_iterable(hue_to_hsvs(hue) for hue in hues)
    rgbs = (colorsys.hsv_to_rgb(*hsv) for hsv in hsvs)
    return list((float(r), float(g), float(b)) for r, g, b in itertools.islice(rgbs, n))


def draw_detections(
    imgs, calibs, box_corners, T_world_object, T_pseudoCam_world, T_camera_pseudoCam, labels, tags, num_semcls, mask=None
):
    # make sure the number of cla
    id2color = get_colors(num_semcls)
    imgs_output = {}
    # default order in the img list: rgb, slaml, slamr
    for j in range(len(imgs)):
        cam = calibs[j]
        T = cam.shape[0]
        imgs_mviews = []
        imgs_withbox_mviews = []
        for t in range(T):
            cam_t = cam[t : t + 1]
            # log the first img in T
            if j == 0:
                img = imgs[j][t].permute(1, 2, 0).data.cpu().numpy()
            else:
                img = imgs[j][t, 0].data.cpu().numpy()
                img = np.stack([img, img, img], axis=-1)
            # normalization
            img = (img - img.min()) / (img.max() - img.min())
            img_withbox = img.copy()
            T_pseudoCam_object = T_pseudoCam_world[t : t + 1] @ T_world_object
            T_camera_pseudoCam_ = T_camera_pseudoCam[t : t + 1]
            for m, (box_corners_, T_pseudoCam_object_, sem_) in enumerate(zip(
                box_corners, T_pseudoCam_object, labels
            )):
                if sem_ != num_semcls:
                    if mask is not None and mask[m] == False:
                        continue
                    color = tuple((np.array(id2color[sem_])).tolist())
                    # compute projections of the corners into the image
                    box_corners_pseudo = T_pseudoCam_object_.transform(box_corners_)
                    box_corners_c = T_camera_pseudoCam_.transform(box_corners_pseudo)
                    corners_im, valid = cam_t.project(box_corners_c)

                    corners_im = corners_im[0].data.cpu().numpy()
                    valid = valid[0].data.cpu().numpy()
                    # img_withbox = img_withbox.copy()
                    faces = get_faces()
                    for face in faces:
                        for i in range(len(face) - 1):
                            if valid[face[i]] and valid[face[i + 1]]:
                                img_withbox = cv2.line(
                                    img_withbox,
                                    tuple(
                                        corners_im[face[i], :].astype(np.int).tolist()
                                    ),
                                    tuple(
                                        corners_im[face[i + 1], :]
                                        .astype(np.int)
                                        .tolist()
                                    ),
                                    color,
                                    thickness=2,
                                )
            img = torch.from_numpy(img)
            img_withbox = torch.from_numpy(img_withbox)
            if len(imgs) != 1:
                img = torch.rot90(img, -1, [0, 1])
            imgs_mviews.append(img)
            imgs_withbox_mviews.append(img_withbox)
        img = torch.cat(imgs_mviews, dim=0)
        img = img.permute((2, 0, 1)).contiguous()
        img_withbox = torch.cat(imgs_withbox_mviews, dim=0)
        img_withbox = img_withbox.permute((2, 0, 1)).contiguous()
        if "gt" not in tags[j]:
            imgs_output[tags[j]] = img
        imgs_output[tags[j] + 'withbox'] = img_withbox
    return imgs_output


def roty(t, device="cpu"):
    """Rotation about the y-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return torch.Tensor([[c, 0, s], [0, 1, 0], [-s, 0, c]]).to(device)


def rotz(t, device="cpu"):
    """Rotation about the z-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return torch.Tensor([[c, -s, 0], [s, c, 0], [0, 0, 1]]).to(device)

utils/parq_utils.py
ymingxie/PARQ
"""
Copyright (c) Meta Platforms, Inc. and its affiliates. All Rights Reserved
This file is derived from [pytorch-cosine-annealing-with-warmup](https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup).

Original header:
Copyright (c) 2022 Naoki Katsura

This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""
import math

import torch
from torch.optim.lr_scheduler import _LRScheduler


# https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup/blob/master/cosine_annealing_warmup/scheduler.py
class CosineAnnealingWarmupRestarts(_LRScheduler):
    """
    optimizer (Optimizer): Wrapped optimizer.
    first_cycle_steps (int): First cycle step size.
    cycle_mult(float): Cycle steps magnification. Default: -1.
    max_lr(float): First cycle's max learning rate. Default: 0.1.
    min_lr(float): Min learning rate. Default: 0.001.
    warmup_steps(int): Linear warmup step size. Default: 0.
    gamma(float): Decrease rate of max learning rate by cycle. Default: 1.
    last_epoch (int): The index of last epoch. Default: -1.

    """

    def __init__(
        self,
        optimizer: torch.optim.Optimizer,
        first_cycle_steps: int,
        cycle_mult: float = 1.0,
        max_lr: float = 0.1,
        min_lr: float = 0.001,
        warmup_steps: int = 0,
        gamma: float = 1.0,
        last_epoch: int = -1,
    ):

        assert warmup_steps < first_cycle_steps

        self.first_cycle_steps = first_cycle_steps  # first cycle step size
        self.cycle_mult = cycle_mult  # cycle steps magnification
        self.base_max_lr = max_lr  # first max learning rate
        self.max_lr = max_lr  # max learning rate in the current cycle
        self.min_lr = min_lr  # min learning rate
        self.warmup_steps = warmup_steps  # warmup step size
        self.gamma = gamma  # decrease rate of max learning rate by cycle

        self.cur_cycle_steps = first_cycle_steps  # first cycle step size
        self.cycle = 0  # cycle count
        self.step_in_cycle = last_epoch  # step size of the current cycle

        super(CosineAnnealingWarmupRestarts, self).__init__(optimizer, last_epoch)

        # set learning rate min_lr
        self.init_lr()

    def init_lr(self):
        self.base_lrs = []
        self.max_lrs = []
        for param_group in self.optimizer.param_groups:
            # Start with max learning rate then anealing if no warmup steps.
            param_group["lr"] = self.min_lr if self.warmup_steps != 0 else self.max_lr
            self.base_lrs.append(self.min_lr)
            self.max_lrs.append(self.max_lr)

    def get_lr(self):
        if self.step_in_cycle == -1:
            if self.warmup_steps == 0:
                return self.max_lrs
            else:
                return self.base_lrs
        elif self.step_in_cycle < self.warmup_steps:
            return [
                (self.max_lr - base_lr) * self.step_in_cycle / self.warmup_steps
                + base_lr
                for base_lr in self.base_lrs
            ]
        else:
            return [
                base_lr
                + (self.max_lr - base_lr)
                * (
                    1
                    + math.cos(
                        math.pi
                        * (self.step_in_cycle - self.warmup_steps)
                        / (self.cur_cycle_steps - self.warmup_steps)
                    )
                )
                / 2
                for base_lr in self.base_lrs
            ]

    def step(self, epoch=None):
        if epoch is None:
            epoch = self.last_epoch + 1
            self.step_in_cycle = self.step_in_cycle + 1
            if self.step_in_cycle >= self.cur_cycle_steps:
                self.cycle += 1
                self.step_in_cycle = self.step_in_cycle - self.cur_cycle_steps
                self.cur_cycle_steps = (
                    int((self.cur_cycle_steps - self.warmup_steps) * self.cycle_mult)
                    + self.warmup_steps
                )

        else:
            if epoch >= self.first_cycle_steps:
                if self.cycle_mult == 1.0:

                    self.step_in_cycle = epoch % self.first_cycle_steps
                    self.cycle = epoch // self.first_cycle_steps
                else:
                    n = int(
                        math.log(
                            (
                                epoch / self.first_cycle_steps * (self.cycle_mult - 1)
                                + 1
                            ),
                            self.cycle_mult,
                        )
                    )

                    self.cycle = n
                    self.step_in_cycle = epoch - int(
                        self.first_cycle_steps
                        * (self.cycle_mult**n - 1)
                        / (self.cycle_mult - 1)
                    )
                    self.cur_cycle_steps = self.first_cycle_steps * self.cycle_mult ** (
                        n
                    )

            else:
                self.cur_cycle_steps = self.first_cycle_steps
                self.step_in_cycle = epoch

        self.max_lr = self.base_max_lr * (self.gamma**self.cycle)
        self.last_epoch = math.floor(epoch)
        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
            param_group["lr"] = lr

utils/train_utils.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and its affiliates. All Rights Reserved

import torch


def pca_compress(img):
    B, C, H, W = img.shape
    assert B == 1
    x = img.view(C, -1).transpose(1, 0)
    x -= torch.mean(x, 0, True)
    U, S, V = torch.pca_lowrank(x, center=False)
    x = torch.matmul(x, V[:, :3])
    return x.transpose(1, 0).view(1, 3, H, W)


def normalize(img):
    return (img - img.min()) / (img.max() - img.min())
utils/vis_utils.py
ymingxie/PARQ
import torch
from collections import OrderedDict

# my_model = torch.load('/work/vig/yimingx/deper_logs_test/release/parq_23-06-09-22-03-49-epoch0000-val_loss0.000-train_loss0.000.ckpt')
official_model = torch.load('/work/vig/yimingx/deper_log_compare/model_perceiverIO_23-02-17-10-51-32/checkpoints/epoch0-step15259.ckpt')

official_model = official_model['state_dict']
new_weight_dict = OrderedDict()
new_model = {}
for key, value in official_model.items():
    if 'input_preprocessors.resnet_fpn' in key:
        ind = len('input_preprocessors.resnet_fpn')
        key = 'backbone2d' + key[ind:]
    if 'input_tokenizers.rgb_snippet.to_tokens.project' in key:
        continue
    if 'input_tokenizers.rgb_snippet' in key:
        ind = len('input_tokenizers.rgb_snippet.token_position_encoder')
        key = 'add_ray_pe' + key[ind:]
    if 'ray_offset_scale' in key:
        continue
    if 'query_tokenizers.deper.transformer' in key:
        ind = len('query_tokenizers.deper.transformer')
        key = 'box3d_decoder.parq_module' + key[ind:]
    if 'query_tokenizers.deper' in key:
        ind = len('query_tokenizers.deper')
        key = 'box3d_decoder' + key[ind:]
    new_weight_dict[key] = value

new_model['state_dict'] = new_weight_dict
torch.save(new_model, '/work/vig/yimingx/deper_log_compare/parq_release.ckpt')
utils/weight_convert.py
ymingxie/PARQ
# Copyright (c) Meta Platforms, Inc. and its affiliates. All Rights Reserved
# This file is derived from [Pixloc](https://github.com/cvg/pixloc).
# Originating Author: Paul-Edouard Sarlin
# Modified for [PARQ] by Yiming Xie.

# Original header:
# Copyright 2021 Paul-Edouard Sarlin

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import List, Tuple, Union

import torch

import functools
import inspect
import numpy as np
import math
from torch.utils.data._utils.collate import (default_collate_err_msg_format,
                                             np_str_obj_array_pattern)
from torch._six import string_classes
import collections


# https://github.com/cvg/pixloc/blob/master/pixloc/pixlib/datasets/base_dataset.py#L44
def collate(batch):
    """Difference with PyTorch default_collate: it can stack of other objects.
    """
    if not isinstance(batch, list):  # no batching
        return batch
    elem = batch[0]
    elem_type = type(elem)
    if isinstance(elem, torch.Tensor):
        out = None
        if torch.utils.data.get_worker_info() is not None:
            # If we're in a background process, concatenate directly into a
            # shared memory tensor to avoid an extra copy
            numel = sum([x.numel() for x in batch])
            storage = elem.storage()._new_shared(numel)
            out = elem.new(storage)
        return torch.stack(batch, 0, out=out)
    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
            and elem_type.__name__ != 'string_':
        if elem_type.__name__ == 'ndarray' or elem_type.__name__ == 'memmap':
            # array of string classes and object
            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
                raise TypeError(default_collate_err_msg_format.format(elem.dtype))

            return collate([torch.as_tensor(b) for b in batch])
        elif elem.shape == ():  # scalars
            return torch.as_tensor(batch)
    elif isinstance(elem, float):
        return torch.tensor(batch, dtype=torch.float64)
    elif isinstance(elem, int):
        return torch.tensor(batch)
    elif isinstance(elem, string_classes):
        return batch
    elif isinstance(elem, collections.abc.Mapping):
        return {key: collate([d[key] for d in batch]) for key in elem}
    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
        return elem_type(*(collate(samples) for samples in zip(*batch)))
    elif isinstance(elem, collections.abc.Sequence):
        # check to make sure that the elements in batch have consistent size
        it = iter(batch)
        elem_size = len(next(it))
        if not all(len(elem) == elem_size for elem in it):
            raise RuntimeError('each element in list of batch should be of equal size')
        transposed = zip(*batch)
        return [collate(samples) for samples in transposed]
    else:
        # try to stack anyway in case the object implements stacking.
        return torch.stack(batch, 0)
    

# https://github.com/cvg/pixloc/blob/master/pixloc/pixlib/geometry/wrappers.py#L17
def autocast(func):
    """Cast the inputs of a TensorWrapper method to PyTorch tensors
       if they are numpy arrays. Use the device and dtype of the wrapper.
    """
    @functools.wraps(func)
    def wrap(self, *args):
        device = torch.device('cpu')
        dtype = None
        if isinstance(self, TensorWrapper):
            if self._data is not None:
                device = self.device
                dtype = self.dtype
        elif not inspect.isclass(self) or not issubclass(self, TensorWrapper):
            raise ValueError(self)

        cast_args = []
        for arg in args:
            if isinstance(arg, np.ndarray):
                arg = torch.from_numpy(arg)
                arg = arg.to(device=device, dtype=dtype)
            cast_args.append(arg)
        return func(self, *cast_args)

    return wrap


# https://github.com/cvg/pixloc/blob/master/pixloc/pixlib/geometry/wrappers.py#L43
# Modified for [PARQ] by Yiming Xie: add squeeze, unsqueeze, clone, view
class TensorWrapper:
    _data = None

    @autocast
    def __init__(self, data: torch.Tensor):
        self._data = data

    @property
    def shape(self):
        return self._data.shape[:-1]

    @property
    def device(self):
        return self._data.device
    
    @property
    def ndim(self):
        return self._data.ndim
    
    @property
    def dtype(self):
        return self._data.dtype

    def __getitem__(self, index):
        return self.__class__(self._data[index])

    def __setitem__(self, index, item):
        self._data[index] = item.data

    def to(self, *args, **kwargs):
        return self.__class__(self._data.to(*args, **kwargs))

    def cpu(self):
        return self.__class__(self._data.cpu())

    def cuda(self):
        return self.__class__(self._data.cuda())

    def pin_memory(self):
        return self.__class__(self._data.pin_memory())

    def float(self):
        return self.__class__(self._data.float())

    def double(self):
        return self.__class__(self._data.double())

    def detach(self):
        return self.__class__(self._data.detach())
    
    def squeeze(self, dim=None):
        assert dim != -1 and dim != self._data.dim() - 1
        return self.__class__(self._data.squeeze(dim=dim))

    def unsqueeze(self, dim=None):
        assert dim != -1 and dim != self._data.dim()
        return self.__class__(self._data.unsqueeze(dim=dim))

    def clone(self):
        return self.__class__(self._data.clone())
    
    def view(self, *shape):
        assert shape[-1] == -1 or shape[-1] == self._data.shape[-1]
        return self.__class__(self._data.view(*shape))
    
    @classmethod
    def stack(cls, objects: List, dim=0, *, out=None):
        data = torch.stack([obj._data for obj in objects], dim=dim, out=out)
        return cls(data)

    def __torch_function__(self, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        if func is torch.stack:
            return self.stack(*args, **kwargs)
        else:
            return NotImplemented


# https://github.com/cvg/pixloc/blob/master/pixloc/pixlib/geometry/wrappers.py#L103
class Pose(TensorWrapper):
    def __init__(self, data: torch.Tensor):
        assert data.shape[-1] == 12
        super().__init__(data)

    @classmethod
    @autocast
    def from_Rt(cls, R: torch.Tensor, t: torch.Tensor):
        '''Pose from a rotation matrix and translation vector.
        Accepts numpy arrays or PyTorch tensors.

        Args:
            R: rotation matrix with shape (..., 3, 3).
            t: translation vector with shape (..., 3).
        '''
        assert R.shape[-2:] == (3, 3)
        assert t.shape[-1] == 3
        assert R.shape[:-2] == t.shape[:-1]
        data = torch.cat([R.flatten(start_dim=-2), t], -1)
        return cls(data)

    @classmethod
    def from_4x4mat(cls, T: torch.Tensor):
        '''Pose from an SE(3) transformation matrix.
        Args:
            T: transformation matrix with shape (..., 4, 4).
        '''
        assert T.shape[-2:] == (4, 4)
        R, t = T[..., :3, :3], T[..., :3, 3]
        return cls.from_Rt(R, t)

    @property
    def R(self) -> torch.Tensor:
        '''Underlying rotation matrix with shape (..., 3, 3).'''
        rvec = self._data[..., :9]
        return rvec.reshape(rvec.shape[:-1]+(3, 3))

    @property
    def t(self) -> torch.Tensor:
        '''Underlying translation vector with shape (..., 3).'''
        return self._data[..., -3:]
    
    @property
    def matrix(self) -> torch.Tensor:
        """Underlying transformation matrix with shape (..., 4, 4)."""
        rvec = self._data[..., :9]
        rmat = rvec.reshape(rvec.shape[:-1] + (3, 3))
        tvec = self._data[..., -3:].unsqueeze(-1)
        T_3x4 = torch.cat([rmat, tvec], dim=-1)
        bot_row = T_3x4.new_zeros(T_3x4.shape[:-2] + (1, 4))
        bot_row[..., 0, 3] = 1
        return torch.cat([T_3x4, bot_row], dim=-2)

    def inverse(self) -> 'Pose':
        '''Invert an SE(3) pose.'''
        R = self.R.transpose(-1, -2)
        t = -(R @ self.t.unsqueeze(-1)).squeeze(-1)
        return self.__class__.from_Rt(R, t)

    def compose(self, other: 'Pose') -> 'Pose':
        '''Chain two SE(3) poses: T_B2C.compose(T_A2B) -> T_A2C.'''
        R = self.R @ other.R
        t = self.t + (self.R @ other.t.unsqueeze(-1)).squeeze(-1)
        return self.__class__.from_Rt(R, t)

    @autocast
    def transform(self, p3d: torch.Tensor) -> torch.Tensor:
        '''Transform a set of 3D points.
        Args:
            p3d: 3D points, numpy array or PyTorch tensor with shape (..., 3).
        '''
        assert p3d.shape[-1] == 3
        # assert p3d.shape[:-2] == self.shape  # allow broadcasting
        return p3d @ self.R.transpose(-1, -2) + self.t.unsqueeze(-2)

    def __mul__(self, p3D: torch.Tensor) -> torch.Tensor:
        '''Transform a set of 3D points: T_A2B * p3D_A -> p3D_B.'''
        return self.transform(p3D)

    def __matmul__(self, other: 'Pose') -> 'Pose':
        '''Chain two SE(3) poses: T_B2C @ T_A2B -> T_A2C.'''
        return self.compose(other)

    def numpy(self) -> Tuple[np.ndarray]:
        return self.R.numpy(), self.t.numpy()

    def magnitude(self) -> Tuple[torch.Tensor]:
        '''Magnitude of the SE(3) transformation.
        Returns:
            dr: rotation anngle in degrees.
            dt: translation distance in meters.
        '''
        trace = torch.diagonal(self.R, dim1=-1, dim2=-2).sum(-1)
        cos = torch.clamp((trace - 1) / 2, -1, 1)
        dr = torch.acos(cos).abs() / math.pi * 180
        dt = torch.norm(self.t, dim=-1)
        return dr, dt

    def __repr__(self):
        return f'Pose: {self.shape} {self.dtype} {self.device}'



class Obb3D(TensorWrapper):
    def __init__(self, data: torch.Tensor):
        '''
            bb3_object (..., 6): 3D bounding box [xmin,xmax,ymin,ymax,zmin,zmax] in object coord frame
            T_world_object (..., 12): 3D SE3 transform from object to world coords
            sem_id (..., 1): semantic id
        '''
        assert data.shape[-1] == 19
        super().__init__(data)

    @classmethod
    def separate_init(cls, bb3_object: torch.Tensor, T_world_object: Pose, sem_id: torch.Tensor):
        if sem_id.dim() != bb3_object.dim():
            sem_id = sem_id.unsqueeze(-1)
        data = torch.cat(
            [
                bb3_object, # ..., 6
                T_world_object, # ..., 12
                sem_id, # ..., 1
            ],
            dim=-1,
        )
        return cls(data)
    
    @property
    def bb3_object(self) -> torch.Tensor:
        """3D bounding box [xmin,xmax,ymin,ymax,zmin,zmax] in object coord frame, with shape (..., 6)."""
        return self._data[..., :6]

    @property
    def bb3_min_object(self) -> torch.Tensor:
        """3D bounding box minimum corner [xmin,ymin,zmin] in object coord frame, with shape (..., 3)."""
        return self._data[..., 0:6:2]

    @property
    def bb3_max_object(self) -> torch.Tensor:
        """3D bounding box maximum corner [xmax,ymax,zmax] in object coord frame, with shape (..., 3)."""
        return self._data[..., 1:6:2]

    @property
    def bb3_center_object(self) -> torch.Tensor:
        """3D bounding box center in object coord frame, with shape (..., 3)."""
        return 0.5 * (self.bb3_min_object + self.bb3_max_object)

    @property
    def bb3_size(self) -> torch.Tensor:
        """3D bounding box diangonal, with shape (..., 3)."""
        return self.bb3_max_object - self.bb3_min_object

    @property
    def T_world_object(self) -> torch.Tensor:
        """3D SE3 transform from object to world coords, with shape (..., 12)."""
        return Pose(self._data[..., 6:18])

    @property
    def sem_id(self) -> torch.Tensor:
        """semantic id, with shape (..., 1)."""
        return self._data[..., 18].unsqueeze(-1)

    @property
    def bb3corners_object(self) -> torch.Tensor:
        """return the 8 corners of the 3D BB in object coord frame (..., 8, 3)."""
        b3o = self.bb3_object
        x_min, x_max = b3o[..., 0], b3o[..., 1]
        y_min, y_max = b3o[..., 2], b3o[..., 3]
        z_min, z_max = b3o[..., 4], b3o[..., 5]
        c3o = torch.stack(
            [
                x_min,
                y_min,
                z_min,
                x_max,
                y_min,
                z_min,
                x_max,
                y_max,
                z_min,
                x_min,
                y_max,
                z_min,
                x_min,
                y_min,
                z_max,
                x_max,
                y_min,
                z_max,
                x_max,
                y_max,
                z_max,
                x_min,
                y_max,
                z_max,
            ],
            dim=-1,
        )
        c3o = c3o.reshape(*c3o.shape[:-1], 8, 3)
        return c3o

    def add_padding(self, max_box: int = 100) -> "Obb3D":
        """
        Adds padding to Obbs to make them all the same size. Returns a new Obb3D.
        """
        assert self._data.ndim <= 2
        boxes = self._data
        num_pad = max_box - len(boxes)
        # All -1's denotes a pad box.
        pad_box = -1 * self._data.new_ones(self._data.shape[-1])
        if num_pad > 0:
            rep_box = torch.stack([pad_box for _ in range(num_pad)], dim=0)
            boxes = torch.cat([boxes, rep_box], dim=0)
        elif num_pad < 0:
            boxes = boxes[:max_box]
        return self.__class__(boxes)

    def remove_padding(self) -> List["Obb3D"]:
        """
        Removes any padding by finding Obbs with all -1s. Returns a list.
        """
        assert self._data.ndim <= 3

        if self._data.ndim == 1:
            return self

        # All -1's denotes a pad box.
        pad_box = (-1 * self._data.new_ones(self._data.shape[-1])).unsqueeze(-2)
        not_pad = ~torch.all(self._data == pad_box, dim=-1)

        if self.ndim == 2:
            num_valid = not_pad.sum()
            new_data = self.__class__(self._data[:num_valid])
        else:
            B = self._data.shape[0]
            new_data = []
            for b in range(B):
                num_valid = not_pad[b].sum()
                new_data.append(self.__class__(self._data[b][:num_valid]))
        return new_data

    def __repr__(self):
        return f"Obb3D {self.shape} {self.dtype} {self.device}"


# https://github.com/cvg/pixloc/blob/master/pixloc/pixlib/geometry/wrappers.py#L224
# Modified for [PARQ] by Yiming Xie.
class Camera(TensorWrapper):
    eps = 1e-3

    def __init__(self, data: torch.Tensor):
        assert data.shape[-1] in {6}
        super().__init__(data)

    @classmethod
    def separate_init(cls, fx, fy, cx, cy, width, height, dist=None):
        '''
        '''
        if dist is None:
            data = torch.cat([width, height, fx, fy, cx, cy], dim=-1)
        else:
            data = torch.cat([width, height, fx, fy, cx, cy, dist], dim=-1)
        return cls(data)

    @property
    def size(self) -> torch.Tensor:
        '''Size (width height) of the images, with shape (..., 2).'''
        return self._data[..., :2]

    @property
    def f(self) -> torch.Tensor:
        '''Focal lengths (fx, fy) with shape (..., 2).'''
        return self._data[..., 2:4]

    @property
    def c(self) -> torch.Tensor:
        '''Principal points (cx, cy) with shape (..., 2).'''
        return self._data[..., 4:6]

    @property
    def dist(self) -> torch.Tensor:
        '''Distortion parameters, with shape (..., {0, 2, 4}).'''
        return self._data[..., 6:]

    def scale(self, scales: Union[float, int, Tuple[Union[float, int]]]):
        '''Update the camera parameters after resizing an image.'''
        if isinstance(scales, (int, float)):
            scales = (scales, scales)
        s = self._data.new_tensor(scales)
        data = torch.cat([
            self.size*s,
            self.f*s,
            (self.c+0.5)*s-0.5,
            self.dist], -1)
        return self.__class__(data)

    def crop(self, left_top: Tuple[float], size: Tuple[int]):
        '''Update the camera parameters after cropping an image.'''
        left_top = self._data.new_tensor(left_top)
        size = self._data.new_tensor(size)
        data = torch.cat([
            size,
            self.f,
            self.c - left_top,
            self.dist], -1)
        return self.__class__(data)

    @autocast
    def in_image(self, p2d: torch.Tensor):
        '''Check if 2D points are within the image boundaries.'''
        assert p2d.shape[-1] == 2
        # assert p2d.shape[:-2] == self.shape  # allow broadcasting
        size = self.size.unsqueeze(-2)
        valid = torch.all((p2d >= 0) & (p2d <= (size - 1)), -1)
        return valid

    @autocast
    def project(self, p3d: torch.Tensor) -> Tuple[torch.Tensor]:
        '''Project 3D points into the camera plane and check for visibility.'''
        z = p3d[..., -1]
        in_front = z > self.eps
        z = z.clamp(min=self.eps)
        p2d = p3d[..., :-1] / z.unsqueeze(-1)
        
        p2d = p2d * self.f.unsqueeze(-2) + self.c.unsqueeze(-2)

        valid = in_front & self.in_image(p2d)

        return p2d, valid

    @autocast
    def unproject(self, uv: torch.Tensor) -> torch.Tensor:
        """
        Batched implementation of the Pinhole (aka Linear) camera
        model. 

        Inputs:
            uv: BxNx3 tensor of 2D pixels to be projected
            params: Bx4 tensor of Pinhole parameters formatted like this:
                    [f_u f_v c_u c_v]
        Outputs:
            xyz: BxNx3 tensor of 3D rays of uv points with z = 1.

        """

        assert uv.ndim == 3, "Expected batched input shaped BxNx3"
        B, N = uv.shape[0], uv.shape[1]

        # Focal length and principal point
        fx_fy = self.f.reshape(B, 1, 2)
        cx_cy = self.c.reshape(B, 1, 2)

        uv_dist = (uv - cx_cy) / fx_fy

        p3d = torch.cat([uv_dist, uv.new_ones(B, N, 1)], dim=2)

        return p3d

    def __repr__(self):
        return f'Camera {self.shape} {self.dtype} {self.device}'
utils/wrappers.py
ymingxie/PARQ