Mask4D: End-to-End Mask-Based 4D Panoptic Segmentation for LiDAR Sequences, RA-L, 2023 
import os
import random

import numpy as np
import torch
import yaml
from mask_4d.utils.data_util import data_prepare
from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, Dataset


class SemanticDatasetModule(LightningDataModule):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.things_ids = []
        self.color_map = []

    def prepare_data(self):
        pass

    def setup(self, stage=None):
        train_set = SemanticDataset(
            self.cfg.KITTI.PATH + "/sequences/",
            self.cfg.KITTI.CONFIG,
            split="train",
        )
        train_mask_set = MaskSemanticDataset(
            dataset=train_set,
            split="train",
            min_pts=self.cfg.KITTI.MIN_POINTS,
            space=self.cfg.KITTI.SPACE,
            num_pts=self.cfg.KITTI.SUB_NUM_POINTS,
            voxel_size=self.cfg.BACKBONE.VOXEL_SIZE,
            voxel_max=self.cfg.BACKBONE.VOXEL_MAX,
        )
        self.train_seq_mask = SequenceMaskDataset(
            train_mask_set,
            n_scans=self.cfg.TRAIN.N_SCANS,
            interval=self.cfg.TRAIN.INTERVAL,
        )

        val_set = SemanticDataset(
            self.cfg.KITTI.PATH + "/sequences/",
            self.cfg.KITTI.CONFIG,
            split="valid",
        )
        self.val_mask_set = MaskSemanticDataset(
            dataset=val_set,
            split="valid",
            min_pts=self.cfg.KITTI.MIN_POINTS,
            space=self.cfg.KITTI.SPACE,
            num_pts=0,
            voxel_size=self.cfg.BACKBONE.VOXEL_SIZE,
            voxel_max=self.cfg.BACKBONE.VOXEL_MAX,
        )

        test_set = SemanticDataset(
            self.cfg.KITTI.PATH + "/sequences/",
            self.cfg.KITTI.CONFIG,
            split="test",
        )
        self.test_mask_set = MaskSemanticDataset(
            dataset=test_set,
            split="test",
            min_pts=self.cfg.KITTI.MIN_POINTS,
            space=self.cfg.KITTI.SPACE,
            num_pts=0,
            voxel_size=self.cfg.BACKBONE.VOXEL_SIZE,
            voxel_max=self.cfg.BACKBONE.VOXEL_MAX,
        )

        self.things_ids = train_set.things_ids
        self.color_map = train_set.color_map

    def train_dataloader(self):
        dataset = self.train_seq_mask
        collate_fn = SphericalSequenceCollation()
        self.train_loader = DataLoader(
            dataset=dataset,
            batch_size=1,
            collate_fn=collate_fn,
            shuffle=True,
            num_workers=self.cfg.TRAIN.NUM_WORKERS,
            pin_memory=True,
            drop_last=False,
            timeout=0,
        )
        self.train_iter = iter(self.train_loader)
        return self.train_loader

    def val_dataloader(self):
        dataset = self.val_mask_set
        collate_fn = SphericalBatchCollation()
        self.valid_loader = DataLoader(
            dataset=dataset,
            batch_size=self.cfg.TRAIN.BATCH_SIZE,
            collate_fn=collate_fn,
            shuffle=False,
            num_workers=self.cfg.TRAIN.NUM_WORKERS,
            pin_memory=True,
            drop_last=False,
            timeout=0,
        )
        self.valid_iter = iter(self.valid_loader)
        return self.valid_loader

    def test_dataloader(self):
        dataset = self.test_mask_set
        collate_fn = SphericalBatchCollation()
        self.test_loader = DataLoader(
            dataset=dataset,
            batch_size=self.cfg.TRAIN.BATCH_SIZE,
            collate_fn=collate_fn,
            shuffle=False,
            num_workers=self.cfg.TRAIN.NUM_WORKERS,
            pin_memory=True,
            drop_last=False,
            timeout=0,
        )
        self.test_iter = iter(self.test_loader)
        return self.test_loader


class SemanticDataset(Dataset):
    def __init__(self, data_path, cfg_path, split="train"):
        yaml_path = cfg_path
        with open(yaml_path, "r") as stream:
            semyaml = yaml.safe_load(stream)

        self.things = get_things()
        self.stuff = get_stuff()

        self.things_ids = get_things_ids()

        self.color_map = semyaml["color_map_learning"]
        self.labels = semyaml["labels"]
        self.learning_map = semyaml["learning_map"]
        self.inv_learning_map = semyaml["learning_map_inv"]
        self.split = split
        split = semyaml["split"][self.split]

        self.im_idx = []
        pose_files = []
        calib_files = []
        fill = 2
        for i_folder in split:
            self.im_idx += absoluteFilePaths(
                "/".join([data_path, str(i_folder).zfill(fill), "velodyne"])
            )
            pose_files.append(
                absoluteDirPath(
                    "/".join([data_path, str(i_folder).zfill(fill), "poses.txt"])
                )
            )
            calib_files.append(
                absoluteDirPath(
                    "/".join([data_path, str(i_folder).zfill(fill), "calib.txt"])
                )
            )

        self.im_idx.sort()
        self.poses, self.n_scans = load_poses(pose_files, calib_files)

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

    def __getitem__(self, index):
        fname = self.im_idx[index]
        pose = self.poses[index]
        points = np.fromfile(self.im_idx[index], dtype=np.float32).reshape((-1, 4))
        xyz = points[:, :3]
        intensity = points[:, 3]
        if len(intensity.shape) == 2:
            intensity = np.squeeze(intensity)
        if self.split == "test":
            annotated_data = np.expand_dims(
                np.zeros_like(points[:, 0], dtype=int), axis=1
            )
            sem_labels = annotated_data
            ins_labels = annotated_data
        else:
            annotated_data = np.fromfile(
                self.im_idx[index].replace("velodyne", "labels")[:-3] + "label",
                dtype=np.int32,
            ).reshape((-1, 1))
            sem_labels = annotated_data & 0xFFFF
            ins_labels = annotated_data >> 16
            sem_labels = np.vectorize(self.learning_map.__getitem__)(sem_labels)

        return (xyz, sem_labels, ins_labels, intensity, fname, pose)


class MaskSemanticDataset(Dataset):
    def __init__(
        self,
        dataset,
        split,
        min_pts,
        space,
        num_pts,
        voxel_size,
        voxel_max,
    ):
        self.dataset = dataset
        self.num_points = num_pts
        self.split = split
        self.min_points = min_pts
        self.th_ids = dataset.things_ids
        self.xlim = space[0]
        self.ylim = space[1]
        self.zlim = space[2]
        self.voxel_size = np.array(voxel_size)
        self.voxel_max = voxel_max

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

    def __getitem__(self, index):
        data = self.dataset[index]
        xyz, sem_labels, ins_labels, intensity, fname, pose = data
        full_xyz = xyz.shape[0]
        keep_xyz = np.argwhere(
            (self.xlim[0] < xyz[:, 0])
            & (xyz[:, 0] < self.xlim[1])
            & (self.ylim[0] < xyz[:, 1])
            & (xyz[:, 1] < self.ylim[1])
            & (self.zlim[0] < xyz[:, 2])
            & (xyz[:, 2] < self.zlim[1])
        )[:, 0]
        xyz = xyz[keep_xyz]
        sem_labels = sem_labels[keep_xyz]
        ins_labels = ins_labels[keep_xyz]
        intensity = intensity[keep_xyz]

        feats = np.concatenate((xyz, np.expand_dims(intensity, axis=1)), axis=1)

        # Subsample
        if self.split == "train" and len(xyz) > self.num_points:
            idx = np.random.choice(np.arange(len(xyz)), self.num_points, replace=False)
            xyz = xyz[idx]
            sem_labels = sem_labels[idx]
            ins_labels = ins_labels[idx]
            feats = feats[idx]
            intensity = intensity[idx]

        # SphereFormer
        sp_xyz = xyz.copy()
        sp_coords, sp_xyz, sp_feats, sp_labels, sp_idx_recons = data_prepare(
            sp_xyz,
            feats,
            sem_labels,
            self.split,
            self.voxel_size,
            self.voxel_max,
        )

        if self.split == "test":
            return (
                xyz,
                sem_labels,
                ins_labels,
                torch.tensor([]),
                torch.tensor([]),
                torch.tensor([]),
                fname,
                pose,
                keep_xyz,
                full_xyz,
                sp_coords,
                sp_xyz,
                sp_feats,
                sp_labels,
                sp_idx_recons,
            )

        stuff_masks = np.array([]).reshape(0, xyz.shape[0])
        stuff_masks_ids = []
        things_masks = np.array([]).reshape(0, xyz.shape[0])
        things_cls = np.array([], dtype=int)
        things_masks_ids = []

        stuff_labels = np.asarray(
            [0 if s in self.th_ids else s for s in sem_labels[:, 0]]
        )
        stuff_cls, st_cnt = np.unique(stuff_labels, return_counts=True)
        # filter small masks
        keep_st = np.argwhere(st_cnt > self.min_points)[:, 0]
        stuff_cls = stuff_cls[keep_st][1:]
        if len(stuff_cls):
            stuff_masks = np.array(
                [np.where(stuff_labels == i, 1.0, 0.0) for i in stuff_cls]
            )
            stuff_masks_ids = [0 for m in stuff_masks]
        # things masks
        ins_sems = np.where(ins_labels == 0, 0, sem_labels)
        _ins_labels = ins_sems + ((ins_labels << 16) & 0xFFFF0000).reshape(-1, 1)
        things_ids, th_idx, th_cnt = np.unique(
            _ins_labels[:, 0], return_index=True, return_counts=True
        )
        # filter small instances
        keep_th = np.argwhere(th_cnt > self.min_points)[:, 0]
        things_ids = things_ids[keep_th]
        th_idx = th_idx[keep_th]
        # remove instances with wrong sem class
        keep_th = np.array(
            [i for i, idx in enumerate(th_idx) if sem_labels[idx] in self.th_ids],
            dtype=int,
        )
        things_ids = things_ids[keep_th]
        th_idx = th_idx[keep_th]
        if len(th_idx):
            things_masks = np.array(
                [np.where(_ins_labels[:, 0] == i, 1.0, 0.0) for i in things_ids]
            )
            things_cls = np.array([sem_labels[i] for i in th_idx]).squeeze(1)
            things_masks_ids = [t for t in things_ids]

        masks = torch.from_numpy(np.concatenate((stuff_masks, things_masks)))
        masks_cls = torch.from_numpy(np.concatenate((stuff_cls, things_cls)))
        stuff_masks_ids.extend(things_masks_ids)
        masks_ids = torch.tensor(stuff_masks_ids)

        assert (
            masks.shape[0] == masks_cls.shape[0]
        ), f"not same number masks and classes: masks {masks.shape[0]}, classes {masks_cls.shape[0]} "

        return (
            xyz,
            sem_labels,
            ins_labels,
            masks,
            masks_cls,
            masks_ids,
            fname,
            pose,
            keep_xyz,
            full_xyz,
            sp_coords,
            sp_xyz,
            sp_feats,
            sp_labels,
            sp_idx_recons,
        )


class SequenceMaskDataset(Dataset):
    def __init__(self, dataset, n_scans, interval):
        super().__init__()
        self.dataset = dataset
        self.n_scans = n_scans
        self.interval = interval

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

    def __getitem__(self, index):
        data = self.dataset[index]
        fname = data[6]
        _fname = fname.split("/")
        scan = int(_fname[-1][:-4])
        seq = _fname[-3]
        max_scan = self.dataset.dataset.n_scans[seq] - 1

        # select n_scans randomly between scan-interval/2 and scan+interval/2
        s_before = random.sample(
            set(np.arange(scan - int(self.interval / 2), scan)), round(self.n_scans / 2)
        )
        s_after = random.sample(
            set(np.arange(scan, scan + int(self.interval / 2))), self.n_scans // 2
        )
        scans = s_before + s_after
        scans.sort(reverse=True)
        while scans[-1] < 0:
            scans.pop()
        scans.sort()
        while scans[-1] > max_scan:
            scans.pop()
        idx = index + np.array(scans) - scan
        np.append(idx, np.array([index]))
        idx.sort()
        out_data = []
        for i in idx:
            out_data.append(self.dataset[i])

        return out_data


class SphericalBatchCollation:
    def __init__(self):
        self.keys = [
            "pt_coord",
            "sem_label",
            "ins_label",
            "masks",
            "masks_cls",
            "masks_ids",
            "fname",
            "pose",
            "keep_xyz",
            "full_xyz",
            "sp_coord",
            "sp_xyz",
            "sp_feat",
            "sp_label",
            "sp_idx_recons",
        ]

    def __call__(self, data):
        return {self.keys[i]: list(x) for i, x in enumerate(zip(*data))}


class SphericalSequenceCollation:
    def __init__(self):
        self.keys = [
            "pt_coord",
            "sem_label",
            "ins_label",
            "masks",
            "masks_cls",
            "masks_ids",
            "fname",
            "pose",
            "keep_xyz",
            "full_xyz",
            "sp_coord",
            "sp_xyz",
            "sp_feat",
            "sp_label",
            "sp_idx_recons",
        ]

    def __call__(self, data):
        _data = data[0]  # always bs=0
        return {self.keys[i]: list(x) for i, x in enumerate(zip(*_data))}


def absoluteFilePaths(directory):
    for dirpath, _, filenames in os.walk(directory):
        for f in filenames:
            yield os.path.abspath(os.path.join(dirpath, f))


def absoluteDirPath(directory):
    return os.path.abspath(directory)


def parse_calibration(filename):
    calib = {}
    calib_file = open(filename)
    for line in calib_file:
        key, content = line.strip().split(":")
        values = [float(v) for v in content.strip().split()]
        pose = np.zeros((4, 4))
        pose[0, 0:4] = values[0:4]
        pose[1, 0:4] = values[4:8]
        pose[2, 0:4] = values[8:12]
        pose[3, 3] = 1.0
        calib[key] = pose
    calib_file.close()
    return calib


def parse_poses(filename, calibration):
    file = open(filename)
    poses = []
    Tr = calibration["Tr"]
    Tr_inv = np.linalg.inv(Tr)
    for line in file:
        values = [float(v) for v in line.strip().split()]
        pose = np.zeros((4, 4))
        pose[0, 0:4] = values[0:4]
        pose[1, 0:4] = values[4:8]
        pose[2, 0:4] = values[8:12]
        pose[3, 3] = 1.0
        poses.append(np.matmul(Tr_inv, np.matmul(pose, Tr)))
    return poses


def load_poses(pose_files, calib_files):
    poses = []
    n_scans = {}
    # go through every file and get all poses
    # add them to match im_idx
    for i in range(len(pose_files)):
        calib = parse_calibration(calib_files[i])
        seq_poses_f64 = parse_poses(pose_files[i], calib)
        seq_poses = [pose.astype(np.float32) for pose in seq_poses_f64]
        poses += seq_poses
        seq = pose_files[i].split("/")[-2]
        n_scans[seq] = len(seq_poses)
    return poses, n_scans


def getDir(obj):
    return os.path.dirname(os.path.abspath(obj))


def get_things():
    things = {
        1: "car",
        2: "bicycle",
        3: "motorcycle",
        4: "truck",
        5: "other-vehicle",
        6: "person",
        7: "bicyclist",
        8: "motorcyclist",
    }
    return things


def get_stuff():
    stuff = {
        9: "road",
        10: "parking",
        11: "sidewalk",
        12: "other-ground",
        13: "building",
        14: "fence",
        15: "vegetation",
        16: "trunk",
        17: "terrain",
        18: "pole",
        19: "traffic-sign",
    }
    return stuff


def get_things_ids():
    return [1, 2, 3, 4, 5, 6, 7, 8]

mask_4d/datasets/kitti_dataset.py
PRBonn/Mask4D
# Modified by Rodrigo Marcuzzi from https://github.com/dvlab-research/SphereFormer/blob/master/model/unet_spherical_transformer.py
import functools
from collections import OrderedDict

import numpy as np
import spconv.pytorch as spconv
import torch
import torch.nn as nn
import torch.nn.functional as F
from mask_4d.models.spherical_transformer import SphereFormer
from spconv.core import ConvAlgo
from spconv.pytorch.modules import SparseModule
from torch_scatter import scatter_mean


class ResidualBlock(SparseModule):
    def __init__(self, in_channels, out_channels, norm_fn, indice_key=None):
        super().__init__()
        if in_channels == out_channels:
            self.i_branch = spconv.SparseSequential(nn.Identity())
        else:
            self.i_branch = spconv.SparseSequential(
                spconv.SubMConv3d(in_channels, out_channels, kernel_size=1, bias=False)
            )
        self.conv_branch = spconv.SparseSequential(
            norm_fn(in_channels),
            nn.ReLU(),
            spconv.SubMConv3d(
                in_channels,
                out_channels,
                kernel_size=3,
                padding=1,
                bias=False,
                indice_key=indice_key,
            ),
            norm_fn(out_channels),
            nn.ReLU(),
            spconv.SubMConv3d(
                out_channels,
                out_channels,
                kernel_size=3,
                padding=1,
                bias=False,
                indice_key=indice_key,
            ),
        )

    def forward(self, input):
        identity = spconv.SparseConvTensor(
            input.features, input.indices, input.spatial_shape, input.batch_size
        )
        output = self.conv_branch(input)
        output = output.replace_feature(
            output.features + self.i_branch(identity).features
        )
        return output


def get_downsample_info(xyz, batch, indice_pairs):
    pair_in, pair_out = indice_pairs[0], indice_pairs[1]
    valid_mask = pair_in != -1
    valid_pair_in, valid_pair_out = (
        pair_in[valid_mask].long(),
        pair_out[valid_mask].long(),
    )
    xyz_next = scatter_mean(xyz[valid_pair_in], index=valid_pair_out, dim=0)
    batch_next = scatter_mean(batch.float()[valid_pair_in], index=valid_pair_out, dim=0)
    return xyz_next, batch_next


class UBlock(nn.Module):
    def __init__(
        self,
        nPlanes,
        norm_fn,
        block_reps,
        block,
        window_size,
        window_size_sphere,
        quant_size,
        quant_size_sphere,
        head_dim=16,
        window_size_scale=[2.0, 2.0],
        rel_query=True,
        rel_key=True,
        rel_value=True,
        drop_path=0.0,
        indice_key_id=1,
        grad_checkpoint_layers=[],
        sphere_layers=[1, 2, 3, 4, 5],
        a=0.05 * 0.25,
    ):
        super().__init__()

        self.nPlanes = nPlanes
        self.indice_key_id = indice_key_id
        self.grad_checkpoint_layers = grad_checkpoint_layers
        self.sphere_layers = sphere_layers

        blocks = {
            "block{}".format(i): block(
                nPlanes[0],
                nPlanes[0],
                norm_fn,
                indice_key="subm{}".format(indice_key_id),
            )
            for i in range(block_reps)
        }
        blocks = OrderedDict(blocks)
        self.blocks = spconv.SparseSequential(blocks)

        if indice_key_id in sphere_layers:
            self.window_size = window_size
            self.window_size_sphere = window_size_sphere
            num_heads = nPlanes[0] // head_dim
            self.transformer_block = SphereFormer(
                nPlanes[0],
                num_heads,
                window_size,
                window_size_sphere,
                quant_size,
                quant_size_sphere,
                indice_key="sphereformer{}".format(indice_key_id),
                rel_query=rel_query,
                rel_key=rel_key,
                rel_value=rel_value,
                drop_path=drop_path[0],
                a=a,
            )

        if len(nPlanes) > 1:
            self.conv = spconv.SparseSequential(
                norm_fn(nPlanes[0]),
                nn.ReLU(),
                spconv.SparseConv3d(
                    nPlanes[0],
                    nPlanes[1],
                    kernel_size=2,
                    stride=2,
                    bias=False,
                    indice_key="spconv{}".format(indice_key_id),
                    algo=ConvAlgo.Native,
                ),
            )

            window_size_scale_cubic, window_size_scale_sphere = window_size_scale
            window_size_next = np.array(
                [
                    window_size[0] * window_size_scale_cubic,
                    window_size[1] * window_size_scale_cubic,
                    window_size[2] * window_size_scale_cubic,
                ]
            )
            quant_size_next = np.array(
                [
                    quant_size[0] * window_size_scale_cubic,
                    quant_size[1] * window_size_scale_cubic,
                    quant_size[2] * window_size_scale_cubic,
                ]
            )
            window_size_sphere_next = np.array(
                [
                    window_size_sphere[0] * window_size_scale_sphere,
                    window_size_sphere[1] * window_size_scale_sphere,
                    window_size_sphere[2],
                ]
            )
            quant_size_sphere_next = np.array(
                [
                    quant_size_sphere[0] * window_size_scale_sphere,
                    quant_size_sphere[1] * window_size_scale_sphere,
                    quant_size_sphere[2],
                ]
            )
            self.u = UBlock(
                nPlanes[1:],
                norm_fn,
                block_reps,
                block,
                window_size_next,
                window_size_sphere_next,
                quant_size_next,
                quant_size_sphere_next,
                window_size_scale=window_size_scale,
                rel_query=rel_query,
                rel_key=rel_key,
                rel_value=rel_value,
                drop_path=drop_path[1:],
                indice_key_id=indice_key_id + 1,
                grad_checkpoint_layers=grad_checkpoint_layers,
                sphere_layers=sphere_layers,
                a=a,
            )

            self.deconv = spconv.SparseSequential(
                norm_fn(nPlanes[1]),
                nn.ReLU(),
                spconv.SparseInverseConv3d(
                    nPlanes[1],
                    nPlanes[0],
                    kernel_size=2,
                    bias=False,
                    indice_key="spconv{}".format(indice_key_id),
                    algo=ConvAlgo.Native,
                ),
            )

            blocks_tail = {}
            for i in range(block_reps):
                blocks_tail["block{}".format(i)] = block(
                    nPlanes[0] * (2 - i),
                    nPlanes[0],
                    norm_fn,
                    indice_key="subm{}".format(indice_key_id),
                )
            blocks_tail = OrderedDict(blocks_tail)
            self.blocks_tail = spconv.SparseSequential(blocks_tail)

    def forward(self, inp, xyz, batch):
        assert (inp.indices[:, 0] == batch).all()

        output = self.blocks(inp)

        # transformer
        if self.indice_key_id in self.sphere_layers:
            if self.indice_key_id in self.grad_checkpoint_layers:

                def run(feats_, xyz_, batch_):
                    return self.transformer_block(feats_, xyz_, batch_)

                transformer_features = torch.utils.checkpoint.checkpoint(
                    run, output.features, xyz, batch
                )
            else:
                transformer_features = self.transformer_block(
                    output.features, xyz, batch
                )
            output = output.replace_feature(transformer_features)

        identity = spconv.SparseConvTensor(
            output.features, output.indices, output.spatial_shape, output.batch_size
        )

        if len(self.nPlanes) > 1:
            output_decoder = self.conv(output)

            # downsample
            indice_pairs = output_decoder.indice_dict[
                "spconv{}".format(self.indice_key_id)
            ].indice_pairs
            xyz_next, batch_next = get_downsample_info(xyz, batch, indice_pairs)

            output_decoder = self.u(output_decoder, xyz_next, batch_next.long())
            output_decoder = self.deconv(output_decoder)
            output = output.replace_feature(
                torch.cat((identity.features, output_decoder.features), dim=1)
            )
            output = self.blocks_tail(output)

        return output


class SphericalEncoderDecoder(nn.Module):
    def __init__(self, cfg, data_cfg):
        super().__init__()
        input_c = cfg.IN_CHANNELS
        m = cfg.CHANNELS[0]
        classes = data_cfg.NUM_CLASSES
        block_reps = cfg.BLOCK_REPS
        channels = cfg.CHANNELS
        window_size = np.array(cfg.WINDOW_SIZE)
        window_size_sphere = np.array(cfg.WINDOW_SIZE_SPHERE)
        quant_size = 1 / np.array(cfg.QUANT_SIZE)
        quant_size_sphere = 1 / np.array(cfg.QUANT_SIZE_SPHERE)
        rel_query = True
        rel_key = True
        rel_value = True
        drop_path_rate = cfg.DROP_PATH_RATE
        window_size_scale = cfg.WINDOW_SIZE_SCALE
        grad_checkpoint_layers = []
        sphere_layers = cfg.SPHERE_LAYERS
        a = 0.05 * 0.25

        norm_fn = functools.partial(nn.BatchNorm1d, eps=1e-4, momentum=0.1)

        block = ResidualBlock

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, 7)]

        #### backbone
        self.input_conv = spconv.SparseSequential(
            spconv.SubMConv3d(
                input_c, m, kernel_size=3, padding=1, bias=False, indice_key="subm1"
            )
        )

        self.unet = UBlock(
            channels,
            norm_fn,
            block_reps,
            block,
            window_size,
            window_size_sphere,
            quant_size,
            quant_size_sphere,
            window_size_scale=window_size_scale,
            rel_query=rel_query,
            rel_key=rel_key,
            rel_value=rel_value,
            drop_path=dpr,
            indice_key_id=1,
            grad_checkpoint_layers=grad_checkpoint_layers,
            sphere_layers=sphere_layers,
            a=a,
        )

        self.output_layer = spconv.SparseSequential(norm_fn(m), nn.ReLU())

        #### semantic segmentation
        self.linear = nn.Linear(m, classes - 1)  # bias(default): True

        self.apply(self.set_bn_init)

    @staticmethod
    def set_bn_init(m):
        classname = m.__class__.__name__
        if classname.find("BatchNorm") != -1:
            m.weight.data.fill_(1.0)
            m.bias.data.fill_(0.0)

    def forward(self, x):
        coord = torch.cat(x["sp_coord"])
        xyz = torch.cat(x["sp_xyz"])
        feat = torch.cat(x["sp_feat"])
        idx_recon = x["sp_idx_recons"]
        _coord = np.insert(
            np.cumsum(np.array([cc.shape[0] for cc in x["sp_coord"]])), 0, 0
        )[:-1]
        # remap idx_recon to stack tensor
        idx_recon = torch.cat([ii + cc for ii, cc in zip(idx_recon, _coord)])

        offset_ = [cc.shape[0] for cc in x["sp_coord"]]
        batch = (
            torch.cat([torch.tensor([ii] * o) for ii, o in enumerate(offset_)], 0)
            .long()
            .cuda()
        )

        coord = torch.cat([batch.unsqueeze(-1), coord], -1)
        coord[:, 1:] += (torch.rand(3) * 2).type_as(coord)
        spatial_shape = np.clip((coord.max(0)[0][1:] + 1).cpu().numpy(), 128, None)

        sinput = spconv.SparseConvTensor(
            feat, coord.int(), spatial_shape, len(x["pt_coord"])
        )

        output0 = self.input_conv(sinput)
        output1 = self.unet(output0, xyz, batch)

        # reconstruct full pcd to get per point features
        pt_feat = output1.features[idx_recon, :]

        coors = torch.from_numpy(x["pt_coord"][0]).to(pt_feat.device).unsqueeze(0)
        logits = self.linear(pt_feat)
        logits = torch.hstack(
            (torch.zeros(logits.shape[0]).unsqueeze(1).cuda(), logits)
        )
        return pt_feat, coors, logits

mask_4d/models/backbone.py
PRBonn/Mask4D
from typing import Optional

import torch
from mask_4d.models.position_attention import PositionAttention
from torch import Tensor, nn
from torch.nn import functional as F


class SelfAttentionLayer(nn.Module):
    def __init__(self, d_model, nhead, dropout=0.0, pre_norm=False, activation="relu"):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(
            d_model, nhead, dropout=dropout, batch_first=True
        )

        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.pre_norm = pre_norm

        self._reset_parameters()

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

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

    def forward(
        self,
        q_embed,
        attn_mask: Optional[Tensor] = None,
        padding_mask: Optional[Tensor] = None,
        query_pos: Optional[Tensor] = None,
    ):
        if self.pre_norm:
            q_embed = self.norm(q_embed)
            q = k = self.with_pos_embed(q_embed, query_pos)
            q_embed2 = self.self_attn(
                q, k, value=q_embed, attn_mask=attn_mask, key_padding_mask=padding_mask
            )[0]
            q_embed = q_embed + self.dropout(q_embed2)
        else:
            q = k = self.with_pos_embed(q_embed, query_pos)
            q_embed2 = self.self_attn(
                q, k, value=q_embed, attn_mask=attn_mask, key_padding_mask=padding_mask
            )[0]
            q_embed = q_embed + self.dropout(q_embed2)
            q_embed = self.norm(q_embed)
        return q_embed


class CrossAttentionLayer(nn.Module):
    def __init__(self, d_model, nhead, dropout=0.0, pre_norm=False, activation="relu"):
        super().__init__()
        self.multihead_attn = nn.MultiheadAttention(
            d_model, nhead, dropout=dropout, batch_first=True
        )

        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.pre_norm = pre_norm

        self._reset_parameters()

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

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

    def with_pos_embed2(self, tensor, pos: Optional[Tensor]):
        out = torch.cat((tensor, pos.unsqueeze(0)), dim=-1)
        return out

    def forward(
        self,
        q_embed,
        bb_feat,
        attn_mask: Optional[Tensor] = None,
        padding_mask: Optional[Tensor] = None,
        pos: Optional[Tensor] = None,
        query_pos: Optional[Tensor] = None,
    ):
        if self.pre_norm:
            q_embed2 = self.multihead_attn(
                query=self.with_pos_embed(q_embed, query_pos),
                key=self.with_pos_embed(bb_feat, pos),
                value=self.with_pos_embed(bb_feat, pos),
                # value=bb_feat,
                attn_mask=attn_mask,
                key_padding_mask=padding_mask,
            )[0]
            q_embed = q_embed + self.dropout(q_embed2)
            q_embed = self.norm(q_embed)
        else:
            q_embed = self.norm(q_embed)
            q_embed2 = self.multihead_attn(
                query=self.with_pos_embed(q_embed, query_pos),
                key=self.with_pos_embed(bb_feat, pos),
                value=self.with_pos_embed(bb_feat, pos),
                # value=bb_feat,
                attn_mask=attn_mask,
                key_padding_mask=padding_mask,
            )[0]
            q_embed = q_embed + self.dropout(q_embed2)
        return q_embed


class FFNLayer(nn.Module):
    def __init__(
        self,
        d_model,
        dim_feedforward=2048,
        dropout=0.0,
        pre_norm=False,
        activation="relu",
    ):
        super().__init__()
        # 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.norm = nn.LayerNorm(d_model)

        self.activation = _get_activation_fn(activation)
        self.pre_norm = pre_norm

        self._reset_parameters()

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

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

    def forward(self, tgt):
        if self.pre_norm:
            tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
            tgt = tgt + self.dropout(tgt2)
            tgt = self.norm(tgt)
        else:
            tgt = self.norm(tgt)
            tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
            tgt = tgt + self.dropout(tgt2)
        return tgt


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}.")


class MLP(nn.Module):
    """Very simple multi-layer perceptron (also called FFN)"""

    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(
            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
        )

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x


class PositionCrossAttentionLayer(nn.Module):
    def __init__(
        self,
        d_model,
        nhead,
        dropout=0.0,
        activation="relu",
    ):
        super().__init__()
        self.multihead_attn = PositionAttention(d_model, nhead, dropout=dropout)

        self.nhead = nhead

        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)

        self._reset_parameters()

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

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

    def forward(
        self,
        q_embed,
        kernel,
        bb_feat,
        attn_mask: Optional[Tensor] = None,
        padding_mask: Optional[Tensor] = None,
        pos: Optional[Tensor] = None,
        query_pos: Optional[Tensor] = None,
    ):
        q_embed = self.norm(q_embed)
        q_embed2 = self.multihead_attn(
            query=self.with_pos_embed(q_embed, query_pos),
            key=self.with_pos_embed(bb_feat, pos),
            value=self.with_pos_embed(bb_feat, pos),
            attn_mask=attn_mask,
            key_padding_mask=padding_mask,
            kernel=kernel,
        )
        q_embed = q_embed + self.dropout(q_embed2)
        return q_embed

mask_4d/models/blocks.py
PRBonn/Mask4D
# Modified from https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/transformer_decoder/mask2former_transformer_decoder.py
import mask_4d.models.blocks as blocks
import mask_4d.utils.misc as misc
import torch
from mask_4d.models.positional_encoder import PositionalEncoder
from torch import nn


class MaskedTransformerDecoder(nn.Module):
    def __init__(self, cfg, bb_cfg, data_cfg):
        super().__init__()
        hidden_dim = cfg.HIDDEN_DIM
        self.cfg = cfg
        self.pe_layer = PositionalEncoder(cfg.POS_ENC)
        self.num_layers = cfg.FEATURE_LEVELS * cfg.DEC_BLOCKS
        self.nheads = cfg.NHEADS
        self.num_queries = cfg.NUM_QUERIES
        self.num_feature_levels = cfg.FEATURE_LEVELS

        self.decoder_norm = nn.LayerNorm(hidden_dim)
        self.mask_feat_proj = nn.Sequential()
        self.transformer_self_attention_layers = nn.ModuleList()
        self.transformer_cross_attention_layers = nn.ModuleList()
        self.transformer_ffn_layers = nn.ModuleList()

        for _ in range(self.num_layers):
            self.transformer_self_attention_layers.append(
                blocks.SelfAttentionLayer(hidden_dim, self.nheads)
            )
            self.transformer_cross_attention_layers.append(
                blocks.PositionCrossAttentionLayer(hidden_dim, self.nheads)
            )

            self.transformer_ffn_layers.append(
                blocks.FFNLayer(hidden_dim, cfg.DIM_FEEDFORWARD)
            )

        if bb_cfg.CHANNELS[0] != hidden_dim:
            self.mask_feat_proj = nn.Linear(bb_cfg.CHANNELS[0], hidden_dim)
        in_channels = self.num_feature_levels * [bb_cfg.CHANNELS[0]]

        self.input_proj = nn.ModuleList()
        for ch in in_channels:
            if ch != hidden_dim:  # linear projection to hidden_dim
                self.input_proj.append(nn.Linear(ch, hidden_dim))
            else:
                self.input_proj.append(nn.Sequential())

        self.class_embed = nn.Linear(hidden_dim, data_cfg.NUM_CLASSES + 1)
        self.mask_embed = blocks.MLP(hidden_dim, hidden_dim, hidden_dim, 3)

    def forward(self, feats, coors, track_ins):
        mask_features = self.mask_feat_proj(feats) + self.pe_layer(coors)
        src = []
        pos = []
        queries = []

        for i in range(self.num_feature_levels):
            pos.append(self.pe_layer(coors))
            feat = self.input_proj[i](feats)
            src.append(feat)

        query_embed = track_ins.query_pe.unsqueeze(0)
        output = track_ins.query.unsqueeze(0)
        q_centers = track_ins.center
        size = track_ins.size_xy
        angle = track_ins.angle

        predictions_class = []
        predictions_mask = []

        if output.shape[-1] > self.num_queries:
            mask_kernel = self.compute_kernel(q_centers, coors, size, angle, mask=True)
        else:
            mask_kernel = torch.zeros((1, coors.shape[1], q_centers.shape[0])).to(
                output.device
            )

        outputs_class, outputs_mask, attn_mask = self.forward_prediction_heads(
            output, mask_features, mask_kernel
        )
        predictions_class.append(outputs_class)
        predictions_mask.append(outputs_mask)
        queries.append(output)

        for i in range(self.num_layers):
            level_index = i % self.num_feature_levels
            if attn_mask is not None:
                attn_mask[attn_mask.sum(-1) == attn_mask.shape[-1]] = False

            kernel = self.compute_kernel(q_centers, coors, size, angle)

            output = self.transformer_cross_attention_layers[i](
                output,
                kernel,
                src[level_index],
                attn_mask=attn_mask,
                padding_mask=None,
                pos=pos[level_index],
                query_pos=query_embed,
            )

            output = self.transformer_self_attention_layers[i](
                output, attn_mask=None, padding_mask=None, query_pos=query_embed
            )

            output = self.transformer_ffn_layers[i](output)

            if output.shape[-1] > self.num_queries:
                mask_kernel = self.compute_kernel(
                    q_centers, coors, size, angle, mask=True
                )
            else:
                mask_kernel = torch.zeros_like(outputs_mask)

            outputs_class, outputs_mask, attn_mask = self.forward_prediction_heads(
                output, mask_features, mask_kernel
            )
            predictions_class.append(outputs_class)
            predictions_mask.append(outputs_mask)
            queries.append(output)

        assert len(predictions_class) == self.num_layers + 1
        out = {
            "pred_logits": predictions_class[-1],
            "pred_masks": predictions_mask[-1],
            "queries": queries[-1],
        }

        out["aux_outputs"] = self.set_aux(predictions_class, predictions_mask, queries)

        return out

    def forward_prediction_heads(
        self,
        output,
        mask_features,
        kernel,
    ):
        decoder_output = self.decoder_norm(output)  # Layer norm
        outputs_class = self.class_embed(decoder_output)  # Linear

        mask_embed = self.mask_embed(decoder_output)  # MLP

        outputs_mask = torch.einsum("bqc,bpc->bpq", mask_embed, mask_features)

        attn_mask = outputs_mask.sigmoid() + kernel
        attn_mask = attn_mask - attn_mask.min()
        attn_mask = attn_mask / attn_mask.max()
        attn_mask = (attn_mask < 0.5).detach().bool()

        # Create binary mask
        attn_mask = (
            attn_mask.unsqueeze(1)
            .repeat(1, self.nheads, 1, 1)
            .flatten(0, 1)
            .permute(0, 2, 1)
        )

        return outputs_class, outputs_mask, attn_mask

    @torch.jit.unused
    def set_aux(self, outputs_class, outputs_seg_masks, queries):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [
            {"pred_logits": a, "pred_masks": b, "queries": c}
            for a, b, c in zip(outputs_class[:-1], outputs_seg_masks[:-1], queries[:-1])
        ]

    def compute_kernel(self, q_centers, coors, size, angle, mask=False):
        dx = torch.cdist(
            q_centers[self.num_queries :][:, 0].unsqueeze(1), coors[:, :, 0].T
        )
        dy = torch.cdist(
            q_centers[self.num_queries :][:, 1].unsqueeze(1), coors[:, :, 1].T
        )

        sx = size[self.num_queries :, 0]
        sy = size[self.num_queries :, 1]
        angles = angle[self.num_queries :]

        x_mu = torch.cat((dx.unsqueeze(-1), dy.unsqueeze(-1)), dim=-1)
        covs = misc.getcovs(sx, sy, angles)
        inv_covs = torch.linalg.inv(covs)
        k_weights = -0.5 * torch.einsum("bnd,bdd,bnd->bn", x_mu, inv_covs, x_mu)
        if mask:
            k_weights = torch.exp(k_weights)
        kernel = torch.zeros((1, coors.shape[1], q_centers.shape[0])).to(
            q_centers.device
        )
        kernel[0, :, self.num_queries :] = k_weights.T
        return kernel

mask_4d/models/decoder.py
PRBonn/Mask4D
# Modified by Rodrigo Marcuzzi from https://github.com/facebookresearch/Mask2Former

from itertools import filterfalse

import torch
import torch.nn.functional as F
from mask_4d.utils.misc import pad_stack
from torch import nn
from torch.autograd import Variable


class MaskLoss(nn.Module):
    """The loss computation happens in two steps:
    1) we compute hungarian assignment between ground truth boxes and the outputs of the model
    2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
    """

    def __init__(self, cfg, data_cfg):
        """
        Parameters:
            num_classes: number of object categories
            weight_dict: dict containing as key the names of the losses and as values their relative weight.
            eos_coef: relative classification weight applied to the no-object category
            losses: list of all the losses to be applied. See get_loss for list of available losses.
            num_points: number of points to sample to compute the loss
        """
        super().__init__()
        self.num_classes = data_cfg.NUM_CLASSES
        self.ignore = data_cfg.IGNORE_LABEL

        self.weight_dict = {
            cfg.LOSS_WEIGHTS_KEYS[i]: cfg.LOSS_WEIGHTS[i]
            for i in range(len(cfg.LOSS_WEIGHTS))
        }

        self.eos_coef = cfg.EOS_COEF

        weights = torch.ones(self.num_classes + 1)
        weights[0] = 0.0
        weights[-1] = self.eos_coef
        self.weights = weights

        self.num_points = cfg.NUM_POINTS

    def forward(self, outputs, targets):
        pass

    def get_losses(
        self, outputs, targets, indices, num_masks, step="detect", skip=None
    ):
        classes = self.loss_classes(outputs, targets, indices, step, skip=skip)
        masks = self.loss_masks(outputs, targets, indices, num_masks, step)
        classes.update(masks)
        return classes

    def loss_classes(self, outputs, targets, indices, step="detect", skip=None):
        """Classification loss (NLL)
        targets dicts must contain the key "classes" containing a tensor of dim [nb_target_boxes]
        """
        assert "pred_logits" in outputs
        pred_logits = outputs["pred_logits"].float()

        idx = self._get_pred_permutation_idx(indices)

        target_classes_o = torch.cat(
            [t[J] for t, (_, J) in zip(targets["classes"], indices)]
        ).to(pred_logits.device)

        target_classes = torch.full(
            pred_logits.shape[:2],
            self.num_classes,  # fill with class no_obj
            dtype=torch.int64,
            device=pred_logits.device,
        )
        target_classes[idx] = target_classes_o

        if skip is not None:
            # assign class 0 to ignore it when computing the loss
            target_classes[0][indices[0][0][skip]] = 0

        loss_ce = (
            F.cross_entropy(
                pred_logits.transpose(1, 2),
                target_classes,
                self.weights.to(pred_logits),
                ignore_index=self.ignore,
            )
            * self.weight_dict["loss_ce"]
        )
        losses = {"loss_ce" + step: loss_ce}
        return losses

    def loss_masks(self, outputs, targets, indices, num_masks, step="detect"):
        """Compute the losses related to the masks: the focal loss and the dice loss."""
        assert "pred_masks" in outputs

        masks = [t for t in targets["masks"]]
        n_masks = [m.shape[0] for m in masks]
        min_pts = min([m.shape[1] for m in masks])
        target_masks = pad_stack(masks)

        pred_idx = self._get_pred_permutation_idx(indices)
        tgt_idx = self._get_tgt_permutation_idx(indices, n_masks)
        pred_masks = outputs["pred_masks"]
        pred_masks = pred_masks[pred_idx[0], :, pred_idx[1]]
        target_masks = target_masks.to(pred_masks)
        target_masks = target_masks[tgt_idx]

        with torch.no_grad():
            # Random sampling num_points
            pt_idx = torch.randint(0, min_pts, (self.num_points, 1)).squeeze(1)
            point_labels = target_masks[:, pt_idx]

        point_logits = pred_masks[:, pt_idx]

        del pred_masks
        del target_masks

        losses = {
            "loss_mask"
            + step: sigmoid_ce_loss_jit(point_logits, point_labels, num_masks)
            * self.weight_dict["loss_mask"],
            "loss_dice"
            + step: dice_loss_jit(point_logits, point_labels, num_masks)
            * self.weight_dict["loss_dice"],
        }

        return losses

    def _get_pred_permutation_idx(self, indices):
        # permute predictions following indices
        batch_idx = torch.cat(
            [torch.full_like(src, i) for i, (src, _) in enumerate(indices)]
        )
        src_idx = torch.cat([src for (src, _) in indices])
        return batch_idx, src_idx

    def _get_tgt_permutation_idx(self, indices, n_masks):
        # permute targets following indices
        batch_idx = torch.cat(
            [torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)]
        )
        tgt_idx = torch.cat([tgt for (_, tgt) in indices])
        # From [B,id] to [id] of stacked masks
        cont_id = torch.cat(
            [torch.arange(n, device=tgt_idx.device) for n in n_masks],
        )
        b_id = torch.stack((batch_idx, cont_id), axis=1)
        map_m = torch.zeros((torch.max(batch_idx) + 1, max(n_masks)))
        for i in range(len(b_id)):
            map_m[b_id[i, 0], b_id[i, 1]] = i
        stack_ids = torch.tensor(
            [int(map_m[batch_idx[i], tgt_idx[i]]) for i in range(len(batch_idx))],
            device=tgt_idx.device,
        )
        return stack_ids


def dice_loss(inputs: torch.Tensor, targets: torch.Tensor, num_masks: float):
    """
    Compute the DICE loss, similar to generalized IOU for masks
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    """
    inputs = inputs.sigmoid()
    numerator = 2 * (inputs * targets).sum(-1)
    denominator = inputs.sum(-1) + targets.sum(-1)
    # loss = 1 - (numerator + 0.001) / (denominator + 0.001)
    loss = 1 - (numerator + 1) / (denominator + 1)
    return loss.sum() / num_masks


dice_loss_jit = torch.jit.script(dice_loss)  # type: torch.jit.ScriptModule


def sigmoid_ce_loss(inputs: torch.Tensor, targets: torch.Tensor, num_masks: float):
    """
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    Returns:
        Loss tensor
    """
    loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
    return loss.mean(1).sum() / num_masks


sigmoid_ce_loss_jit = torch.jit.script(sigmoid_ce_loss)  # type: torch.jit.ScriptModule


class SemLoss(nn.Module):
    def __init__(self, w):
        super().__init__()

        self.ce_w, self.lov_w = w

        self.cross_entropy = torch.nn.CrossEntropyLoss(ignore_index=0)

    def forward(self, outputs, targets, n):
        ce = self.cross_entropy(outputs, targets)
        lovasz = self.lovasz_softmax(F.softmax(outputs, dim=1), targets)
        loss = {"sem_ce_" + n: self.ce_w * ce, "sem_lov_" + n: self.lov_w * lovasz}
        return loss

    def lovasz_grad(self, gt_sorted):
        """
        Computes gradient of the Lovasz extension w.r.t sorted errors
        See Alg. 1 in paper
        """
        p = len(gt_sorted)
        gts = gt_sorted.sum()
        intersection = gts - gt_sorted.float().cumsum(0)
        union = gts + (1 - gt_sorted).float().cumsum(0)
        jaccard = 1.0 - intersection / union
        if p > 1:  # cover 1-pixel case
            jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
        return jaccard

    def lovasz_softmax(self, probas, labels, classes="present", ignore=None):
        """
        Multi-class Lovasz-Softmax loss
          probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1).
                  Interpreted as binary (sigmoid) output with outputs of size [B, H, W].
          labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1)
          classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
          per_image: compute the loss per image instead of per batch
          ignore: void class labels
        """
        loss = self.lovasz_softmax_flat(
            *self.flatten_probas(probas, labels, ignore), classes=classes
        )
        return loss

    def lovasz_softmax_flat(self, probas, labels, classes="present"):
        """
        Multi-class Lovasz-Softmax loss
          probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1)
          labels: [P] Tensor, ground truth labels (between 0 and C - 1)
          classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
        """
        if probas.numel() == 0:
            # only void pixels, the gradients should be 0
            return probas * 0.0
        C = probas.size(1)
        losses = []
        class_to_sum = list(range(C)) if classes in ["all", "present"] else classes
        for c in class_to_sum:
            fg = (labels == c).float()  # foreground for class c
            if classes == "present" and fg.sum() == 0:
                continue
            if C == 1:
                if len(classes) > 1:
                    raise ValueError("Sigmoid output possible only with 1 class")
                class_pred = probas[:, 0]
            else:
                class_pred = probas[:, c]
            errors = (Variable(fg) - class_pred).abs()
            errors_sorted, perm = torch.sort(errors, 0, descending=True)
            perm = perm.data
            fg_sorted = fg[perm]
            losses.append(
                torch.dot(errors_sorted, Variable(self.lovasz_grad(fg_sorted)))
            )
        return self.mean(losses)

    def flatten_probas(self, probas, labels, ignore=None):
        """
        Flattens predictions in the batch
        """
        # Probabilities from SparseTensor.features already flattened
        N, C = probas.size()
        probas = probas.contiguous().view(-1, C)
        labels = labels.view(-1)
        if ignore is None:
            return probas, labels
        valid = labels != ignore
        vprobas = probas[torch.nonzero(valid).squeeze()]
        vlabels = labels[valid]
        return vprobas, vlabels

    def isnan(self, x):
        return x != x

    def mean(self, l, ignore_nan=False, empty=0):
        """
        nanmean compatible with generators.
        """
        l = iter(l)
        if ignore_nan:
            l = filterfalse(self.isnan, l)
        try:
            n = 1
            acc = next(l)
        except StopIteration:
            if empty == "raise":
                raise ValueError("Empty mean")
            return empty
        for n, v in enumerate(l, 2):
            acc += v
        if n == 1:
            return acc
        return acc / n


class WrongLoss(nn.Module):
    def __init__(self, cfg, data_cfg):
        """
        Loss to account for correctly/incorrectly matched instances
        """
        super().__init__()
        self.weight_dict = {
            cfg.LOSS_WEIGHTS_KEYS[i]: cfg.LOSS_WEIGHTS[i]
            for i in range(len(cfg.LOSS_WEIGHTS))
        }

    def forward(self, track_ins, pred_idx, tgt_idx, targets, outputs, scan_i):
        """
        Arguments:
            track_ins: list of tracked instances
            pred_idx: matched predictions
            tgt_idx: matched ground truth
            targets: ground truth
            outputs:
            scan_i: number of the scan in the sequence
        Detection queries:
            * Matched GT is not a tracked instance, compute the loss.
            * Matched GT is a tracked instance, force to predict no_obj

        Tracking queries:
            * Correct instance matched, compute loss.
            * Incorrect instance matched, don't compute class loss and
              compute disimilarity loss
        """
        losses = {}
        tr_ids = track_ins.id[pred_idx]
        tgt_ids = targets["id"][0][tgt_idx]
        tracking_ids = torch.unique(track_ins.id)[1:]

        tracked = torch.isin(tgt_ids, tracking_ids)
        new_det = tr_ids == -1
        det_wrong_match = new_det * tracked
        det_wrong_idx = torch.nonzero(det_wrong_match).T[0]

        tr_q = tr_ids != -1
        unmatched_ids = tr_ids != tgt_ids
        tr_wrong_match = tr_q * unmatched_ids
        skip = tr_wrong_match
        tr_wrong_idx = torch.nonzero(tr_wrong_match).T[0]

        idx_delete = torch.cat((det_wrong_idx, tr_wrong_idx))

        if tr_wrong_match.sum() > 1:
            tr_wrong_pred_mask = outputs["pred_masks"][0][:, pred_idx[tr_wrong_match]]
            tr_wrong_tgt_mask = targets["masks"][0][tgt_idx[tr_wrong_match]]
            loss_wrong_mask = self.get_loss(
                tr_wrong_pred_mask, tr_wrong_tgt_mask, str(scan_i)
            )
            losses.update(loss_wrong_mask)
        return losses, idx_delete, skip

    def get_loss(self, pred_masks, tgt_masks, scan):
        """
        Compute the losses related to the masks: the focal loss and the dice loss.
        The losses are computed for the wrongly matched pairs (tracked_ins,GT)
        The losses only consider the points which lay inside the GT masks and force
        those points in the predicted mask to have low logits.
        """
        # points inside GT masks should be 0
        gt_neg = tgt_masks - 1
        gt_neg[gt_neg < 0] = 1
        # tgt_masks will filter out the GT background points
        # from the loss computation

        loss_ce = (
            wrong_sigmoid_ce_loss_jit(pred_masks.T, gt_neg, tgt_masks.bool())
            * self.weight_dict["loss_wrong_mask"]
        )

        loss_dice = (
            wrong_dice_loss_jit(pred_masks.T, gt_neg, tgt_masks.bool())
            * self.weight_dict["loss_wrong_dice"]
        )

        losses = {
            "loss_wrong_mask_" + scan: loss_ce,
            "loss_wrong_dice_" + scan: loss_dice,
        }

        return losses


def wrong_dice_loss(inputs: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor):
    """
    Compute the DICE loss, similar to generalized IOU for masks
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    """
    inputs = inputs.sigmoid()
    numerator = 2 * (inputs[mask] * targets[mask]).sum(-1)
    denominator = inputs[mask].sum(-1) + targets[mask].sum(-1)
    # loss = 1 - (numerator + 0.001) / (denominator + 0.001)
    loss = 1 - (numerator + 1) / (denominator + 1)
    return loss


wrong_dice_loss_jit = torch.jit.script(wrong_dice_loss)  # type: torch.jit.ScriptModule


def wrong_sigmoid_ce_loss(
    inputs: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor
):
    """
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    Returns:
        Loss tensor
    """
    num_masks = len(mask)
    loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")

    # loss = (loss * mask).mean(1).sum() / num_masks
    loss = loss[mask].mean() / num_masks
    return loss


wrong_sigmoid_ce_loss_jit = torch.jit.script(
    wrong_sigmoid_ce_loss
)  # type: torch.jit.ScriptModule

mask_4d/models/loss.py
PRBonn/Mask4D
import mask_4d.utils.misc as misc
import mask_4d.utils.testing as testing
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from mask_4d.models.backbone import SphericalEncoderDecoder
from mask_4d.models.decoder import MaskedTransformerDecoder
from mask_4d.models.loss import MaskLoss, SemLoss, WrongLoss
from mask_4d.models.matcher import HungarianMatcher
from mask_4d.models.positional_encoder import PositionalEncoder
from mask_4d.utils.evaluate_4dpanoptic import PanopticKitti4DEvaluator
from mask_4d.utils.evaluate_panoptic import PanopticKittiEvaluator
from mask_4d.utils.instances import Tracks
from mask_4d.utils.kalman_filter import KalmanBoxTracker
from pytorch_lightning.core.lightning import LightningModule


class Mask4D(LightningModule):
    def __init__(self, hparams):
        super().__init__()
        self.save_hyperparameters(dict(hparams))
        self.cfg = hparams
        self.n_q = hparams.DECODER.NUM_QUERIES

        self.backbone = SphericalEncoderDecoder(hparams.BACKBONE, hparams.KITTI)
        self.decoder = MaskedTransformerDecoder(
            hparams.DECODER, hparams.BACKBONE, hparams.KITTI
        )

        self.mask_loss = MaskLoss(hparams.LOSS, hparams.KITTI)
        self.wrong_loss = WrongLoss(hparams.LOSS, hparams.KITTI)
        self.sem_loss = SemLoss(hparams.LOSS.SEM.WEIGHTS)

        self.evaluator = PanopticKittiEvaluator(hparams.KITTI)
        self.evaluator4d = PanopticKitti4DEvaluator(hparams.KITTI)

        self.matcher = HungarianMatcher(
            hparams.LOSS.LOSS_WEIGHTS, hparams.LOSS.NUM_POINTS
        )
        self.query_feat = nn.Embedding(
            self.n_q, hparams.DECODER.HIDDEN_DIM, device="cuda"
        )
        self.query_embed = nn.Embedding(
            self.n_q, hparams.DECODER.HIDDEN_DIM, device="cuda"
        )
        self.track_ins = self.init_tracks()
        self.last_ins_id = 1
        self.last_pose = np.eye(4)
        self.trackers = {}
        self.min_mask_pts = hparams.KITTI.MIN_POINTS

    def forward(self, x, track_ins):
        feats, coors, logits = self.backbone(x)
        outputs = self.decoder(feats, coors, track_ins)
        return outputs, logits

    def training_step(self, x: dict, idx):
        track_ins = self.init_tracks()
        losses = {}
        for s_i in range(len(x["pt_coord"])):
            scan = {k: [v[s_i]] for k, v in x.items()}
            outputs, logits = self.forward(scan, track_ins)
            track_ins.query = outputs["queries"][0]
            targets = {
                "classes": scan["masks_cls"],
                "masks": scan["masks"],
                "id": [scan["masks_ids"][0]],
            }
            matched_ins = []
            if (
                s_i != 0
                and len(track_ins) > self.n_q
                and (targets["id"][0] != 0).sum() > 0
            ):
                # Associate detections with previously tracked instances
                track_loss, matched_ins = self.tracking_step(
                    outputs, targets, track_ins, s_i
                )
                losses.update(track_loss)

            # Detect stuff and new instances
            pred_idx, tgt_idx, det_losses = self.detection_step(
                outputs, targets, matched_ins, track_ins, s_i
            )
            losses.update(det_losses)

            # aux losses
            losses.update(self.auxLoss(scan, logits, str(s_i)))

            track_ins = self.update_tr_instances(
                scan, outputs, pred_idx, tgt_idx, targets, track_ins
            )

            if s_i + 1 < len(x["pt_coord"]):
                track_ins.center = self.update_centers_ego(x, s_i, track_ins)

            new_track_ins = self.init_tracks()
            track_ins = Tracks.cat([new_track_ins, track_ins])

        # log losses
        loss_detect = 0
        loss_track = 0
        loss_wrong = 0
        for k, v in losses.items():
            if "detect" in k:
                loss_detect += v
            if "track" in k:
                loss_track += v
            if "wrong" in k:
                loss_wrong += v
        self.log(f"losses/detect", loss_detect, batch_size=self.cfg.TRAIN.BATCH_SIZE)
        self.log(f"losses/track", loss_track, batch_size=self.cfg.TRAIN.BATCH_SIZE)
        self.log(f"losses/wrong", loss_wrong, batch_size=self.cfg.TRAIN.BATCH_SIZE)

        total_loss = sum(losses.values())
        self.log("train_loss", total_loss, batch_size=self.cfg.TRAIN.BATCH_SIZE)
        return total_loss

    def validation_step(self, x: dict, idx):
        self.evaluation_step(x, idx)

    def evaluation_step(self, x: dict, idx):
        sem_pred, ins_pred = self.panoptic4d_inference(x)
        self.evaluator.update(sem_pred, ins_pred, x)
        self.evaluator4d.update(sem_pred, ins_pred, x)

    def validation_epoch_end(self, outputs):
        self.log(
            "metrics/pq",
            self.evaluator.get_mean_pq(),
            batch_size=self.cfg.TRAIN.BATCH_SIZE,
        )
        self.log(
            "metrics/iou",
            self.evaluator.get_mean_iou(),
            batch_size=self.cfg.TRAIN.BATCH_SIZE,
        )
        self.log(
            "metrics/rq",
            self.evaluator.get_mean_rq(),
            batch_size=self.cfg.TRAIN.BATCH_SIZE,
        )
        self.evaluator4d.calculate_metrics()
        self.log(
            "metrics/aq",
            self.evaluator4d.get_mean_aq(),
            batch_size=self.cfg.TRAIN.BATCH_SIZE,
        )
        self.log(
            "metrics/pq4d",
            self.evaluator4d.get_mean_pq4d(),
            batch_size=self.cfg.TRAIN.BATCH_SIZE,
        )

        # reset tracking variables
        self.last_ins_id = 1
        self.track_ins = self.init_tracks()

        if not "EVALUATE" in self.cfg:
            self.evaluator.reset()
            self.evaluator4d.reset()

    def test_step(self, x: dict, idx):
        _sem_pred, _ins_pred = self.panoptic4d_inference(x)
        sem_pred, ins_pred = testing.full_preds(x, _sem_pred, _ins_pred)
        results_dir = self.cfg.RESULTS_DIR
        class_inv_lut = self.evaluator.get_class_inv_lut()
        testing.save_results(sem_pred, ins_pred, results_dir, x, class_inv_lut)

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.cfg.TRAIN.LR)
        scheduler = torch.optim.lr_scheduler.StepLR(
            optimizer, step_size=self.cfg.TRAIN.STEP, gamma=self.cfg.TRAIN.DECAY
        )
        return [optimizer], [scheduler]

    def panoptic4d_inference(self, x):
        sem_pred = []
        ins_pred = []
        for i in range(len(x["pt_coord"])):
            if int(x["fname"][i].split("/")[-1][:-4]) == 0:
                self.track_ins = self.init_tracks()
                self.last_ins_id = 1
                self.last_pose = x["pose"][i]
            elif len(self.track_ins) > self.n_q:
                self.update_centers_motion(x, i)

            scan = {k: [v[i]] for k, v in x.items()}
            outputs, logits = self.forward(scan, self.track_ins)
            self.track_ins.query = outputs["queries"][0]

            sem_pred.append(self.semantic_predictions(outputs))
            ins_pred.append(self.instance_predictions(outputs, scan))

            self.update_motion_model()

            new_track_ins = self.init_tracks()
            self.track_ins = Tracks.cat([new_track_ins, self.track_ins])
        return sem_pred, ins_pred

    def semantic_predictions(self, outputs):
        mask_cls = outputs["pred_logits"][0]
        mask_pred = outputs["pred_masks"][0]
        mask_cls = F.softmax(mask_cls, dim=-1)[..., :-1]
        mask_pred = mask_pred.sigmoid()
        pred = torch.einsum("qc,pq->pc", mask_cls, mask_pred)
        semseg = torch.argmax(pred, dim=1)
        return semseg.cpu().numpy()

    def instance_predictions(self, outputs, x):
        points = torch.from_numpy(x["pt_coord"][0]).to(outputs["pred_masks"])
        pose = x["pose"][0]
        glob_points = misc.apply_pose(points, pose)
        ins_pred = torch.zeros(points.shape[0]).long()
        mask_cls = outputs["pred_logits"][0]
        mask_pred = outputs["pred_masks"][0].sigmoid()
        things_ids = torch.tensor(self.trainer.datamodule.things_ids).cuda()
        num_classes = self.cfg.KITTI.NUM_CLASSES

        scores, labels = mask_cls.max(-1)
        keep = labels.ne(num_classes) & torch.isin(labels, things_ids)

        if keep.sum() == 0:  # no detected instances
            self.track_ins = self.track_ins[self.track_ins.id != -1]
            self.track_ins.life -= 1
            self.track_ins = self.track_ins[self.track_ins.life > 0]
            return ins_pred.cpu().numpy()

        cur_scores = scores[keep]
        cur_classes = labels[keep]
        cur_masks = mask_pred[:, keep]
        cur_prob_masks = cur_scores.unsqueeze(0) * cur_masks
        cur_mask_ids = cur_prob_masks.argmax(1)
        cur_ins_ids = self.track_ins.id[keep]
        ins_idx = torch.where(keep)[0]

        for k in range(cur_classes.shape[0]):
            mask_area = (cur_mask_ids == k).sum().item()  # points in mask k
            original_area = (cur_masks[:, k] >= 0.5).sum().item()  # binary mask
            mask = (cur_mask_ids == k) & (cur_masks[:, k] >= 0.5)
            ins_id = cur_ins_ids[k]

            if mask_area > 0 and original_area > 0 and mask.sum().item() > 0:
                if mask_area / original_area < self.cfg.MODEL.OVERLAP_THRESHOLD:
                    continue
                if mask.sum() < self.min_mask_pts:
                    continue

                self.track_ins.life[ins_idx[k]] += 1
                pts = misc.remove_outliers(points[mask])
                if pts.shape[0] < self.min_mask_pts:
                    continue
                g_pts = misc.remove_outliers(glob_points[mask])
                size, angle = misc.fit_ellipse(pts)
                self.track_ins.size_xy[ins_idx[k]] = size
                self.track_ins.angle[ins_idx[k]] = angle
                self.track_ins.bbox[ins_idx[k]] = misc.bbox_from_pts(g_pts)
                self.track_ins.center[ins_idx[k]] = misc.center_from_points(pts)

                if ins_id == -1:  # if new detected instance
                    self.track_ins.id[ins_idx[k]] = self.last_ins_id
                    ins_id = self.last_ins_id
                    self.last_ins_id += 1
                # assign instance ID to points
                ins_pred[mask == 1] = ins_id
        self.track_ins = self.track_ins[self.track_ins.id != -1]
        self.track_ins.life -= 1
        self.track_ins = self.track_ins[self.track_ins.life > 0]

        return ins_pred.cpu().numpy()

    def init_tracks(self):
        """
        Init list of N tracks with the following params:
            query: [N,C] query feature to decode the mask
            query_pe: [N,C] query positional embedding
            id: [N,1] track identifier
            life: [N,1] for how many steps keep the track alive
            center: [N,3] center of the mask
            size_xy: [N,2,3] size of the instance in X and Y
            angle: [N,1] orientation of the instance (theta)
            bbox: [N,7]: [x,y,z,theta,l,w,h] 3D bounding box around the instance
        """
        tracks = Tracks()
        N = self.query_feat.weight.shape[0]
        dev = self.query_feat.weight.device
        tracks.query = self.query_feat.weight
        tracks.query_pe = self.query_embed.weight
        tracks.id = torch.full((N,), -1, device=dev)
        tracks.life = torch.full((N,), self.cfg.MODEL.TR_LIFE, device=dev)
        tracks.center = torch.full((N, 3), 0.0, device=dev)
        tracks.size_xy = torch.full((N, 3), 0.0, device=dev)
        tracks.angle = torch.full((N,), 0.0, device=dev)
        tracks.bbox = torch.full((N, 7), 0.0, device=dev)
        return tracks

    def detection_step(self, outputs, targets, matched_ins, track_ins, scan_i):
        losses = {}
        indices = self.matcher(outputs, targets)

        pred_idx = self.mask_loss._get_pred_permutation_idx(indices)[1]
        n_masks = [m.shape[0] for m in targets["masks"]]
        tgt_idx = self.mask_loss._get_tgt_permutation_idx(indices, n_masks)

        skip = None
        if len(matched_ins) > 0:
            loss_wrong, idx_delete, skip = self.wrong_loss(
                track_ins, pred_idx, tgt_idx, targets, outputs, scan_i
            )
            losses.update(loss_wrong)
            # delete from predictions and targets
            pred_idx = misc.th_delete(pred_idx, idx_delete)
            tgt_idx = misc.th_delete(tgt_idx, idx_delete)

        # compute loss on pairs matched with Hungarian
        det_loss = self.mask_loss.get_losses(
            outputs,
            targets,
            indices,
            len(indices[0][0]),
            "_detect_" + str(scan_i),
            skip=skip,
        )
        det_loss = {k: self.cfg.LOSS.DET_W * v for k, v in det_loss.items()}
        losses.update(det_loss)
        losses.update(self.aux_outs_loss(outputs, targets, indices, scan_i))
        return pred_idx, tgt_idx, losses

    def tracking_step(self, outputs, targets, track_ins, scan_i):
        matched_ins, matched_outputs, matched_tgt = self.fixed_match(
            track_ins, outputs, targets
        )

        track_loss = {}
        if len(matched_ins) > 0:
            # create indices for the loss function
            idx = torch.arange(len(matched_ins), device=matched_ins.query.device)
            matched_idx = [(idx, idx)]
            # compute loss on fixed matched pairs
            track_loss = self.mask_loss.get_losses(
                matched_outputs,
                matched_tgt,
                matched_idx,
                len(matched_ins),
                "_track_" + str(scan_i),
            )
            track_loss = {k: self.cfg.LOSS.TRACK_W * v for k, v in track_loss.items()}
        return track_loss, matched_ins

    def fixed_match(self, track_ins, outputs, targets):
        """
        Fixed match between tracked instances and targets based on ID.
        Return matched and non-matched instances, outputs and targets
        """
        m_ins_ids, m_tgt_iids = torch.where(track_ins.id[:, None] == targets["id"][0])
        m_out_log = outputs["pred_logits"][:, m_ins_ids, :]
        m_out_masks = outputs["pred_masks"][:, :, m_ins_ids]
        m_tgt_classes = targets["classes"][0][m_tgt_iids]
        m_tgt_masks = targets["masks"][0][m_tgt_iids]
        m_tgt_ids = targets["id"][0][m_tgt_iids]
        if len(m_ins_ids) == 0:
            return [], [], [], [], []

        matched_outputs = {"pred_logits": m_out_log, "pred_masks": m_out_masks}
        matched_targets = {
            "classes": [m_tgt_classes],
            "masks": [m_tgt_masks],
            "id": [m_tgt_ids],
        }
        matched_ins = track_ins[m_ins_ids]

        return matched_ins, matched_outputs, matched_targets

    def update_tr_instances(self, scan, outputs, pred_idx, tgt_idx, targets, track_ins):
        """
        Update tracked instances parameters: angle, size, center and id.
        The ids are according to associations made by the Hungarian matcher.
        """
        for p in range(len(pred_idx)):
            if targets["id"][0][tgt_idx[p]] > 0:  # ignore stuff
                coors = torch.from_numpy(scan["pt_coord"][0]).cuda()
                _mask = (outputs["pred_masks"][0, :, pred_idx[p]]).sigmoid() > 0.5
                pts = coors[_mask]
                if pts.shape[0] <= self.min_mask_pts:
                    continue
                track_ins.id[pred_idx[p]] = targets["id"][0][tgt_idx[p]]
                size, angle = misc.fit_ellipse(pts)
                track_ins.size_xy[pred_idx[p]] = size
                track_ins.angle[pred_idx[p]] = angle
                track_ins.center[pred_idx[p]] = misc.center_from_points(pts)
        # remove track_ins with id=-1 (out queries not matched with GT instance)
        track_ins = track_ins[track_ins.id != -1]
        return track_ins

    def update_centers_ego(self, x, scan_i, track_ins):
        """
        Update instance centers compensating the ego motion
        """
        cent = track_ins.center
        p0 = x["pose"][scan_i]
        p1 = x["pose"][scan_i + 1]
        t0_1 = np.linalg.inv(p1) @ p0
        center = misc.apply_pose(cent, t0_1)
        return center

    def update_centers_motion(self, x, i):
        # update centers using the motion model
        inv_p1 = np.linalg.inv(x["pose"][i])
        for tr in self.trackers:
            self.trackers[tr].predict()
        for j in range(len(self.track_ins.center[self.n_q :])):
            id = self.track_ins.id[self.n_q + j].item()
            pred_center = torch.from_numpy(
                self.trackers[id].get_state()[:3][None, :]
            ).cuda()
            # transform from global to current frame
            local_center = misc.apply_pose(pred_center, inv_p1)[0]
            self.track_ins.center[self.n_q + j] = local_center

    def update_motion_model(self):
        for id, bbox in zip(self.track_ins.id, self.track_ins.bbox):
            if id.item() not in self.trackers.keys():
                self.trackers[id.item()] = KalmanBoxTracker(bbox.cpu())
            else:
                self.trackers[id.item()].update(bbox.cpu())

    def aux_outs_loss(self, outputs, targets, indices, scan_i):
        all_losses = {}
        for aux_i, aux_outputs in enumerate(outputs["aux_outputs"]):
            # detection aux losses
            aux_indices = self.matcher(aux_outputs, targets)
            l_dict = self.mask_loss.get_losses(
                aux_outputs,
                targets,
                aux_indices,
                len(indices[0][0]),
                "_detect_" + str(scan_i),
            )
            l_dict = {
                f"{aux_i}_" + k: self.cfg.LOSS.DET_W_AUX * v for k, v in l_dict.items()
            }
            all_losses.update(l_dict)
        return all_losses

    def auxLoss(self, x, bb_log, n):
        loss = {}
        labels = torch.from_numpy(x["sem_label"][0]).type(torch.LongTensor).cuda()[:, 0]
        loss_sem_bb = self.sem_loss(bb_log, labels, n)
        loss.update(loss_sem_bb)
        return loss

mask_4d/models/mask_model.py
PRBonn/Mask4D
# Modified by Rodrigo Marcuzzi from https://github.com/facebookresearch/Mask2Former
"""
Modules to compute the matching cost and solve the corresponding LSAP.
"""
import torch
import torch.nn.functional as F
from scipy.optimize import linear_sum_assignment as lsa
from torch import nn
from torch.cuda.amp import autocast


class HungarianMatcher(nn.Module):
    """This class computes an assignment between the targets and the predictions of the network

    For efficiency reasons, the targets don't include the no_object. Because of this, in general,
    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
    while the others are un-matched (and thus treated as non-objects).
    """

    def __init__(self, costs: list, num_points):
        """Creates the matcher

        Params:
            weight_class: This is the relative weight of the classification error in the matching cost
            weight_mask: This is the relative weight of the focal loss of the binary mask in the matching cost
            weight_dice: This is the relative weight of the dice loss of the binary mask in the matching cost
        """
        super().__init__()
        self.weight_class, self.weight_mask, self.weight_dice = costs[:3]

        assert (
            self.weight_class != 0 or self.weight_mask != 0 or self.weight_dice != 0
        ), "all costs cant be 0"

        self.num_points = num_points

    @torch.no_grad()
    def forward(self, outputs, targets):
        """Performs the matching

        Params:
            outputs: This is a dict that contains at least these entries:
                 "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
                 "pred_masks": Tensor of dim [batch_size, num_queries, H_pred, W_pred] with the predicted masks

            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
                 "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth
                           objects in the target) containing the class labels
                 "masks": Tensor of dim [num_target_boxes, H_gt, W_gt] containing the target masks

        Returns:
            A list of size batch_size, containing tuples of (index_i, index_j) where:
                - index_i is the indices of the selected predictions (in order)
                - index_j is the indices of the corresponding selected targets (in order)
            For each batch element, it holds:
                len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
        """
        return self.memory_efficient_forward(outputs, targets)

    @torch.no_grad()
    def memory_efficient_forward(self, outputs, targets):
        """More memory-friendly matching"""
        bs, num_queries, num_classes = outputs["pred_logits"].shape

        indices = []
        # Iterate through batch size
        for b in range(bs):
            out_prob = outputs["pred_logits"][b].softmax(-1)
            tgt_ids = targets["classes"][b].type(torch.int64)
            cost_class = -out_prob[:, tgt_ids]

            out_mask = outputs["pred_masks"][b].permute(1, 0)
            tgt_mask = targets["masks"][b].to(out_mask)
            n_pts_scan = tgt_mask.shape[1]
            pt_idx = torch.randint(0, n_pts_scan, (self.num_points, 1)).squeeze(1)
            tgt_mask = tgt_mask[:, pt_idx]
            out_mask = out_mask[:, pt_idx]

            with autocast(enabled=False):
                out_mask = out_mask.float()
                tgt_mask = tgt_mask.float()

                cost_mask = batch_sigmoid_ce_cost_jit(out_mask, tgt_mask)
                cost_dice = batch_dice_cost_jit(out_mask, tgt_mask)

            C = (
                self.weight_mask * cost_mask
                + self.weight_class * cost_class
                + self.weight_dice * cost_dice
            )
            C = C.reshape(num_queries, -1).cpu()
            indices.append(lsa(C))

        device = cost_mask.device
        return [
            (
                torch.as_tensor(i, dtype=torch.int64, device=device),
                torch.as_tensor(j, dtype=torch.int64, device=device),
            )
            for i, j in indices
        ]


def batch_dice_cost(inputs: torch.Tensor, targets: torch.Tensor):
    """
    Compute the DICE loss, similar to generalized IOU for masks
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    """
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1)
    numerator = 2 * torch.einsum("nc,mc->nm", inputs, targets)
    denominator = inputs.sum(-1)[:, None] + targets.sum(-1)[None, :]
    loss = 1 - (numerator + 1) / (denominator + 1)
    return loss


batch_dice_cost_jit = torch.jit.script(batch_dice_cost)  # type: torch.jit.ScriptModule


def batch_sigmoid_ce_cost(inputs: torch.Tensor, targets: torch.Tensor):
    """
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    Returns:
        Loss tensor
    """
    hw = inputs.shape[1]

    pos = F.binary_cross_entropy_with_logits(
        inputs, torch.ones_like(inputs), reduction="none"
    )
    neg = F.binary_cross_entropy_with_logits(
        inputs, torch.zeros_like(inputs), reduction="none"
    )

    loss = torch.einsum("nc,mc->nm", pos, targets) + torch.einsum(
        "nc,mc->nm", neg, (1 - targets)
    )

    return loss / hw


batch_sigmoid_ce_cost_jit = torch.jit.script(
    batch_sigmoid_ce_cost
)  # type: torch.jit.ScriptModule

mask_4d/models/matcher.py
PRBonn/Mask4D
import math

import torch
import torch.nn.functional as F
from torch import nn


class PositionAttention(nn.Module):
    """Position-aware Mask Attention."""

    def __init__(self, num_hiddens, num_heads, dropout, bias=True):
        super().__init__()
        self.num_heads = num_heads
        self.attention = DotProductAttention(dropout)
        self.in_proj_weight = nn.parameter.Parameter(
            torch.empty(3 * num_hiddens, num_hiddens)
        )
        self.in_proj_bias = nn.parameter.Parameter(torch.empty(3 * num_hiddens))
        self.out_proj = nn.Linear(num_hiddens, num_hiddens, bias=bias)
        self.hd = num_hiddens

    def forward(self, query, key, value, attn_mask, key_padding_mask, kernel):
        if key_padding_mask:
            attn_mask = attn_mask + key_padding_mask
        queries = self.transpose_qkv(
            F.linear(
                query, self.in_proj_weight[: self.hd, :], self.in_proj_bias[: self.hd]
            )
        )
        keys = self.transpose_qkv(
            F.linear(
                key,
                self.in_proj_weight[self.hd : -self.hd, :],
                self.in_proj_bias[self.hd : -self.hd],
            )
        )
        values = self.transpose_qkv(
            F.linear(
                value,
                self.in_proj_weight[-self.hd :, :],
                self.in_proj_bias[-self.hd :],
            )
        )
        output = self.attention(queries, keys, values, attn_mask, kernel)
        output_concat = self.transpose_output(output)
        return self.out_proj(output_concat)

    def transpose_qkv(self, X):
        """Transposition for parallel computation of multiple attention heads."""
        X = X.reshape(X.shape[0], X.shape[1], self.num_heads, -1)
        X = X.permute(0, 2, 1, 3)
        return X.reshape(-1, X.shape[2], X.shape[3])

    def transpose_output(self, X):
        """Reverse the operation of transpose_qkv."""
        X = X.reshape(-1, self.num_heads, X.shape[1], X.shape[2])
        X = X.permute(0, 2, 1, 3)
        return X.reshape(X.shape[0], X.shape[1], -1)


class DotProductAttention(nn.Module):
    def __init__(self, dropout):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

    def forward(self, queries, keys, values, mask=None, kernel=None):
        d = queries.shape[-1]
        scores = torch.bmm(queries, keys.transpose(1, 2)) / math.sqrt(d)
        scores = scores + kernel.permute(0, 2, 1)
        self.attention_weights = masked_softmax(scores, mask)
        return torch.bmm(self.dropout(self.attention_weights), values)


def masked_softmax(X, mask):
    """Perform softmax operation by masking elements on the last axis."""
    if mask is None:
        return nn.functional.softmax(X, dim=-1)
    else:
        mask = torch.zeros_like(mask, dtype=X.dtype).masked_fill_(mask, float("-inf"))
        X = X + mask
        return nn.functional.softmax(X, dim=-1)

mask_4d/models/position_attention.py
PRBonn/Mask4D
import math

import torch
import torch.nn as nn


class PositionalEncoder(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.max_freq = cfg.MAX_FREQ
        self.dimensionality = cfg.DIMENSIONALITY
        self.num_bands = math.floor(cfg.FEAT_SIZE / cfg.DIMENSIONALITY / 2)
        self.base = cfg.BASE
        pad = cfg.FEAT_SIZE - self.num_bands * 2 * cfg.DIMENSIONALITY
        self.zero_pad = nn.ZeroPad2d((pad, 0, 0, 0))  # left padding

    def forward(self, _x):
        """
        _x [B,N,3]: batched point coordinates
        returns: [B,N,C]: positional encoding of dimension C
        """
        x = _x.clone()
        x[:, :, 0] = x[:, :, 0] / 48
        x[:, :, 1] = x[:, :, 1] / 48
        x[:, :, 2] = x[:, :, 2] / 4
        x = x.unsqueeze(-1)
        scales = torch.logspace(
            0.0,
            math.log(self.max_freq / 2) / math.log(self.base),
            self.num_bands,
            base=self.base,
            device=x.device,
            dtype=x.dtype,
        )
        # reshaping
        scales = scales[(*((None,) * (len(x.shape) - 1)), Ellipsis)]
        x = x * scales * math.pi
        x = torch.cat([x.sin(), x.cos()], dim=-1)
        x = x.flatten(2)
        enc = self.zero_pad(x)
        return enc

mask_4d/models/positional_encoder.py
PRBonn/Mask4D
# From https://github.com/dvlab-research/SphereFormer/blob/master/model/spherical_transformer.py
from functools import partial

import numpy as np
import torch
import torch.nn as nn
from sptr import (SparseTrTensor, get_indices_params, sparse_self_attention,
                  to_3d_numpy)
from timm.models.layers import DropPath, trunc_normal_


class Mlp(nn.Module):
    """Multilayer perceptron."""

    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop, inplace=True)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


def cart2sphere(xyz):
    x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2]
    theta = (torch.atan2(y, x) + np.pi) * 180 / np.pi
    beta = torch.atan2(torch.sqrt(x**2 + y**2), z) * 180 / np.pi
    r = torch.sqrt(x**2 + y**2 + z**2)
    return torch.stack([theta, beta, r], -1)


def exponential_split(xyz, index_0, index_1, relative_position_index, a=0.05 * 0.25):
    """
    Mapping functioni from r to idx
    | r         ---> idx    |
    | ...       ---> ...    |
    | [-2a, a)  ---> -2     |
    | [-a, 0)   ---> -1     |
    | [0, a)    ---> 0      |
    | [a, 2a)   ---> 1      |
    | [2a, 4a)  ---> 2      |
    | [4a, 6a)  ---> 3      |
    | [6a, 10a) ---> 4      |
    | [10a, 14a)---> 5      |
    | ...       ---> ...    |
    Starting from 0, the split length will double once used twice.
    """

    r = xyz[:, 2]
    rel_pos = r[index_0.long()] - r[index_1.long()]  # [M,3]
    rel_pos_abs = rel_pos.abs()
    flag_float = (rel_pos >= 0).float()
    idx = 2 * torch.floor(torch.log((rel_pos_abs + 2 * a) / a) / np.log(2)) - 2
    idx = idx + ((3 * (2 ** (idx // 2)) - 2) * a <= rel_pos_abs).float()
    idx = idx * (2 * flag_float - 1) + (flag_float - 1)
    relative_position_index[:, 2] = idx.long() + 24
    return relative_position_index


# Modified From sptr.VarLengthMultiheadSA (https://github.com/dvlab-research/SparseTransformer)
class SparseMultiheadSASphereConcat(nn.Module):
    def __init__(
        self,
        embed_dim,
        num_heads,
        indice_key,
        window_size,
        window_size_sphere,
        shift_win=False,
        pe_type="none",
        dropout=0.0,
        qk_scale=None,
        qkv_bias=True,
        algo="native",
        **kwargs
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.indice_key = indice_key
        self.shift_win = shift_win
        self.pe_type = pe_type
        head_dim = embed_dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        self.window_size = to_3d_numpy(window_size)
        self.window_size_sphere = to_3d_numpy(window_size_sphere)

        if pe_type == "contextual":
            self.rel_query, self.rel_key, self.rel_value = (
                kwargs["rel_query"],
                kwargs["rel_key"],
                kwargs["rel_value"],
            )

            quant_size = kwargs["quant_size"]
            self.quant_size = to_3d_numpy(quant_size)

            quant_size_sphere = kwargs["quant_size_sphere"]
            self.quant_size_sphere = to_3d_numpy(quant_size_sphere)

            self.a = kwargs["a"]

            quant_grid_length = int((window_size[0] + 1e-4) / quant_size[0])
            assert int((window_size[0] + 1e-4) / quant_size[0]) == int(
                (window_size[1] + 1e-4) / quant_size[1]
            )

            # currently only support rel_query, rel_key and rel_value also equal to True
            assert self.rel_query and self.rel_key and self.rel_value
            num_heads_brc1 = num_heads // 2
            self.num_heads_brc1 = num_heads_brc1

            if self.rel_query:
                self.relative_pos_query_table = nn.Parameter(
                    torch.zeros(2 * quant_grid_length - 1, 3, num_heads_brc1, head_dim)
                )
                trunc_normal_(self.relative_pos_query_table, std=0.02)
            if self.rel_key:
                self.relative_pos_key_table = nn.Parameter(
                    torch.zeros(2 * quant_grid_length - 1, 3, num_heads_brc1, head_dim)
                )
                trunc_normal_(self.relative_pos_key_table, std=0.02)
            if self.rel_value:
                self.relative_pos_value_table = nn.Parameter(
                    torch.zeros(2 * quant_grid_length - 1, 3, num_heads_brc1, head_dim)
                )
                trunc_normal_(self.relative_pos_value_table, std=0.02)
            self.quant_grid_length = quant_grid_length

            quant_grid_length_sphere = int(
                (window_size_sphere[0] + 1e-4) / quant_size_sphere[0]
            )
            assert int((window_size_sphere[0] + 1e-4) / quant_size_sphere[0]) == int(
                (window_size_sphere[1] + 1e-4) / quant_size_sphere[1]
            )

            num_heads_brc2 = num_heads - num_heads_brc1
            if self.rel_query:
                self.relative_pos_query_table_sphere = nn.Parameter(
                    torch.zeros(
                        2 * quant_grid_length_sphere, 3, num_heads_brc2, head_dim
                    )
                )
                trunc_normal_(self.relative_pos_query_table_sphere, std=0.02)
            if self.rel_key:
                self.relative_pos_key_table_sphere = nn.Parameter(
                    torch.zeros(
                        2 * quant_grid_length_sphere, 3, num_heads_brc2, head_dim
                    )
                )
                trunc_normal_(self.relative_pos_key_table_sphere, std=0.02)
            if self.rel_value:
                self.relative_pos_value_table_sphere = nn.Parameter(
                    torch.zeros(
                        2 * quant_grid_length_sphere, 3, num_heads_brc2, head_dim
                    )
                )
                trunc_normal_(self.relative_pos_value_table_sphere, std=0.02)
            self.quant_grid_length_sphere = quant_grid_length_sphere

        self.qkv = nn.Linear(embed_dim, embed_dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(dropout, inplace=True)
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.proj_drop = nn.Dropout(dropout, inplace=True)

    def forward(self, sptr_tensor: SparseTrTensor):
        query, key, value = (
            sptr_tensor.query_feats,
            sptr_tensor.key_feats,
            sptr_tensor.value_feats,
        )
        assert key is None and value is None
        xyz = sptr_tensor.query_indices[:, 1:]
        batch = sptr_tensor.query_indices[:, 0]

        assert xyz.shape[1] == 3

        N, C = query.shape

        qkv = (
            self.qkv(query)
            .reshape(N, 3, self.num_heads, C // self.num_heads)
            .permute(1, 0, 2, 3)
            .contiguous()
        )
        query, key, value = qkv[0], qkv[1], qkv[2]  # [N, num_heads, C//num_heads]
        query = query * self.scale

        xyz_sphere = cart2sphere(xyz)
        index_params = sptr_tensor.find_indice_params(self.indice_key)
        if index_params is None:
            (
                index_0,
                index_0_offsets,
                n_max,
                index_1,
                index_1_offsets,
                sort_idx,
            ) = get_indices_params(xyz, batch, self.window_size, self.shift_win)
            (
                index_0_sphere,
                index_0_offsets_sphere,
                n_max_sphere,
                index_1_sphere,
                index_1_offsets_sphere,
                sort_idx_sphere,
            ) = get_indices_params(
                xyz_sphere, batch, self.window_size_sphere, self.shift_win
            )
            sptr_tensor.indice_dict[self.indice_key] = (
                index_0,
                index_0_offsets,
                n_max,
                index_1,
                index_1_offsets,
                sort_idx,
                self.window_size,
                self.shift_win,
                index_0_sphere,
                index_0_offsets_sphere,
                n_max_sphere,
                index_1_sphere,
                index_1_offsets_sphere,
                sort_idx_sphere,
            )
        else:
            (
                index_0,
                index_0_offsets,
                n_max,
                index_1,
                index_1_offsets,
                sort_idx,
                window_size,
                shift_win,
                index_0_sphere,
                index_0_offsets_sphere,
                n_max_sphere,
                index_1_sphere,
                index_1_offsets_sphere,
                sort_idx_sphere,
            ) = index_params
            assert (window_size == self.window_size) and (
                shift_win == self.shift_win
            ), "window_size and shift_win must be the same for sptr_tensors with the same indice_key: {}".format(
                self.indice_key
            )
            assert (
                window_size_sphere == self.window_size_sphere
            ), "window_size and shift_win must be the same for sptr_tensors with the same indice_key: {}".format(
                self.indice_key
            )

        kwargs = {
            "query": query[:, : self.num_heads_brc1].contiguous().float(),
            "key": key[:, : self.num_heads_brc1].contiguous().float(),
            "value": value[:, : self.num_heads_brc1].contiguous().float(),
            "xyz": xyz.float(),
            "index_0": index_0.int(),
            "index_0_offsets": index_0_offsets.int(),
            "n_max": n_max,
            "index_1": index_1.int(),
            "index_1_offsets": index_1_offsets.int(),
            "sort_idx": sort_idx,
            "window_size": self.window_size,
            "shift_win": self.shift_win,
            "pe_type": self.pe_type,
        }
        if self.pe_type == "contextual":
            kwargs.update(
                {
                    "rel_query": self.rel_query,
                    "rel_key": self.rel_key,
                    "rel_value": self.rel_value,
                    "quant_size": self.quant_size,
                    "quant_grid_length": self.quant_grid_length,
                    "relative_pos_query_table": self.relative_pos_query_table.float(),
                    "relative_pos_key_table": self.relative_pos_key_table.float(),
                    "relative_pos_value_table": self.relative_pos_value_table.float(),
                }
            )
        out1 = sparse_self_attention(**kwargs)

        kwargs = {
            "query": query[:, self.num_heads_brc1 :].contiguous().float(),
            "key": key[:, self.num_heads_brc1 :].contiguous().float(),
            "value": value[:, self.num_heads_brc1 :].contiguous().float(),
            "xyz": xyz_sphere.float(),
            "index_0": index_0_sphere.int(),
            "index_0_offsets": index_0_offsets_sphere.int(),
            "n_max": n_max_sphere,
            "index_1": index_1_sphere.int(),
            "index_1_offsets": index_1_offsets_sphere.int(),
            "sort_idx": sort_idx_sphere,
            "window_size": self.window_size_sphere,
            "shift_win": self.shift_win,
            "pe_type": self.pe_type,
        }
        if self.pe_type == "contextual":
            kwargs.update(
                {
                    "rel_query": self.rel_query,
                    "rel_key": self.rel_key,
                    "rel_value": self.rel_value,
                    "quant_size": self.quant_size_sphere,
                    "quant_grid_length": self.quant_grid_length_sphere,
                    "relative_pos_query_table": self.relative_pos_query_table_sphere.float(),
                    "relative_pos_key_table": self.relative_pos_key_table_sphere.float(),
                    "relative_pos_value_table": self.relative_pos_value_table_sphere.float(),
                    "split_func": partial(exponential_split, a=self.a),
                }
            )
        out2 = sparse_self_attention(**kwargs)

        x = torch.cat([out1, out2], 1).view(N, C)

        x = self.proj(x)
        x = self.proj_drop(x)  # [N, C]

        output_tensor = SparseTrTensor(
            x,
            sptr_tensor.query_indices,
            sptr_tensor.spatial_shape,
            sptr_tensor.batch_size,
        )

        return output_tensor


class SphereFormer(nn.Module):
    def __init__(
        self,
        dim,
        num_heads,
        window_size,
        window_size_sphere,
        quant_size,
        quant_size_sphere,
        indice_key,
        pe_type="contextual",
        rel_query=True,
        rel_key=False,
        rel_value=False,
        drop_path=0.0,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        a=0.05 * 0.25,
    ):
        super().__init__()
        self.window_size = window_size

        self.norm1 = norm_layer(dim)
        self.attn = SparseMultiheadSASphereConcat(
            dim,
            num_heads=num_heads,
            indice_key=indice_key,
            window_size=window_size,
            window_size_sphere=window_size_sphere,
            pe_type=pe_type,
            quant_size=quant_size,
            quant_size_sphere=quant_size_sphere,
            rel_query=rel_query,
            rel_key=rel_key,
            rel_value=rel_value,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            a=a,
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer
        )

    def forward(self, feats, xyz, batch):
        short_cut = feats
        feats = self.norm1(feats)

        sptr_tensor = SparseTrTensor(
            feats,
            torch.cat([batch[:, None], xyz], -1),
            spatial_shape=None,
            batch_size=None,
        )
        sptr_tensor = self.attn(sptr_tensor)
        feats = sptr_tensor.query_feats

        feats = short_cut + self.drop_path(feats)
        feats = feats + self.drop_path(self.mlp(self.norm2(feats)))

        return feats

mask_4d/models/spherical_transformer.py
PRBonn/Mask4D
import os
from os.path import join

import click
import torch
import yaml
from easydict import EasyDict as edict
from mask_4d.datasets.kitti_dataset import SemanticDatasetModule
from mask_4d.models.mask_model import Mask4D
from pytorch_lightning import Trainer, seed_everything


def getDir(obj):
    return os.path.dirname(os.path.abspath(obj))


@click.command()
@click.option("--w", type=str, required=True)
@click.option("--save_testset", is_flag=True)
def main(w, save_testset):
    seed_everything(42, workers=True)
    model_cfg = edict(
        yaml.safe_load(open(join(getDir(__file__), "../config/model.yaml")))
    )
    backbone_cfg = edict(
        yaml.safe_load(open(join(getDir(__file__), "../config/backbone.yaml")))
    )
    decoder_cfg = edict(
        yaml.safe_load(open(join(getDir(__file__), "../config/decoder.yaml")))
    )
    cfg = edict({**model_cfg, **backbone_cfg, **decoder_cfg})

    cfg.EVALUATE = True
    if save_testset:
        results_dir = create_dirs()
        print(f"Saving test set predictions in directory {results_dir}")
        cfg.RESULTS_DIR = results_dir

    data = SemanticDatasetModule(cfg)
    model = Mask4D(cfg)
    w = torch.load(w, map_location="cpu")
    model.load_state_dict(w["state_dict"])

    trainer = Trainer(gpus=cfg.TRAIN.N_GPUS)

    if save_testset:
        trainer.test(model, data)
    else:
        trainer.validate(model, data)
    model.evaluator.print_results()
    print("#############################################################")
    model.evaluator4d.print_results()


def create_dirs():
    results_dir = join(getDir(__file__), "..", "output", "test", "sequences")
    if not os.path.exists(results_dir):
        os.makedirs(results_dir, exist_ok=True)
    for i in range(11, 22):
        sub_dir = os.path.join(results_dir, str(i).zfill(2), "predictions")
        if not os.path.exists(sub_dir):
            os.makedirs(sub_dir, exist_ok=True)
    return results_dir


if __name__ == "__main__":
    main()

mask_4d/scripts/evaluate_model.py
PRBonn/Mask4D
import os
import subprocess
from os.path import join

import click
import torch
import yaml
from easydict import EasyDict as edict
from mask_4d.datasets.kitti_dataset import SemanticDatasetModule
from mask_4d.models.mask_model import Mask4D
from pytorch_lightning import Trainer
from pytorch_lightning import loggers as pl_loggers
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint


@click.command()
@click.option("--w", type=str, default=None, required=False)
@click.option("--ckpt", type=str, default=None, required=False)
def main(w, ckpt):
    model_cfg = edict(
        yaml.safe_load(open(join(getDir(__file__), "../config/model.yaml")))
    )
    backbone_cfg = edict(
        yaml.safe_load(open(join(getDir(__file__), "../config/backbone.yaml")))
    )
    decoder_cfg = edict(
        yaml.safe_load(open(join(getDir(__file__), "../config/decoder.yaml")))
    )
    cfg = edict({**model_cfg, **backbone_cfg, **decoder_cfg})

    data = SemanticDatasetModule(cfg)
    model = Mask4D(cfg)
    if w:
        w = torch.load(w, map_location="cpu")
        model.load_state_dict(w["state_dict"])

    tb_logger = pl_loggers.TensorBoardLogger(
        "experiments/" + cfg.EXPERIMENT.ID, default_hp_metric=False
    )

    # Callbacks
    lr_monitor = LearningRateMonitor(logging_interval="step")

    aq_ckpt = ModelCheckpoint(
        monitor="metrics/aq",
        filename=cfg.EXPERIMENT.ID + "_epoch{epoch:02d}_aq{metrics/aq:.2f}",
        auto_insert_metric_name=False,
        mode="max",
        save_last=True,
    )
    pq_ckpt = ModelCheckpoint(
        monitor="metrics/pq",
        filename=cfg.EXPERIMENT.ID + "_epoch{epoch:02d}_pq{metrics/pq:.2f}",
        auto_insert_metric_name=False,
        mode="max",
        save_last=True,
    )
    pq4d_ckpt = ModelCheckpoint(
        monitor="metrics/pq4d",
        filename=cfg.EXPERIMENT.ID + "_epoch{epoch:02d}_pq4d{metrics/pq4d:.2f}",
        auto_insert_metric_name=False,
        mode="max",
        save_last=True,
    )

    callbacks = [lr_monitor, pq_ckpt, pq4d_ckpt, aq_ckpt]

    trainer = Trainer(
        gpus=cfg.TRAIN.N_GPUS,
        accelerator="ddp",
        logger=tb_logger,
        max_epochs=cfg.TRAIN.MAX_EPOCH,
        callbacks=callbacks,
        log_every_n_steps=1,
        gradient_clip_val=0.5,
        accumulate_grad_batches=cfg.TRAIN.BATCH_ACC,
        resume_from_checkpoint=ckpt,
    )

    trainer.fit(model, data)


def getDir(obj):
    return os.path.dirname(os.path.abspath(obj))


if __name__ == "__main__":
    main()

mask_4d/scripts/train_model.py
PRBonn/Mask4D
# Modified by Rodrigo Marcuzzi from https://github.com/dvlab-research/SphereFormer/blob/master/util/data_util.py
import numpy as np
import torch
from mask_4d.utils.voxelize import voxelize
from torch_scatter import scatter_mean


def data_prepare(
    coord,
    feat,
    label,
    split="train",
    voxel_size=np.array([0.1, 0.1, 0.1]),
    voxel_max=None,
):
    coord_min = np.min(coord, 0)
    # coord -= coord_min
    coord_norm = coord - coord_min
    if split == "train":
        uniq_idx, idx_recon = voxelize(coord_norm, voxel_size)
        coord_voxel = np.floor(coord_norm[uniq_idx] / np.array(voxel_size))
        coord, feat, label = coord[uniq_idx], feat[uniq_idx], label[uniq_idx]
        if voxel_max and label.shape[0] > voxel_max:
            init_idx = np.random.randint(label.shape[0])
            crop_idx = np.argsort(np.sum(np.square(coord - coord[init_idx]), 1))[
                :voxel_max
            ]
            coord, feat, label = coord[crop_idx], feat[crop_idx], label[crop_idx]
            coord_voxel = coord_voxel[crop_idx]
    else:
        idx_recon = voxelize(coord_norm, voxel_size, mode=1)

    coord = torch.FloatTensor(coord)
    feat = torch.FloatTensor(feat)
    label = torch.LongTensor(label)
    if split == "train":
        coord_voxel = torch.LongTensor(coord_voxel)
        idx_recon = torch.LongTensor(idx_recon)
        return coord_voxel, coord, feat, label, idx_recon
    else:
        coord_norm = torch.FloatTensor(coord_norm)
        idx_recon = torch.LongTensor(idx_recon)
        coord_norm = scatter_mean(coord_norm, idx_recon, dim=0)
        coord_voxel = torch.floor(coord_norm / torch.from_numpy(voxel_size)).long()
        coord = scatter_mean(coord, idx_recon, dim=0)
        feat = scatter_mean(feat, idx_recon, dim=0)
        return coord_voxel, coord, feat, label, idx_recon

mask_4d/utils/data_util.py
PRBonn/Mask4D
#!/usr/bin/env python3

# This file is covered by the LICENSE file in the root of this project.

import math

import numpy as np


class PanopticEval:
    """Panoptic evaluation using numpy

    authors: Andres Milioto and Jens Behley

    """

    def __init__(
        self, n_classes, device=None, ignore=None, offset=2**32, min_points=30
    ):
        self.n_classes = n_classes
        assert device == None
        self.ignore = np.array(ignore, dtype=np.int64)
        self.include = np.array(
            [n for n in range(self.n_classes) if n not in self.ignore], dtype=np.int64
        )

        # print("[PANOPTIC EVAL] IGNORE: ", self.ignore)
        # print("[PANOPTIC EVAL] INCLUDE: ", self.include)

        self.reset()
        self.offset = offset  # largest number of instances in a given scan
        self.min_points = (
            min_points  # smallest number of points to consider instances in gt
        )
        self.eps = 1e-15

    def num_classes(self):
        return self.n_classes

    def merge(self, evaluator):
        self.px_iou_conf_matrix += evaluator.px_iou_conf_matrix
        self.pan_tp += evaluator.pan_tp
        self.pan_iou += evaluator.pan_iou
        self.pan_fp += evaluator.pan_fp
        self.pan_fn += evaluator.pan_fn

        self.evaluated_fnames += evaluator.evaluated_fnames

    def reset(self):
        # general things
        # iou stuff
        self.px_iou_conf_matrix = np.zeros(
            (self.n_classes, self.n_classes), dtype=np.int64
        )
        # panoptic stuff
        self.pan_tp = np.zeros(self.n_classes, dtype=np.int64)
        self.pan_iou = np.zeros(self.n_classes, dtype=np.double)
        self.pan_fp = np.zeros(self.n_classes, dtype=np.int64)
        self.pan_fn = np.zeros(self.n_classes, dtype=np.int64)

        self.evaluated_fnames = []

    ################################# IoU STUFF ##################################
    def addBatchSemIoU(self, x_sem, y_sem):
        # idxs are labels and predictions
        idxs = np.stack([x_sem, y_sem], axis=0)

        # make confusion matrix (cols = gt, rows = pred)
        np.add.at(self.px_iou_conf_matrix, tuple(idxs), 1)

    def getSemIoUStats(self):
        # clone to avoid modifying the real deal
        conf = self.px_iou_conf_matrix.copy().astype(np.double)
        # remove fp from confusion on the ignore classes predictions
        # points that were predicted of another class, but were ignore
        # (corresponds to zeroing the cols of those classes, since the predictions
        # go on the rows)
        conf[:, self.ignore] = 0

        # get the clean stats
        tp = conf.diagonal()
        fp = conf.sum(axis=1) - tp
        fn = conf.sum(axis=0) - tp
        return tp, fp, fn

    def getSemIoU(self):
        tp, fp, fn = self.getSemIoUStats()
        # print(f"tp={tp}")
        # print(f"fp={fp}")
        # print(f"fn={fn}")
        intersection = tp
        union = tp + fp + fn
        union = np.maximum(union, self.eps)
        iou = intersection.astype(np.double) / union.astype(np.double)
        iou_mean = (
            intersection[self.include].astype(np.double)
            / union[self.include].astype(np.double)
        ).mean()

        return iou_mean, iou  # returns "iou mean", "iou per class" ALL CLASSES

    def getSemAcc(self):
        tp, fp, fn = self.getSemIoUStats()
        total_tp = tp.sum()
        total = tp[self.include].sum() + fp[self.include].sum()
        total = np.maximum(total, self.eps)
        acc_mean = total_tp.astype(np.double) / total.astype(np.double)

        return acc_mean  # returns "acc mean"

    ################################# IoU STUFF ##################################
    ##############################################################################

    #############################  Panoptic STUFF ################################
    def addBatchPanoptic(self, x_sem_row, x_inst_row, y_sem_row, y_inst_row):
        # make sure instances are not zeros (it messes with my approach)
        x_inst_row = x_inst_row + 1
        y_inst_row = y_inst_row + 1

        # only interested in points that are outside the void area (not in excluded classes)
        for cl in self.ignore:
            # make a mask for this class
            gt_not_in_excl_mask = y_sem_row != cl
            # remove all other points
            x_sem_row = x_sem_row[gt_not_in_excl_mask]
            y_sem_row = y_sem_row[gt_not_in_excl_mask]
            x_inst_row = x_inst_row[gt_not_in_excl_mask]
            y_inst_row = y_inst_row[gt_not_in_excl_mask]

        # first step is to count intersections > 0.5 IoU for each class (except the ignored ones)
        for cl in self.include:
            # print("*"*80)
            # print("CLASS", cl.item())
            # get a class mask
            x_inst_in_cl_mask = x_sem_row == cl
            y_inst_in_cl_mask = y_sem_row == cl

            # get instance points in class (makes outside stuff 0)
            x_inst_in_cl = x_inst_row * x_inst_in_cl_mask.astype(np.int64)
            y_inst_in_cl = y_inst_row * y_inst_in_cl_mask.astype(np.int64)

            # generate the areas for each unique instance prediction
            unique_pred, counts_pred = np.unique(
                x_inst_in_cl[x_inst_in_cl > 0], return_counts=True
            )
            id2idx_pred = {id: idx for idx, id in enumerate(unique_pred)}
            matched_pred = np.array([False] * unique_pred.shape[0])
            # print("Unique predictions:", unique_pred)

            # generate the areas for each unique instance gt_np
            unique_gt, counts_gt = np.unique(
                y_inst_in_cl[y_inst_in_cl > 0], return_counts=True
            )
            id2idx_gt = {id: idx for idx, id in enumerate(unique_gt)}
            matched_gt = np.array([False] * unique_gt.shape[0])
            # print("Unique ground truth:", unique_gt)

            # generate intersection using offset
            valid_combos = np.logical_and(x_inst_in_cl > 0, y_inst_in_cl > 0)
            offset_combo = (
                x_inst_in_cl[valid_combos] + self.offset * y_inst_in_cl[valid_combos]
            )
            unique_combo, counts_combo = np.unique(offset_combo, return_counts=True)

            # generate an intersection map
            # count the intersections with over 0.5 IoU as TP
            gt_labels = unique_combo // self.offset
            pred_labels = unique_combo % self.offset
            gt_areas = np.array([counts_gt[id2idx_gt[id]] for id in gt_labels])
            pred_areas = np.array([counts_pred[id2idx_pred[id]] for id in pred_labels])
            intersections = counts_combo
            unions = gt_areas + pred_areas - intersections
            ious = intersections.astype(float) / unions.astype(float)

            tp_indexes = ious > 0.5
            self.pan_tp[cl] += np.sum(tp_indexes)
            self.pan_iou[cl] += np.sum(ious[tp_indexes])

            matched_gt[[id2idx_gt[id] for id in gt_labels[tp_indexes]]] = True
            matched_pred[[id2idx_pred[id] for id in pred_labels[tp_indexes]]] = True

            # count the FN
            self.pan_fn[cl] += np.sum(
                np.logical_and(counts_gt >= self.min_points, matched_gt == False)
            )

            # count the FP
            self.pan_fp[cl] += np.sum(
                np.logical_and(counts_pred >= self.min_points, matched_pred == False)
            )

    def getPQ(self):
        # first calculate for all classes
        sq_all = self.pan_iou.astype(np.double) / np.maximum(
            self.pan_tp.astype(np.double), self.eps
        )
        rq_all = self.pan_tp.astype(np.double) / np.maximum(
            self.pan_tp.astype(np.double)
            + 0.5 * self.pan_fp.astype(np.double)
            + 0.5 * self.pan_fn.astype(np.double),
            self.eps,
        )
        pq_all = sq_all * rq_all

        # then do the REAL mean (no ignored classes)
        SQ = sq_all[self.include].mean()
        RQ = rq_all[self.include].mean()
        PQ = pq_all[self.include].mean()

        return PQ, SQ, RQ, pq_all, sq_all, rq_all

    #############################  Panoptic STUFF ################################
    ##############################################################################

    def addBatch(self, x_sem, x_inst, y_sem, y_inst):  # x=preds, y=targets
        """IMPORTANT: Inputs must be batched. Either [N,H,W], or [N, P]"""
        # add to IoU calculation (for checking purposes)
        self.addBatchSemIoU(x_sem, y_sem)

        # now do the panoptic stuff
        self.addBatchPanoptic(x_sem, x_inst, y_sem, y_inst)

    def addBatch_w_fname(
        self, x_sem, x_inst, y_sem, y_inst, fname
    ):  # x=preds, y=targets
        """IMPORTANT: Inputs must be batched. Either [N,H,W], or [N, P]"""
        # add to IoU calculation (for checking purposes)
        self.addBatchSemIoU(x_sem, y_sem)

        # now do the panoptic stuff
        self.addBatchPanoptic(x_sem, x_inst, y_sem, y_inst)

        self.evaluated_fnames.append(fname)


class Panoptic4DEval:
    """Panoptic evaluation using numpy

    authors: Andres Milioto, Jens Behley, Aljosa Osep

    """

    def __init__(
        self, n_classes, device=None, ignore=None, offset=2**32, min_points=30
    ):
        self.n_classes = n_classes
        # assert (device == None)
        self.ignore = np.array(ignore, dtype=np.int64)
        self.include = np.array(
            [n for n in range(self.n_classes) if n not in self.ignore], dtype=np.int64
        )
        # print("[PANOPTIC4D EVAL] IGNORE: ", self.ignore)
        # print("[PANOPTIC4D EVAL] INCLUDE: ", self.include)

        self.reset()
        self.offset = offset  # largest number of instances in a given scan
        self.min_points = (
            min_points  # smallest number of points to consider instances in gt
        )
        self.eps = 1e-15

    def num_classes(self):
        return self.n_classes

    def reset(self):
        # general things
        # iou stuff
        self.px_iou_conf_matrix = np.zeros(
            (self.n_classes, self.n_classes), dtype=np.int64
        )

        self.sequences = []
        self.preds = {}
        self.gts = {}
        self.intersects = {}
        self.intersects_ovr = {}

        # Per-class association quality collect here
        self.pan_aq = np.zeros(self.n_classes, dtype=np.double)
        self.pan_aq_ovr = 0.0

    ################################# IoU STUFF ##################################
    def addBatchSemIoU(self, x_sem, y_sem):
        # idxs are labels and predictions
        idxs = np.stack([x_sem, y_sem], axis=0)

        # make confusion matrix (cols = gt, rows = pred)
        np.add.at(self.px_iou_conf_matrix, tuple(idxs), 1)

    def getSemIoUStats(self):
        # clone to avoid modifying the real deal
        conf = self.px_iou_conf_matrix.copy().astype(np.double)
        # remove fp from confusion on the ignore classes predictions
        # points that were predicted of another class, but were ignore
        # (corresponds to zeroing the cols of those classes, since the predictions
        # go on the rows)
        conf[:, self.ignore] = 0

        # get the clean stats
        tp = conf.diagonal()
        fp = conf.sum(axis=1) - tp
        fn = conf.sum(axis=0) - tp
        return tp, fp, fn

    def getSemIoU(self):
        tp, fp, fn = self.getSemIoUStats()
        # print(f"tp={tp}")
        # print(f"fp={fp}")
        # print(f"fn={fn}")
        intersection = tp
        union = tp + fp + fn

        num_present_classes = np.count_nonzero(union)

        union = np.maximum(union, self.eps)
        iou = intersection.astype(np.double) / union.astype(np.double)
        iou_mean = np.sum(iou) / num_present_classes

        return iou_mean, iou  # returns "iou mean", "iou per class" ALL CLASSES

    def getSemAcc(self):
        tp, fp, fn = self.getSemIoUStats()
        total_tp = tp.sum()
        total = tp[self.include].sum() + fp[self.include].sum()
        total = np.maximum(total, self.eps)
        acc_mean = total_tp.astype(np.double) / total.astype(np.double)

        return acc_mean  # returns "acc mean"

    ################################# IoU STUFF ##################################
    ##############################################################################

    #############################  Panoptic STUFF ################################

    def update_dict_stat(self, stat_dict, unique_ids, unique_cnts):
        for uniqueid, counts in zip(unique_ids, unique_cnts):
            if uniqueid == 1:
                continue  # 1 -- no instance
            if uniqueid in stat_dict:
                stat_dict[uniqueid] += counts
            else:
                stat_dict[uniqueid] = counts

    def addBatchPanoptic4D(self, seq, x_sem_row, x_inst_row, y_sem_row, y_inst_row):
        # start = time.time()
        if seq not in self.sequences:
            self.sequences.append(seq)
            self.preds[seq] = {}
            self.gts[seq] = [{} for i in range(self.n_classes)]
            self.intersects[seq] = [{} for i in range(self.n_classes)]
            self.intersects_ovr[seq] = [{} for i in range(self.n_classes)]

        # make sure instances are not zeros (it messes with my approach)
        x_inst_row = x_inst_row + 1
        y_inst_row = y_inst_row + 1

        # only interested in points that are outside the void area (not in excluded classes)
        for cl in self.ignore:
            # make a mask for this class
            gt_not_in_excl_mask = y_sem_row != cl
            # remove all other points
            x_sem_row = x_sem_row[gt_not_in_excl_mask]
            y_sem_row = y_sem_row[gt_not_in_excl_mask]
            x_inst_row = x_inst_row[gt_not_in_excl_mask]
            y_inst_row = y_inst_row[gt_not_in_excl_mask]

        for cl in self.include:
            # print("*"*80)
            # print("CLASS", cl.item())

            # Per-class accumulated stats
            cl_preds = self.preds[seq]
            cl_gts = self.gts[seq][cl]
            cl_intersects = self.intersects[seq][cl]

            # get a binary class mask (filter acc. to semantic class!)
            x_inst_in_cl_mask = x_sem_row == cl
            y_inst_in_cl_mask = y_sem_row == cl

            # get instance points in class (mask-out everything but _this_ class)
            x_inst_in_cl = x_inst_row * x_inst_in_cl_mask.astype(np.int64)
            y_inst_in_cl = y_inst_row * y_inst_in_cl_mask.astype(np.int64)

            # generate the areas for each unique instance gt_np (i.e., set2)
            unique_gt, counts_gt = np.unique(
                y_inst_in_cl[y_inst_in_cl > 0], return_counts=True
            )
            self.update_dict_stat(
                cl_gts,
                unique_gt[counts_gt > self.min_points],
                counts_gt[counts_gt > self.min_points],
            )

            valid_combos_min_point = np.zeros_like(
                y_inst_in_cl
            )  # instances which have at least self.min points
            for valid_id in unique_gt[counts_gt > self.min_points]:
                valid_combos_min_point = np.logical_or(
                    valid_combos_min_point, y_inst_in_cl == valid_id
                )

            y_inst_in_cl = y_inst_in_cl * valid_combos_min_point
            # generate the areas for each unique instance prediction (i.e., set1)
            unique_pred, counts_pred = np.unique(
                x_inst_in_cl[x_inst_in_cl > 0], return_counts=True
            )

            # is there better way to do this?
            self.update_dict_stat(cl_preds, unique_pred, counts_pred)

            valid_combos = np.logical_and(
                x_inst_row > 0, y_inst_in_cl > 0
            )  # Convert to boolean and do logical and, based on semantics

            # generate intersection using offset
            offset_combo = (
                x_inst_row[valid_combos] + self.offset * y_inst_in_cl[valid_combos]
            )
            unique_combo, counts_combo = np.unique(offset_combo, return_counts=True)

            self.update_dict_stat(cl_intersects, unique_combo, counts_combo)

    def getPQ4D(self):
        num_tubes = [0] * self.n_classes
        for seq in self.sequences:
            for cl in self.include:
                cl_preds = self.preds[seq]
                cl_gts = self.gts[seq][cl]
                cl_intersects = self.intersects[seq][cl]
                outer_sum = 0.0
                outer_sum_iou = 0.0
                num_tubes[cl] += len(cl_gts)
                for gt_id, gt_size in cl_gts.items():
                    inner_sum = 0.0
                    inner_sum_iou = 0.0
                    for pr_id, pr_size in cl_preds.items():
                        # TODO: pay attention for zero intersection!
                        TPA_key = pr_id + self.offset * gt_id
                        if TPA_key in cl_intersects:
                            TPA = cl_intersects[TPA_key]
                            Prec = TPA / float(
                                pr_size
                            )  # TODO I dont think these can ever be zero, but double check
                            Recall = TPA / float(gt_size)
                            TPA_ovr = self.intersects[seq][cl][TPA_key]
                            inner_sum_iou += TPA_ovr * (
                                TPA_ovr / (gt_size + pr_size - TPA_ovr)
                            )
                            if Prec > 1.0 or Recall > 1.0:
                                print("something wrong !!")
                    outer_sum_iou += 1.0 / float(gt_size) * float(inner_sum_iou)
                self.pan_aq[
                    cl
                ] += outer_sum_iou  # 1.0/float(len(cl_gts)) # Normalize by #tubes
                self.pan_aq_ovr += outer_sum_iou
        # ==========

        AQ_overall = np.sum(self.pan_aq_ovr) / np.sum(num_tubes[1:9])
        AQ = self.pan_aq / np.maximum(num_tubes, self.eps)

        iou_mean, iou = self.getSemIoU()

        PQ4D = math.sqrt(AQ_overall * iou_mean)
        return PQ4D, AQ_overall, AQ, iou, iou_mean

    #############################  Panoptic STUFF ################################
    ##############################################################################

    def addBatch(self, seq, x_sem, x_inst, y_sem, y_inst):  # x=preds, y=targets
        """IMPORTANT: Inputs must be batched. Either [N,H,W], or [N, P]"""
        # add to IoU calculation (for checking purposes)
        self.addBatchSemIoU(x_sem, y_sem)

        # now do the panoptic stuff
        self.addBatchPanoptic4D(seq, x_sem, x_inst, y_sem, y_inst)


def gen_psuedo_labels(n=50):
    # generate ground truth and prediction
    sem_pred = []
    inst_pred = []
    sem_gt = []
    inst_gt = []

    # some ignore stuff
    N_ignore = n
    sem_pred.extend([0 for i in range(N_ignore)])
    inst_pred.extend([0 for i in range(N_ignore)])
    sem_gt.extend([0 for i in range(N_ignore)])
    inst_gt.extend([0 for i in range(N_ignore)])

    # grass segment
    N_grass = n + 1
    N_grass_pred = np.random.randint(0, N_grass)  # rest is sky
    sem_pred.extend([1 for i in range(N_grass_pred)])  # grass
    sem_pred.extend([2 for i in range(N_grass - N_grass_pred)])  # sky
    inst_pred.extend([0 for i in range(N_grass)])
    sem_gt.extend([1 for i in range(N_grass)])  # grass
    inst_gt.extend([0 for i in range(N_grass)])

    # sky segment
    N_sky = n + 2
    N_sky_pred = np.random.randint(0, N_sky)  # rest is grass
    sem_pred.extend([2 for i in range(N_sky_pred)])  # sky
    sem_pred.extend([1 for i in range(N_sky - N_sky_pred)])  # grass
    inst_pred.extend([0 for i in range(N_sky)])  # first instance
    sem_gt.extend([2 for i in range(N_sky)])  # sky
    inst_gt.extend([0 for i in range(N_sky)])  # first instance

    # wrong dog as person prediction
    N_dog = n + 3
    N_person = N_dog
    sem_pred.extend([3 for i in range(N_person)])
    inst_pred.extend([35 for i in range(N_person)])
    sem_gt.extend([4 for i in range(N_dog)])
    inst_gt.extend([22 for i in range(N_dog)])

    # two persons in prediction, but three in gt
    N_person = n + 4
    sem_pred.extend([3 for i in range(6 * N_person)])
    inst_pred.extend([8 for i in range(4 * N_person)])
    inst_pred.extend([95 for i in range(2 * N_person)])
    sem_gt.extend([3 for i in range(6 * N_person)])
    inst_gt.extend([33 for i in range(3 * N_person)])
    inst_gt.extend([42 for i in range(N_person)])
    inst_gt.extend([11 for i in range(2 * N_person)])

    # gt and pred to numpy
    sem_pred = np.array(sem_pred, dtype=np.int64).reshape(1, -1)
    inst_pred = np.array(inst_pred, dtype=np.int64).reshape(1, -1)
    sem_gt = np.array(sem_gt, dtype=np.int64).reshape(1, -1)
    inst_gt = np.array(inst_gt, dtype=np.int64).reshape(1, -1)

    return sem_pred, inst_pred, sem_gt, inst_gt


def test_4D():
    classes = 3  # ignore, car, truck
    cl_strings = ["ignore", "car", "truck"]
    ignore = [0]  # only ignore ignore class

    sem_gt = np.zeros(20, dtype=np.int32)
    sem_gt[5:10] = 1
    sem_gt[10:] = 2

    inst_gt = np.zeros(20, dtype=np.int32)
    inst_gt[5:10] = 1
    inst_gt[10:] = 1
    inst_gt[15:] = 2

    # we have 3 instance 1 car, 2 truck as gt
    sem_pred = np.zeros(20, dtype=np.int32)
    sem_pred[5:10] = 1
    sem_pred[10:15] = 2
    sem_pred[15:] = 1

    inst_pred = np.zeros(20, dtype=np.int32)
    inst_pred[5:10] = 1
    inst_pred[10:] = 2

    # evaluator
    class_evaluator = Panoptic4DEval(3, None, ignore, offset=2**32, min_points=1)
    class_evaluator.addBatch(1, sem_pred, inst_pred, sem_gt, inst_gt)
    PQ4D, AQ_ovr, AQ, iou, iou_mean = class_evaluator.getPQ4D()

    print("=== Results ===")
    print("PQ4D:", PQ4D)
    print("AQ_mean:", AQ_ovr)
    print("AQ:", AQ)
    print("iou:", iou)
    print("iou_mean:", iou_mean)


if __name__ == "__main__":
    # generate problem from He paper (https://arxiv.org/pdf/1801.00868.pdf)
    classes = 5  # ignore, grass, sky, person, dog
    cl_strings = ["ignore", "grass", "sky", "person", "dog"]
    ignore = [0]  # only ignore ignore class
    min_points = 1  # for this example we care about all points

    sem_pred_0, inst_pred_0, sem_gt_0, inst_gt_0 = gen_psuedo_labels(50)
    sem_pred_1, inst_pred_1, sem_gt_1, inst_gt_1 = gen_psuedo_labels(51)

    # evaluator
    evaluator = PanopticEval(classes, ignore=ignore, min_points=1)
    evaluator.addBatch(sem_pred_0, inst_pred_0, sem_gt_0, inst_gt_0)
    evaluator.addBatch(sem_pred_1, inst_pred_1, sem_gt_1, inst_gt_1)

    evaluator_0 = PanopticEval(classes, ignore=ignore, min_points=1)
    evaluator_0.addBatch(sem_pred_0, inst_pred_0, sem_gt_0, inst_gt_0)

    evaluator_1 = PanopticEval(classes, ignore=ignore, min_points=1)
    evaluator_1.addBatch(sem_pred_1, inst_pred_1, sem_gt_1, inst_gt_1)

    pq, sq, rq, all_pq, all_sq, all_rq = evaluator.getPQ()
    iou, all_iou = evaluator.getSemIoU()

    # [PANOPTIC EVAL] IGNORE:  [0]
    # [PANOPTIC EVAL] INCLUDE:  [1 2 3 4]
    # TOTALS
    # PQ: 0.47916666666666663
    # SQ: 0.5520833333333333
    # RQ: 0.6666666666666666
    # IoU: 0.5476190476190476
    # Class ignore 	 PQ: 0.0 SQ: 0.0 RQ: 0.0 IoU: 0.0
    # Class grass 	 PQ: 0.6666666666666666 SQ: 0.6666666666666666 RQ: 1.0 IoU: 0.6666666666666666
    # Class sky 	 PQ: 0.6666666666666666 SQ: 0.6666666666666666 RQ: 1.0 IoU: 0.6666666666666666
    # Class person 	 PQ: 0.5833333333333333 SQ: 0.875 RQ: 0.6666666666666666 IoU: 0.8571428571428571
    # Class dog 	 PQ: 0.0 SQ: 0.0 RQ: 0.0 IoU: 0.0

    print("TOTALS")
    print("PQ:", pq.item(), pq.item() == 0.47916666666666663)
    print("SQ:", sq.item(), sq.item() == 0.5520833333333333)
    print("RQ:", rq.item(), rq.item() == 0.6666666666666666)
    print("IoU:", iou.item(), iou.item() == 0.5476190476190476)
    for i, (pq, sq, rq, iou) in enumerate(zip(all_pq, all_sq, all_rq, all_iou)):
        print(
            "Class",
            cl_strings[i],
            "\t",
            "PQ:",
            pq.item(),
            "SQ:",
            sq.item(),
            "RQ:",
            rq.item(),
            "IoU:",
            iou.item(),
        )

    pq, sq, rq, all_pq, all_sq, all_rq = evaluator_0.getPQ()
    iou, all_iou = evaluator_0.getSemIoU()

    # [PANOPTIC EVAL] IGNORE:  [0]
    # [PANOPTIC EVAL] INCLUDE:  [1 2 3 4]
    # TOTALS
    # PQ: 0.47916666666666663
    # SQ: 0.5520833333333333
    # RQ: 0.6666666666666666
    # IoU: 0.5476190476190476
    # Class ignore 	 PQ: 0.0 SQ: 0.0 RQ: 0.0 IoU: 0.0
    # Class grass 	 PQ: 0.6666666666666666 SQ: 0.6666666666666666 RQ: 1.0 IoU: 0.6666666666666666
    # Class sky 	 PQ: 0.6666666666666666 SQ: 0.6666666666666666 RQ: 1.0 IoU: 0.6666666666666666
    # Class person 	 PQ: 0.5833333333333333 SQ: 0.875 RQ: 0.6666666666666666 IoU: 0.8571428571428571
    # Class dog 	 PQ: 0.0 SQ: 0.0 RQ: 0.0 IoU: 0.0

    print("TOTALS")
    print("PQ:", pq.item(), pq.item() == 0.47916666666666663)
    print("SQ:", sq.item(), sq.item() == 0.5520833333333333)
    print("RQ:", rq.item(), rq.item() == 0.6666666666666666)
    print("IoU:", iou.item(), iou.item() == 0.5476190476190476)
    for i, (pq, sq, rq, iou) in enumerate(zip(all_pq, all_sq, all_rq, all_iou)):
        print(
            "Class",
            cl_strings[i],
            "\t",
            "PQ:",
            pq.item(),
            "SQ:",
            sq.item(),
            "RQ:",
            rq.item(),
            "IoU:",
            iou.item(),
        )

    pq, sq, rq, all_pq, all_sq, all_rq = evaluator_1.getPQ()
    iou, all_iou = evaluator_1.getSemIoU()

    # [PANOPTIC EVAL] IGNORE:  [0]
    # [PANOPTIC EVAL] INCLUDE:  [1 2 3 4]
    # TOTALS
    # PQ: 0.47916666666666663
    # SQ: 0.5520833333333333
    # RQ: 0.6666666666666666
    # IoU: 0.5476190476190476
    # Class ignore 	 PQ: 0.0 SQ: 0.0 RQ: 0.0 IoU: 0.0
    # Class grass 	 PQ: 0.6666666666666666 SQ: 0.6666666666666666 RQ: 1.0 IoU: 0.6666666666666666
    # Class sky 	 PQ: 0.6666666666666666 SQ: 0.6666666666666666 RQ: 1.0 IoU: 0.6666666666666666
    # Class person 	 PQ: 0.5833333333333333 SQ: 0.875 RQ: 0.6666666666666666 IoU: 0.8571428571428571
    # Class dog 	 PQ: 0.0 SQ: 0.0 RQ: 0.0 IoU: 0.0

    print("TOTALS")
    print("PQ:", pq.item(), pq.item() == 0.47916666666666663)
    print("SQ:", sq.item(), sq.item() == 0.5520833333333333)
    print("RQ:", rq.item(), rq.item() == 0.6666666666666666)
    print("IoU:", iou.item(), iou.item() == 0.5476190476190476)
    for i, (pq, sq, rq, iou) in enumerate(zip(all_pq, all_sq, all_rq, all_iou)):
        print(
            "Class",
            cl_strings[i],
            "\t",
            "PQ:",
            pq.item(),
            "SQ:",
            sq.item(),
            "RQ:",
            rq.item(),
            "IoU:",
            iou.item(),
        )

    import pickle

    with open("test.pkl", "wb") as f:
        pickle.dump(evaluator_0, f)

    with open("test.pkl", "rb") as f:
        evaluator_read = pickle.load(f)
    evaluator_1.merge(evaluator_read)

    pq, sq, rq, all_pq, all_sq, all_rq = evaluator_1.getPQ()
    iou, all_iou = evaluator_1.getSemIoU()

    # [PANOPTIC EVAL] IGNORE:  [0]
    # [PANOPTIC EVAL] INCLUDE:  [1 2 3 4]
    # TOTALS
    # PQ: 0.47916666666666663
    # SQ: 0.5520833333333333
    # RQ: 0.6666666666666666
    # IoU: 0.5476190476190476
    # Class ignore 	 PQ: 0.0 SQ: 0.0 RQ: 0.0 IoU: 0.0
    # Class grass 	 PQ: 0.6666666666666666 SQ: 0.6666666666666666 RQ: 1.0 IoU: 0.6666666666666666
    # Class sky 	 PQ: 0.6666666666666666 SQ: 0.6666666666666666 RQ: 1.0 IoU: 0.6666666666666666
    # Class person 	 PQ: 0.5833333333333333 SQ: 0.875 RQ: 0.6666666666666666 IoU: 0.8571428571428571
    # Class dog 	 PQ: 0.0 SQ: 0.0 RQ: 0.0 IoU: 0.0

    print("TOTALS")
    print("PQ:", pq.item(), pq.item() == 0.47916666666666663)
    print("SQ:", sq.item(), sq.item() == 0.5520833333333333)
    print("RQ:", rq.item(), rq.item() == 0.6666666666666666)
    print("IoU:", iou.item(), iou.item() == 0.5476190476190476)
    for i, (pq, sq, rq, iou) in enumerate(zip(all_pq, all_sq, all_rq, all_iou)):
        print(
            "Class",
            cl_strings[i],
            "\t",
            "PQ:",
            pq.item(),
            "SQ:",
            sq.item(),
            "RQ:",
            rq.item(),
            "IoU:",
            iou.item(),
        )

mask_4d/utils/eval_np.py
PRBonn/Mask4D
import time

import numpy as np
import yaml
from mask_4d.utils.eval_np import Panoptic4DEval


class PanopticKitti4DEvaluator:
    def __init__(self, cfg):
        dataset_config_file = cfg.CONFIG
        self.load_kitti_config(dataset_config_file)
        min_points = 50
        self.evaluator = Panoptic4DEval(
            self.nr_classes,
            None,
            self.ignore_class,
            offset=2**32,
            min_points=min_points,
        )
        self.class_metrics = {}
        self.mean_metrics = {}

    def reset(self):
        self.evaluator.reset()
        self.class_metrics = {}
        self.mean_metrics = {}

    def load_kitti_config(self, config_file):
        # Load semantic-kitti config
        # data = yaml.safe_load(open('datasets/semantic-kitti.yaml', 'r'))
        data = yaml.safe_load(open(config_file, "r"))
        # get number of interest classes, and the label mappings
        class_remap = data["learning_map"]
        self.class_inv_remap = data["learning_map_inv"]
        class_ignore = data["learning_ignore"]
        self.nr_classes = len(self.class_inv_remap)
        self.class_strings = data["labels"]
        # make lookup table for mapping
        maxkey = max(class_remap.keys())
        # +100 hack making lut bigger just in case there are unknown labels
        class_lut = np.zeros((maxkey + 100), dtype=np.int32)
        class_lut[list(class_remap.keys())] = list(class_remap.values())
        self.ignore_class = [cl for cl, ignored in class_ignore.items() if ignored]

        self.class_inv_lut = np.zeros((20), dtype=np.int32)
        self.class_inv_lut[list(self.class_inv_remap.keys())] = list(
            self.class_inv_remap.values()
        )

        self.things = get_things()
        self.stuff = get_stuff()
        self.all_classes = self.things + self.stuff

    def update(self, sem_preds, ins_preds, inputs):
        for i in range(len(sem_preds)):
            self.evaluator.addBatch(
                inputs["fname"][i].split("/")[-3],
                sem_preds[i],
                ins_preds[i],
                inputs["sem_label"][i].reshape(-1),
                inputs["ins_label"][i].reshape(-1),
            )

    def get_class_inv_lut(self):
        return self.class_inv_lut

    def calculate_metrics(self):
        PQ4D, AQ_mean, AQ, iou, iou_mean = self.evaluator.getPQ4D()
        things_iou = iou[1:9].mean()
        stuff_iou = iou[9:].mean()

        # print("=== Results ===")
        # print("PQ4D:", PQ4D)
        # print("AQ_mean:", AQ_mean)
        # print("AQ:", AQ)
        # print("iou:", iou)
        # print("iou_mean:", iou_mean)
        # print("things_iou:", things_iou)
        # print("stuff_iou:", stuff_iou)

        metrics_dict = {}
        metrics_dict["PQ4D"] = float(PQ4D)
        metrics_dict["AQ_mean"] = float(AQ_mean)
        metrics_dict["iou_mean"] = float(iou_mean)
        metrics_dict["things_iou"] = things_iou
        metrics_dict["stuff_iou"] = stuff_iou

        # Save mean metrics
        self.mean_metrics = metrics_dict

        # Save to a file
        # with open(output_filename, 'w') as outfile:
        # yaml.dump(codalab_output, outfile, default_flow_style=False)

        class_all_IoU = iou.flatten().tolist()
        class_all_AQ = AQ.flatten().tolist()

        classes_dict = {}
        for idx, (aq, iou) in enumerate(zip(class_all_AQ, class_all_IoU)):
            class_str = self.class_strings[self.class_inv_remap[idx]]
            classes_dict[class_str] = {}
            classes_dict[class_str]["AQ"] = aq
            classes_dict[class_str]["IoU"] = iou

        # Save per class metrics
        self.class_metrics = classes_dict

    def get_mean_pq4d(self):
        return self.mean_metrics["PQ4D"]

    def get_mean_aq(self):
        return self.mean_metrics["AQ_mean"]

    def get_mean_iou(self):
        return self.mean_metrics["iou_mean"]

    def get_class_metrics(self):
        return self.class_metrics

    def print_results(self):
        print("|        |   AQ   |   IoU  |")
        for k, v in self.class_metrics.items():
            print(
                "|{}| {:.4f} | {:.4f} |".format(
                    k.ljust(8)[-8:],
                    v["AQ"],
                    v["IoU"],
                )
            )
        for key in self.mean_metrics.keys():
            print("{}:\t{}".format(key, self.mean_metrics[key]))

    def print_fp_fn(self):
        print("True Positive: ")
        print("\t|\t".join([str(x) for x in self.evaluator.pan_tp]))
        print("False Positive: ")
        print("\t|\t".join([str(x) for x in self.evaluator.pan_fp]))
        print("False Negative: ")
        print("\t|\t".join([str(x) for x in self.evaluator.pan_fn]))


def get_things():
    things = [
        "car",
        "bicycle",
        "motorcycle",
        "truck",
        "other-vehicle",
        "person",
        "bicyclist",
        "motorcyclist",
    ]
    return things


def get_stuff():
    stuff = [
        "road",
        "parking",
        "sidewalk",
        "other-ground",
        "building",
        "fence",
        "vegetation",
        "trunk",
        "terrain",
        "pole",
        "traffic-sign",
    ]
    return stuff

mask_4d/utils/evaluate_4dpanoptic.py
PRBonn/Mask4D
import time

import numpy as np
import yaml
from mask_4d.utils.eval_np import PanopticEval


class PanopticKittiEvaluator:
    def __init__(self, cfg):
        dataset_config_file = cfg.CONFIG
        self.load_kitti_config(dataset_config_file)
        min_points = 50
        self.evaluator = PanopticEval(
            self.nr_classes, None, self.ignore_class, min_points=min_points
        )
        self.class_metrics = {}
        self.mean_metrics = {}

    def reset(self):
        self.evaluator.reset()
        self.class_metrics = {}
        self.mean_metrics = {}

    def load_kitti_config(self, config_file):
        # Load semantic-kitti config
        # data = yaml.safe_load(open('datasets/semantic-kitti.yaml', 'r'))
        data = yaml.safe_load(open(config_file, "r"))
        # get number of interest classes, and the label mappings
        class_remap = data["learning_map"]
        self.class_inv_remap = data["learning_map_inv"]
        class_ignore = data["learning_ignore"]
        self.nr_classes = len(self.class_inv_remap)
        self.class_strings = data["labels"]
        # make lookup table for mapping
        maxkey = max(class_remap.keys())
        # +100 hack making lut bigger just in case there are unknown labels
        class_lut = np.zeros((maxkey + 100), dtype=np.int32)
        class_lut[list(class_remap.keys())] = list(class_remap.values())
        self.ignore_class = [cl for cl, ignored in class_ignore.items() if ignored]

        self.class_inv_lut = np.zeros((20), dtype=np.int32)
        self.class_inv_lut[list(self.class_inv_remap.keys())] = list(
            self.class_inv_remap.values()
        )

        self.things = get_things()
        self.stuff = get_stuff()
        self.all_classes = self.things + self.stuff

    def update(self, sem_preds, ins_preds, inputs):
        for i in range(len(sem_preds)):
            self.evaluator.addBatch_w_fname(
                sem_preds[i],
                ins_preds[i],
                inputs["sem_label"][i].reshape(-1),
                inputs["ins_label"][i].reshape(-1),
                inputs["fname"][i],
            )
        self.update_metrics()

    def get_class_inv_lut(self):
        return self.class_inv_lut

    def update_metrics(self):
        (
            class_PQ,
            class_SQ,
            class_RQ,
            class_all_PQ,
            class_all_SQ,
            class_all_RQ,
        ) = self.evaluator.getPQ()
        class_IoU, class_all_IoU = self.evaluator.getSemIoU()

        # now make a nice dictionary
        output_dict = {}

        # make python variables
        class_PQ = class_PQ.item()
        class_SQ = class_SQ.item()
        class_RQ = class_RQ.item()
        class_all_PQ = class_all_PQ.flatten().tolist()
        class_all_SQ = class_all_SQ.flatten().tolist()
        class_all_RQ = class_all_RQ.flatten().tolist()
        class_IoU = class_IoU.item()
        class_all_IoU = class_all_IoU.flatten().tolist()

        output_dict["all"] = {}
        output_dict["all"]["PQ"] = class_PQ
        output_dict["all"]["SQ"] = class_SQ
        output_dict["all"]["RQ"] = class_RQ
        output_dict["all"]["IoU"] = class_IoU

        classwise_tables = {}

        for idx, (pq, rq, sq, iou) in enumerate(
            zip(class_all_PQ, class_all_RQ, class_all_SQ, class_all_IoU)
        ):
            class_str = self.class_strings[self.class_inv_remap[idx]]
            output_dict[class_str] = {}
            output_dict[class_str]["PQ"] = pq
            output_dict[class_str]["SQ"] = sq
            output_dict[class_str]["RQ"] = rq
            output_dict[class_str]["IoU"] = iou

        # Save per class metrics
        self.class_metrics = output_dict

        PQ_all = np.mean([float(output_dict[c]["PQ"]) for c in self.all_classes])
        PQ_dagger = np.mean(
            [float(output_dict[c]["PQ"]) for c in self.things]
            + [float(output_dict[c]["IoU"]) for c in self.stuff]
        )
        RQ_all = np.mean([float(output_dict[c]["RQ"]) for c in self.all_classes])
        SQ_all = np.mean([float(output_dict[c]["SQ"]) for c in self.all_classes])

        PQ_things = np.mean([float(output_dict[c]["PQ"]) for c in self.things])
        RQ_things = np.mean([float(output_dict[c]["RQ"]) for c in self.things])
        SQ_things = np.mean([float(output_dict[c]["SQ"]) for c in self.things])

        PQ_stuff = np.mean([float(output_dict[c]["PQ"]) for c in self.stuff])
        RQ_stuff = np.mean([float(output_dict[c]["RQ"]) for c in self.stuff])
        SQ_stuff = np.mean([float(output_dict[c]["SQ"]) for c in self.stuff])
        mIoU = output_dict["all"]["IoU"]

        codalab_output = {}
        codalab_output["pq_mean"] = float(PQ_all)
        codalab_output["pq_dagger"] = float(PQ_dagger)
        codalab_output["sq_mean"] = float(SQ_all)
        codalab_output["rq_mean"] = float(RQ_all)
        codalab_output["iou_mean"] = float(mIoU)
        codalab_output["pq_stuff"] = float(PQ_stuff)
        codalab_output["rq_stuff"] = float(RQ_stuff)
        codalab_output["sq_stuff"] = float(SQ_stuff)
        codalab_output["pq_things"] = float(PQ_things)
        codalab_output["rq_things"] = float(RQ_things)
        codalab_output["sq_things"] = float(SQ_things)

        # Save mean metrics
        self.mean_metrics = codalab_output

    def get_mean_pq(self):
        return self.mean_metrics["pq_mean"]

    def get_mean_iou(self):
        return self.mean_metrics["iou_mean"]

    def get_mean_rq(self):
        return self.mean_metrics["rq_mean"]

    def get_class_metrics(self):
        return self.class_metrics

    def print_results(self):
        evaluated_fnames = self.evaluator.evaluated_fnames
        print(
            "Evaluated {} frames. Duplicated frame number: {}".format(
                len(evaluated_fnames),
                len(evaluated_fnames) - len(set(evaluated_fnames)),
            )
        )
        print("|        |   PQ   |   RQ   |   SQ   |  IoU   |")
        for k, v in self.class_metrics.items():
            print(
                "|{}| {:.4f} | {:.4f} | {:.4f} | {:.4f} |".format(
                    k.ljust(8)[-8:], v["PQ"], v["RQ"], v["SQ"], v["IoU"]
                )
            )
        for key in self.mean_metrics.keys():
            print("{}:\t{}".format(key, self.mean_metrics[key]))

    def print_fp_fn(self):
        print("True Positive: ")
        print("\t|\t".join([str(x) for x in self.evaluator.pan_tp]))
        print("False Positive: ")
        print("\t|\t".join([str(x) for x in self.evaluator.pan_fp]))
        print("False Negative: ")
        print("\t|\t".join([str(x) for x in self.evaluator.pan_fn]))


def get_things():
    things = [
        "car",
        "bicycle",
        "motorcycle",
        "truck",
        "other-vehicle",
        "person",
        "bicyclist",
        "motorcyclist",
    ]
    return things


def get_stuff():
    stuff = [
        "road",
        "parking",
        "sidewalk",
        "other-ground",
        "building",
        "fence",
        "vegetation",
        "trunk",
        "terrain",
        "pole",
        "traffic-sign",
    ]
    return stuff

mask_4d/utils/evaluate_panoptic.py
PRBonn/Mask4D
# Modified by Rodrigo Marcuzzi from https://github.com/megvii-research/MOTR/blob/main/models/structures/instances.py
import itertools
from typing import Any, Dict, List, Union

import torch


class Tracks:
    """
    This class represents a list of tracks.
    It stores the attributes of tracks (e.g., id, labels, scores) as "fields".
    All fields must have the same ``__len__`` which is the number of tracks.

    All other (non-field) attributes of this class are considered private:
    they must start with '_' and are not modifiable by a user.

    Some basic usage:

    1. Set/get/check a field:

       .. code-block:: python

          tracks.gt_boxes = Boxes(...)
          print(tracks.pred_masks)  # a tensor of shape (N, H, W)
          print('gt_masks' in tracks)

    2. ``len(tracks)`` returns the number of tracks
    3. Indexing: ``tracks[indices]`` will apply the indexing on all the fields
       and returns a new :class:`Tracks`.
       Typically, ``indices`` is a integer vector of indices,
       or a binary mask of length ``num_instances``

       .. code-block:: python

          category_3_detections = tracks[tracks.pred_classes == 3]
          confident_detections = tracks[tracks.scores > 0.9]
    """

    def __init__(self, **kwargs: Any):
        """
        Args:
            kwargs: fields to add to this `Tracks`.
        """
        self._fields: Dict[str, Any] = {}
        for k, v in kwargs.items():
            self.set(k, v)

    def __setattr__(self, name: str, val: Any) -> None:
        if name.startswith("_"):
            super().__setattr__(name, val)
        else:
            self.set(name, val)

    def __getattr__(self, name: str) -> Any:
        if name == "_fields" or name not in self._fields:
            raise AttributeError(
                "Cannot find field '{}' in the given Tracks!".format(name)
            )
        return self._fields[name]

    def set(self, name: str, value: Any) -> None:
        """
        Set the field named `name` to `value`.
        The length of `value` must be the number of Tracks,
        and must agree with other existing fields in this object.
        """
        data_len = len(value)
        if len(self._fields):
            assert (
                len(self) == data_len
            ), "Adding a field of length {} to a Tracks of length {}".format(
                data_len, len(self)
            )
        self._fields[name] = value

    def has(self, name: str) -> bool:
        """
        Returns:
            bool: whether the field called `name` exists.
        """
        return name in self._fields

    def remove(self, name: str) -> None:
        """
        Remove the field called `name`.
        """
        del self._fields[name]

    def get(self, name: str) -> Any:
        """
        Returns the field called `name`.
        """
        return self._fields[name]

    def get_fields(self) -> Dict[str, Any]:
        """
        Returns:
            dict: a dict which maps names (str) to data of the fields

        Modifying the returned dict will modify this instance.
        """
        return self._fields

    # Tensor-like methods
    def to(self, *args: Any, **kwargs: Any) -> "Tracks":
        """
        Returns:
            Tracks: all fields are called with a `to(device)`, if the field has this method.
        """
        ret = Tracks()
        for k, v in self._fields.items():
            if hasattr(v, "to"):
                v = v.to(*args, **kwargs)
            ret.set(k, v)
        return ret

    def numpy(self):
        ret = Tracks()
        for k, v in self._fields.items():
            if hasattr(v, "numpy"):
                v = v.numpy()
            ret.set(k, v)
        return ret

    def __getitem__(self, item: Union[int, slice, torch.BoolTensor]):
        """
        Args:
            item: an index-like object and will be used to index all the fields.

        Returns:
            If `item` is a string, return the data in the corresponding field.
            Otherwise, returns an `Tracks` where all fields are indexed by `item`.
        """
        if type(item) == int:
            if item >= len(self) or item < -len(self):
                raise IndexError("Tracks index out of range!")
            else:
                item = slice(item, None, len(self))

        ret = Tracks()
        for k, v in self._fields.items():
            ret.set(k, v[item])
        return ret

    def __len__(self) -> int:
        for v in self._fields.values():
            # use __len__ because len() has to be int and is not friendly to tracing
            return v.__len__()
        raise NotImplementedError("Empty Tracks does not support __len__!")

    def __iter__(self):
        raise NotImplementedError("`Tracks` object is not iterable!")

    @staticmethod
    def cat(instance_lists: List["Tracks"]) -> "Tracks":
        """
        Args:
            instance_lists (list[Tracks])

        Returns:
            Tracks
        """
        assert all(isinstance(i, Tracks) for i in instance_lists)
        assert len(instance_lists) > 0
        if len(instance_lists) == 1:
            return instance_lists[0]

        ret = Tracks()
        for k in instance_lists[0]._fields.keys():
            values = [i.get(k) for i in instance_lists]
            v0 = values[0]
            if isinstance(v0, torch.Tensor):
                values = torch.cat(values, dim=0)
            elif isinstance(v0, list):
                values = list(itertools.chain(*values))
            elif hasattr(type(v0), "cat"):
                values = type(v0).cat(values)
            else:
                raise ValueError(
                    "Unsupported type {} for concatenation".format(type(v0))
                )
            ret.set(k, values)
        return ret

    def __str__(self) -> str:
        s = self.__class__.__name__ + "("
        s += "num_instances={}, ".format(len(self))
        s += "fields=[{}])".format(
            ", ".join((f"{k}: {v}" for k, v in self._fields.items()))
        )
        return s

    __repr__ = __str__

mask_4d/utils/instances.py
PRBonn/Mask4D
# Author: Xinshuo Weng
# email: xinshuo.weng@gmail.com

import numpy as np
from filterpy.kalman import KalmanFilter


class KalmanBoxTracker(object):
    """
    This class represents the internel state of individual tracked objects observed as bbox.
    """

    count = 0

    def __init__(self, bbox3D):
        """
        Initialises a tracker using initial bounding box.
        """
        # define constant velocity model
        self.kf = KalmanFilter(dim_x=10, dim_z=7)
        self.kf.F = np.array(
            [
                [1, 0, 0, 0, 0, 0, 0, 1, 0, 0],  # state transition matrix
                [0, 1, 0, 0, 0, 0, 0, 0, 1, 0],
                [0, 0, 1, 0, 0, 0, 0, 0, 0, 1],
                [0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
                [0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
                [0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
                [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
                [0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
                [0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
                [0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
            ]
        )

        self.kf.H = np.array(
            [
                [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],  # measurement function,
                [0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
                [0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
                [0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
                [0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
                [0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
                [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
            ]
        )

        # # with angular velocity
        # self.kf = KalmanFilter(dim_x=11, dim_z=7)
        # self.kf.F = np.array([[1,0,0,0,0,0,0,1,0,0,0],      # state transition matrix
        #                       [0,1,0,0,0,0,0,0,1,0,0],
        #                       [0,0,1,0,0,0,0,0,0,1,0],
        #                       [0,0,0,1,0,0,0,0,0,0,1],
        #                       [0,0,0,0,1,0,0,0,0,0,0],
        #                       [0,0,0,0,0,1,0,0,0,0,0],
        #                       [0,0,0,0,0,0,1,0,0,0,0],
        #                       [0,0,0,0,0,0,0,1,0,0,0],
        #                       [0,0,0,0,0,0,0,0,1,0,0],
        #                       [0,0,0,0,0,0,0,0,0,1,0],
        #                       [0,0,0,0,0,0,0,0,0,0,1]])

        # self.kf.H = np.array([[1,0,0,0,0,0,0,0,0,0,0],      # measurement function,
        #                       [0,1,0,0,0,0,0,0,0,0,0],
        #                       [0,0,1,0,0,0,0,0,0,0,0],
        #                       [0,0,0,1,0,0,0,0,0,0,0],
        #                       [0,0,0,0,1,0,0,0,0,0,0],
        #                       [0,0,0,0,0,1,0,0,0,0,0],
        #                       [0,0,0,0,0,0,1,0,0,0,0]])

        # self.kf.R[0:,0:] *= 10.   # measurement uncertainty
        self.kf.P[
            7:, 7:
        ] *= 1000.0  # state uncertainty, give high uncertainty to the unobservable initial velocities, covariance matrix
        self.kf.P *= 10.0

        # self.kf.Q[-1,-1] *= 0.01    # process uncertainty
        self.kf.Q[7:, 7:] *= 0.01
        self.kf.x[:7] = bbox3D.reshape((7, 1))

        self.time_since_update = 0
        self.id = KalmanBoxTracker.count
        KalmanBoxTracker.count += 1
        self.history = []
        self.hits = 1  # number of total hits including the first detection
        self.hit_streak = 1  # number of continuing hit considering the first detection
        self.first_continuing_hit = 1
        self.still_first = True
        self.age = 0

    def update(self, bbox3D):
        """
        Updates the state vector with observed bbox.
        """
        self.time_since_update = 0
        self.history = []
        self.hits += 1
        self.hit_streak += 1  # number of continuing hit
        if self.still_first:
            self.first_continuing_hit += 1  # number of continuing hit in the fist time

        ######################### orientation correction
        if self.kf.x[3] >= np.pi:
            self.kf.x[3] -= np.pi * 2  # make the theta still in the range
        if self.kf.x[3] < -np.pi:
            self.kf.x[3] += np.pi * 2

        new_theta = bbox3D[3]
        if new_theta >= np.pi:
            new_theta -= np.pi * 2  # make the theta still in the range
        if new_theta < -np.pi:
            new_theta += np.pi * 2
        bbox3D[3] = new_theta

        predicted_theta = self.kf.x[3]
        if (
            abs(new_theta - predicted_theta) > np.pi / 2.0
            and abs(new_theta - predicted_theta) < np.pi * 3 / 2.0
        ):  # if the angle of two theta is not acute angle
            self.kf.x[3] += np.pi
            if self.kf.x[3] > np.pi:
                self.kf.x[3] -= np.pi * 2  # make the theta still in the range
            if self.kf.x[3] < -np.pi:
                self.kf.x[3] += np.pi * 2

        # now the angle is acute: < 90 or > 270, convert the case of > 270 to < 90
        if abs(new_theta - self.kf.x[3]) >= np.pi * 3 / 2.0:
            if new_theta > 0:
                self.kf.x[3] += np.pi * 2
            else:
                self.kf.x[3] -= np.pi * 2

        #########################     # flip

        self.kf.update(bbox3D)

        if self.kf.x[3] >= np.pi:
            self.kf.x[3] -= np.pi * 2  # make the theta still in the rage
        if self.kf.x[3] < -np.pi:
            self.kf.x[3] += np.pi * 2

    def predict(self):
        """
        Advances the state vector and returns the predicted bounding box estimate.
        """
        self.kf.predict()
        if self.kf.x[3] >= np.pi:
            self.kf.x[3] -= np.pi * 2
        if self.kf.x[3] < -np.pi:
            self.kf.x[3] += np.pi * 2

        self.age += 1
        if self.time_since_update > 0:
            self.hit_streak = 0
            self.still_first = False
        self.time_since_update += 1
        self.history.append(self.kf.x)
        return self.history[-1]

    def get_state(self):
        """
        Returns the current bounding box estimate.
        """
        return self.kf.x[:7].reshape((7,))

    def get_xyvel(self):
        """
        Returns the current velocity estimate.
        """
        return self.kf.x[-3:-1].reshape((2,))

    def get_pose_covariance(self):
        return self.kf.P[:3, :3]

mask_4d/utils/kalman_filter.py
PRBonn/Mask4D
from typing import List

import numpy as np
import torch
import torch.nn.functional as F
from torch import Tensor


def pad_stack(tensor_list: List[Tensor]):
    """
    pad each tensor on the input to the max value in shape[1] and
    concatenate them in a single tensor.
    Input:
        list of tensors [Ni,Pi]
    Output:
        tensor [sum(Ni),max(Pi)]
    """
    _max = max([t.shape[1] for t in tensor_list])
    # valid = [F.pad(torch.ones_like(t), (0, _max - t.shape[1])) for t in tensor_list]
    batched = torch.cat([F.pad(t, (0, _max - t.shape[1])) for t in tensor_list])
    return batched


def apply_pose(points, pose):
    """
    points: [N,3]
    pose: [4,4]
    """
    hpts = torch.hstack((points[:, :3], torch.ones_like(points[:, :1]))).to(points)
    t_pose = torch.tensor(pose).to(points).T
    tr_pts = torch.mm(hpts, t_pose)
    shifted_pts = tr_pts[:, :3]

    return shifted_pts


def center_from_points(points):
    """
    points: instance points Nx3
    return: 3D center [x,y,z]
    """
    xc = (torch.max(points[:, 0]) + torch.min(points[:, 0])) / 2
    yc = (torch.max(points[:, 1]) + torch.min(points[:, 1])) / 2
    zc = (torch.max(points[:, 2]) + torch.min(points[:, 2])) / 2
    return torch.tensor([xc, yc, zc]).to(points)


def fit_ellipse(points):
    """
    fit a 2D ellipse in the XY plane for a set of 3D points
    points: instance points Nx3
    return: axis sizes [size_x, size_y, size_z], angle
    """
    N = points.shape[0]
    X = points[:, :2]
    X = X - X.mean(0)
    _out = torch.linalg.svd(
        (X - X.mean(axis=0)) / torch.sqrt(torch.tensor(N)), full_matrices=False
    )
    angle = np.arctan2(_out.Vh[0, 1].cpu(), _out.Vh[0, 0].cpu()).cuda()
    zc = (torch.max(points[:, 2]) + torch.min(points[:, 2])) / 2
    dz = torch.cdist(points[:, 2].unsqueeze(1), zc.unsqueeze(0).unsqueeze(1))
    sz = 2 * dz.mean()
    size = torch.cat((_out.S, torch.tensor([sz]).to(_out.S)))
    return size, angle


def getcovs(scale_x, scale_y, theta):
    coss = torch.cos(theta)
    sinn = torch.sin(theta)
    rots = torch.stack((torch.stack((coss, -sinn), 1), torch.stack((sinn, coss), 1)), 1)
    covs = torch.stack(
        (
            torch.stack((scale_x, torch.zeros_like(scale_x)), 1),
            torch.stack((torch.zeros_like(scale_y), scale_y), 1),
        ),
        1,
    )
    covs = rots @ covs @ rots.permute(0, 2, 1)
    return covs


def remove_outliers(points):
    m, _ = torch.median(points, 0)
    d = ((points - m) ** 2).sum(1).sqrt()
    return points[d < torch.median(d) * 1.5 * 2]


def bbox_from_pts(points):
    """
    Runs the loss on outputs of the model
    :param points: instance points Nx3
    :return: 3D bbox [x,y,z,theta,l,w,h]
    """

    x1 = torch.min(points[:, 0])
    x2 = torch.max(points[:, 0])
    y1 = torch.min(points[:, 1])
    y2 = torch.max(points[:, 1])
    z1 = torch.min(points[:, 2])
    z2 = torch.max(points[:, 2])

    return torch.tensor(
        [
            x1 + (x2 - x1) / 2,
            y1 + (y2 - y1) / 2,
            z1 + (z2 - z1) / 2,
            0,
            x2 - x1,
            y2 - y1,
            z2 - z1,
        ]
    ).to(points.device)


def th_delete(tensor, indices):
    if len(indices) == 0:
        return tensor
    # remove elements of a tensor at certain indices
    mask = torch.ones(tensor.shape[0], dtype=torch.bool)
    mask[indices] = False
    return tensor[mask]

mask_4d/utils/misc.py
PRBonn/Mask4D
import os

import numpy as np


def save_results(sem_preds, ins_preds, output_dir, batch, class_inv_lut):
    for i in range(len(sem_preds)):
        sem = sem_preds[i]
        ins = ins_preds[i]
        sem_inv = class_inv_lut[sem].astype(np.uint32)
        label = sem_inv.reshape(-1, 1) + (
            (ins.astype(np.uint32) << 16) & 0xFFFF0000
        ).reshape(-1, 1)

        pcd_path = batch["fname"][i]
        seq = pcd_path.split("/")[-3]
        pcd_fname = pcd_path.split("/")[-1].split(".")[-2] + ".label"
        fname = os.path.join(output_dir, seq, "predictions", pcd_fname)
        label.reshape(-1).astype(np.uint32).tofile(fname)


def full_preds(x, _sem_pred, _ins_pred):
    sem_pred = []
    ins_pred = []
    for i in range(len(x["full_xyz"])):
        sem = np.zeros(x["full_xyz"][i]).astype(np.int64) * 9
        sem[x["keep_xyz"][i]] = _sem_pred[i]
        sem_pred.append(sem)
        ins = np.zeros(x["full_xyz"][i]).astype(np.int64) * 9
        ins[x["keep_xyz"][i]] = _ins_pred[i]
        ins_pred.append(ins)
    return sem_pred, ins_pred

mask_4d/utils/testing.py
PRBonn/Mask4D
# Modified by Rodrigo Marcuzzi from https://github.com/dvlab-research/SphereFormer/blob/master/util/voxelize.py
import numpy as np


def fnv_hash_vec(arr):
    """
    FNV64-1A
    """
    assert arr.ndim == 2
    # Floor first for negative coordinates
    arr = arr.copy()
    arr = arr.astype(np.uint64, copy=False)
    hashed_arr = np.uint64(14695981039346656037) * np.ones(
        arr.shape[0], dtype=np.uint64
    )
    for j in range(arr.shape[1]):
        hashed_arr *= np.uint64(1099511628211)
        hashed_arr = np.bitwise_xor(hashed_arr, arr[:, j])
    return hashed_arr


def ravel_hash_vec(arr):
    """
    Ravel the coordinates after subtracting the min coordinates.
    """
    assert arr.ndim == 2
    arr = arr.copy()
    arr -= arr.min(0)
    arr = arr.astype(np.uint64, copy=False)
    arr_max = arr.max(0).astype(np.uint64) + 1

    keys = np.zeros(arr.shape[0], dtype=np.uint64)
    # Fortran style indexing
    for j in range(arr.shape[1] - 1):
        keys += arr[:, j]
        keys *= arr_max[j + 1]
    keys += arr[:, -1]
    return keys


def voxelize(coord, voxel_size=0.05, hash_type="fnv", mode=0):
    discrete_coord = np.floor(coord / np.array(voxel_size))
    if hash_type == "ravel":
        key = ravel_hash_vec(discrete_coord)
    else:
        key = fnv_hash_vec(discrete_coord)

    idx_sort = np.argsort(key)
    key_sort = key[idx_sort]
    _, count = np.unique(key_sort, return_counts=True)
    if mode == 0:  # train mode
        idx_select = (
            np.cumsum(np.insert(count, 0, 0)[0:-1])
            + np.random.randint(0, count.max(), count.size) % count
        )
        idx_unique = idx_sort[idx_select]
        sorted_idx = np.zeros(key.shape[0]).astype(int)
        sorted_idx[idx_select] = 1
        sorted_idx = np.cumsum(sorted_idx) - 1
        idx_recon = np.zeros(key.shape[0]).astype(int)
        idx_recon[idx_sort] = sorted_idx
        return idx_unique, idx_recon
    else:  # val mode
        idx_start = np.cumsum(np.insert(count, 0, 0)[0:-1])
        idx_select = (
            idx_start  # + np.random.randint(0, count.max(), count.size) % count
        )
        # idx_unique = idx_sort[idx_select]
        sorted_idx = np.zeros(key.shape[0]).astype(int)
        sorted_idx[idx_start] = 1
        sorted_idx = np.cumsum(sorted_idx) - 1
        idx_recon = np.zeros(key.shape[0]).astype(int)
        idx_recon[idx_sort] = sorted_idx
        return idx_recon

mask_4d/utils/voxelize.py
PRBonn/Mask4D
### To not hack your paths
# install via `pip3 install -U -e .` in this dir
# https://stackoverflow.com/questions/6323860/sibling-package-imports/50193944#50193944

from setuptools import find_packages, setup

pkg = "mask_4d"
setup(name=pkg, version="1.0", packages=find_packages())

setup.py
PRBonn/Mask4D