OccupancyDETR: Making Semantic Scene Completion as Straightforward as Object Detection
import os
import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import time

import torch
from datasets import load_dataset
from PIL import Image

from src import OccupancyDetr, OccupancyDetrOutput, OccupancyDetrProcessor
from tools.common import *
from tools.data.semantic_kitti import read_calib
from tools.eval import eval_2d_objs, eval_voxels

seq_path = "/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/08/labels.json"
output_dir = "/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/08/outputs"
gt_dir = "/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/08/restore"
preprocessed_dir = "/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/08/preprocessed"
model_path = "saved_models/base_all_08_19_occ2d_4"
device = torch.device("cuda")
confidence_threshold = 0.5
occupancy_threshold = 0.5
mode = "occupancy"


def prepare_input(example):
    image = Image.open(example["image_path"])
    batch = image_processor([image], return_tensors="pt")
    calib = read_calib(example["calib_path"])
    batch["cam_Es"] = torch.tensor([calib["Tr"].tolist()]).to(device)
    batch["cam_Ks"] = torch.tensor([calib["P2"].tolist()]).to(device)
    batch["pixel_values"] = batch["pixel_values"].to(device)
    batch["pixel_mask"] = batch["pixel_mask"].to(device)
    return {
        "input": batch,
        "target": [],
        "info": example["image_path"],
    }


raw_datasets = load_dataset("json", data_files={"test": [seq_path]})
image_processor = OccupancyDetrProcessor.from_pretrained("SenseTime/deformable-detr")
model = OccupancyDetr.from_pretrained(model_path).to(device)
model.mode = mode
model.eval()
model.config.confidence_threshold = confidence_threshold
model.config.occupancy_threshold = occupancy_threshold
print(model.num_parameters())
image_processor.config = model.config
os.makedirs(output_dir, exist_ok=True)
dts = []
for example in raw_datasets["test"]:
    example = prepare_input(example)
    t1 = time.time()
    with torch.no_grad():
        output: OccupancyDetrOutput = model(**example["input"])
    t2 = time.time()
    print(f"{example['info']}: {t2 - t1}s")
    dts.append(t2 - t1)
    results, voxel = image_processor.post_process_for_OccupancyDetr(output)
    id = example["info"].split("/")[-1].split(".")[0]
    result_path = os.path.join(output_dir, f"{id}.json")
    results["image_path"] = example["info"]
    write_json(result_path, results)
    voxel_path = os.path.join(output_dir, f"{id}.npy")
    np.save(voxel_path, voxel.astype(np.uint16))

print(f"Average time: {np.mean(dts)}s")

AP50, _ = eval_2d_objs(output_dir, preprocessed_dir)
print(f"AP50: {AP50}")

miou = eval_voxels(output_dir, gt_dir)
print(f"IOU: {miou}")
print(f"MIOU: {np.mean(miou)}")

scripts/eval.py
jypjypjypjyp/OccupancyDETR
from .models.OccupancyDetr import OccupancyDetr, OccupancyDetrOutput
from .models.OccupancyDetrConfig import OccupancyDetrConfig
from .models.SpatialTransformer3D import SpatialTransformer3D
from .models.SpatialTransformerBEV import SpatialTransformerBEV
from .readers.DataCollatorForKitti import DataCollatorForKitti
from .utils.OccupancyDetrProcessor import OccupancyDetrProcessor

src/__init__.py
jypjypjypjyp/OccupancyDETR
import torch
import torch.nn as nn
from transformers.models.deformable_detr.modeling_deformable_detr import *

from .OccupancyDetrConfig import OccupancyDetrConfig


class ContrastiveDeNoising(nn.Module):
    def __init__(self, config: OccupancyDetrConfig) -> None:
        super().__init__()
        self.dn_number = config.dn_number
        self.label_noise_ratio = config.dn_label_noise_ratio
        self.box_noise_scale = config.dn_box_noise_scale
        self.num_queries = config.num_queries
        self.num_labels = config.num_labels
        self.d_model = config.d_model
        self.label_enc = nn.Embedding(config.num_labels, config.d_model)

    def forward(self, targets):
        if not self.training or self.dn_number == 0:
            return None, None, None, None

        # positive and negative dn queries
        device = targets[0]["class_labels"].device
        dn_number = self.dn_number * 2
        known = [torch.ones_like(t["class_labels"]) for t in targets]
        batch_size = len(known)
        known_num = [sum(k) for k in known]
        if int(max(known_num)) == 0:
            return None, None, None, None

        dn_number = dn_number // (int(max(known_num) * 2))
        if dn_number == 0:
            dn_number = 1
        unmask_bbox = unmask_label = torch.cat(known)
        labels = torch.cat([t["class_labels"] for t in targets])
        boxes = torch.cat([t["boxes"] for t in targets])
        batch_idx = torch.cat([torch.full_like(t["class_labels"].long(), i) for i, t in enumerate(targets)])

        known_indice = torch.nonzero(unmask_label + unmask_bbox)
        known_indice = known_indice.view(-1)
        known_indice = known_indice.repeat(2 * dn_number, 1).view(-1)
        known_labels = labels.repeat(2 * dn_number, 1).view(-1)
        known_bid = batch_idx.repeat(2 * dn_number, 1).view(-1)
        known_bboxs = boxes.repeat(2 * dn_number, 1)
        known_labels_expaned = known_labels.clone()
        known_bbox_expand = known_bboxs.clone()

        # noise label
        if self.label_noise_ratio > 0:
            p = torch.rand_like(known_labels_expaned.float())
            # half of bbox prob
            chosen_indice = torch.nonzero(p < (self.label_noise_ratio * 0.5)).view(-1)
            # randomly put a new one here
            new_label = torch.randint_like(chosen_indice, 0, self.num_labels)
            known_labels_expaned.scatter_(0, chosen_indice, new_label)
        single_pad = int(max(known_num))

        # noise boxes
        pad_size = int(single_pad * 2 * dn_number)
        positive_idx = torch.arange(0, len(boxes), device=device)
        positive_idx = positive_idx.unsqueeze(0).repeat(dn_number, 1)
        torch.arange(0, dn_number, device=device)
        positive_idx += torch.arange(0, dn_number, device=device).unsqueeze(1) * len(boxes) * 2
        positive_idx = positive_idx.flatten()
        negative_idx = positive_idx + len(boxes)
        if self.box_noise_scale > 0:
            known_bbox_ = torch.zeros_like(known_bboxs)
            known_bbox_[:, :2] = known_bboxs[:, :2] - known_bboxs[:, 2:] / 2
            known_bbox_[:, 2:] = known_bboxs[:, :2] + known_bboxs[:, 2:] / 2

            diff = torch.zeros_like(known_bboxs)
            diff[:, :2] = known_bboxs[:, 2:] / 2
            diff[:, 2:] = known_bboxs[:, 2:] / 2

            rand_sign = torch.randint_like(known_bboxs, low=0, high=2) * 2.0 - 1.0
            rand_part = torch.rand_like(known_bboxs)
            rand_part[negative_idx] += 1.0
            rand_part *= rand_sign
            known_bbox_ = known_bbox_ + torch.mul(rand_part, diff) * self.box_noise_scale
            known_bbox_ = known_bbox_.clamp(min=0.0, max=1.0)
            known_bbox_expand[:, :2] = (known_bbox_[:, :2] + known_bbox_[:, 2:]) / 2
            known_bbox_expand[:, 2:] = known_bbox_[:, 2:] - known_bbox_[:, :2]

        input_label_embed = self.label_enc(known_labels_expaned)
        input_bbox_embed = inverse_sigmoid(known_bbox_expand)

        padding_label = torch.zeros(pad_size, self.d_model).to(device)
        padding_bbox = torch.zeros(pad_size, 4).to(device)

        input_query_label = padding_label.repeat(batch_size, 1, 1)
        input_query_bbox = padding_bbox.repeat(batch_size, 1, 1)

        map_known_indice = torch.tensor([], device=device)
        if len(known_num):
            map_known_indice = torch.cat([torch.arange(0, num) for num in known_num]).to(device)
            map_known_indice = torch.cat([map_known_indice + single_pad * i for i in range(2 * dn_number)]).long()
        if len(known_bid):
            input_query_label[(known_bid, map_known_indice)] = input_label_embed
            input_query_bbox[(known_bid, map_known_indice)] = input_bbox_embed

        # build dn mask
        tgt_size = pad_size + self.num_queries
        dn_mask = torch.ones(tgt_size, tgt_size, device=device) < 0
        # 1. match query cannot see the reconstruct
        dn_mask[pad_size:, :pad_size] = True
        dn_mask[:pad_size, pad_size:] = True
        # 2. reconstruct cannot see each other
        for i in range(dn_number):
            if i == 0:
                dn_mask[single_pad * 2 * i : single_pad * 2 * (i + 1), single_pad * 2 * (i + 1) : pad_size] = True
            if i == dn_number - 1:
                dn_mask[single_pad * 2 * i : single_pad * 2 * (i + 1), : single_pad * i * 2] = True
            else:
                dn_mask[single_pad * 2 * i : single_pad * 2 * (i + 1), single_pad * 2 * (i + 1) : pad_size] = True
                dn_mask[single_pad * 2 * i : single_pad * 2 * (i + 1), : single_pad * 2 * i] = True
        dn_mask = dn_mask.unsqueeze(0).repeat(batch_size, 1, 1)
        # for huggingface attention
        dn_mask2 = torch.zeros(dn_mask.shape, device=device)
        dn_mask2[dn_mask] = float("-inf")

        dn_meta = {
            "pad_size": pad_size,
            "num_dn_group": dn_number,
        }
        return input_query_label, input_query_bbox, dn_mask2, dn_meta

    @staticmethod
    def dn_post_process(outputs_class, outputs_boxes, outputs_boxes3d, dn_meta, aux_loss, _set_aux_loss):
        if dn_meta and dn_meta["pad_size"] > 0:
            output_known_class = outputs_class[:, :, : dn_meta["pad_size"], :]
            output_known_boxes = outputs_boxes[:, :, : dn_meta["pad_size"], :]
            output_known_boxes3d = outputs_boxes3d[:, : dn_meta["pad_size"], :]
            outputs_class = outputs_class[:, :, dn_meta["pad_size"] :, :]
            outputs_boxes = outputs_boxes[:, :, dn_meta["pad_size"] :, :]
            outputs_boxes3d = outputs_boxes3d[:, dn_meta["pad_size"] :, :]
            out = {"logits": output_known_class[-1], "pred_boxes": output_known_boxes[-1], "pred_boxes3d": output_known_boxes3d}
            if aux_loss:
                out["aux_outputs"] = _set_aux_loss(output_known_class, output_known_boxes)
            dn_meta["output_known_lbs_bboxes"] = out
        return outputs_class, outputs_boxes, outputs_boxes3d

src/models/ContrastiveDeNoising.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.modeling_deformable_detr import *

from .OccupancyDetrConfig import OccupancyDetrConfig


def multi_scale_deformable_attention_3d(value, value_spatial_shapes, sampling_locations, attention_weights):
    batch_size, _, num_heads, hidden_dim = value.shape
    _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape
    value_list = value.split([d.item() * w.item() * h.item() for d, w, h in value_spatial_shapes], dim=1)
    sampling_grids = 2 * sampling_locations - 1
    sampling_value_list = []
    for level_id, (d, w, h) in enumerate(value_spatial_shapes):
        # batch_size, height*width, num_heads, hidden_dim
        # -> batch_size, height*width, num_heads*hidden_dim
        # -> batch_size, num_heads*hidden_dim, height*width
        # -> batch_size*num_heads, hidden_dim, height, width
        value_l_ = value_list[level_id].flatten(2).transpose(1, 2).reshape(batch_size * num_heads, hidden_dim, d, w, h)
        # batch_size, num_queries, num_heads, num_points, 2
        # -> batch_size, num_heads, num_queries, num_points, 2
        # -> batch_size*num_heads, 1, num_queries, num_points, 2
        sampling_grid_l_ = sampling_grids[:, :, :, level_id].transpose(1, 2).flatten(0, 1)
        sampling_grid_l_ = sampling_grid_l_.unsqueeze(1)
        # batch_size*num_heads, hidden_dim, 1, num_queries, num_points
        sampling_value_l_ = nn.functional.grid_sample(value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False).squeeze(2)
        sampling_value_list.append(sampling_value_l_)
    # (batch_size, num_queries, num_heads, num_levels, num_points)
    # -> (batch_size, num_heads, num_queries, num_levels, num_points)
    # -> (batch_size, num_heads, 1, num_queries, num_levels*num_points)
    attention_weights = attention_weights.transpose(1, 2).reshape(batch_size * num_heads, 1, num_queries, num_levels * num_points)
    output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(batch_size, num_heads * hidden_dim, num_queries)
    return output.transpose(1, 2).contiguous()


class DeformableAttention3D(nn.Module):
    def __init__(self, config: OccupancyDetrConfig, num_heads: int, n_points: int):
        super().__init__()
        if config.d_model % num_heads != 0:
            raise ValueError(f"embed_dim (d_model) must be divisible by num_heads, but got {config.d_model} and {num_heads}")

        self.d_model = config.d_model
        self.n_levels = config.num_feature_levels
        self.n_heads = num_heads
        self.n_points = n_points

        self.sampling_offsets = nn.Linear(config.d_model, num_heads * self.n_levels * n_points * 3)
        self.attention_weights = nn.Linear(config.d_model, num_heads * self.n_levels * n_points)
        self.value_proj = nn.Linear(config.d_model, config.d_model)
        self.output_proj = nn.Linear(config.d_model, config.d_model)

        self._reset_parameters()

    def _reset_parameters(self):
        nn.init.constant_(self.sampling_offsets.weight.data, 0.0)
        thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
        grid_init = torch.stack([thetas.cos(), thetas.sin(), (thetas.sin() + thetas.cos()) / 2], -1)
        grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 3).repeat(1, self.n_levels, self.n_points, 1)
        for i in range(self.n_points):
            grid_init[:, :, i, :] *= i + 1
        with torch.no_grad():
            self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
        nn.init.constant_(self.attention_weights.weight.data, 0.0)
        nn.init.constant_(self.attention_weights.bias.data, 0.0)
        nn.init.xavier_uniform_(self.value_proj.weight.data)
        nn.init.constant_(self.value_proj.bias.data, 0.0)
        nn.init.xavier_uniform_(self.output_proj.weight.data)
        nn.init.constant_(self.output_proj.bias.data, 0.0)

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

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        position_embeddings: Optional[torch.Tensor] = None,
        reference_points=None,
        spatial_shapes=None,
        level_start_index=None,
        output_attentions: bool = False,
    ):

        # add position embeddings to the hidden states before projecting to queries and keys
        if position_embeddings is not None:
            hidden_states = self.with_pos_embed(hidden_states, position_embeddings)

        batch_size, num_queries, _ = hidden_states.shape
        batch_size, sequence_length, _ = encoder_hidden_states.shape
        if (spatial_shapes[:, 0] * spatial_shapes[:, 1] * spatial_shapes[:, 2]).sum() != sequence_length:
            raise ValueError("Make sure to align the spatial shapes with the sequence length of the encoder hidden states")

        value = self.value_proj(encoder_hidden_states)
        if attention_mask is not None:
            # we invert the attention_mask
            value = value.masked_fill(~attention_mask[..., None], float(0))
        value = value.view(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads)
        sampling_offsets = self.sampling_offsets(hidden_states).view(batch_size, num_queries, self.n_heads, self.n_levels, self.n_points, 3)
        attention_weights = self.attention_weights(hidden_states).view(batch_size, num_queries, self.n_heads, self.n_levels * self.n_points)
        attention_weights = F.softmax(attention_weights, -1).view(batch_size, num_queries, self.n_heads, self.n_levels, self.n_points)
        # batch_size, num_queries, n_heads, n_levels, n_points, 3
        offset_normalizer = torch.stack([spatial_shapes[..., 0], spatial_shapes[..., 1], spatial_shapes[..., 2]], -1)
        sampling_locations = reference_points[:, :, None, :, None, :] + sampling_offsets / offset_normalizer[None, None, None, :, None, :]

        output = multi_scale_deformable_attention_3d(value, spatial_shapes, sampling_locations, attention_weights)
        output = self.output_proj(output)

        return output, attention_weights

src/models/DeformableAttention3D.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.modeling_deformable_detr import *
from transformers.models.detr.modeling_detr import DetrAttention

from . import MLP
from .OccupancyDetrConfig import OccupancyDetrConfig
from .OccupancyDetrPretrainedModel import OccupancyDetrPretrainedModel


class DinoDecoderLayer(nn.Module):
    def __init__(self, config: OccupancyDetrConfig):
        super().__init__()
        self.embed_dim = config.d_model

        # self-attention
        self.self_attn = DetrAttention(
            embed_dim=self.embed_dim,
            num_heads=config.decoder_attention_heads,
            dropout=config.attention_dropout,
        )
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout

        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        # cross-attention
        self.encoder_attn = DeformableDetrMultiscaleDeformableAttention(
            config,
            num_heads=config.decoder_attention_heads,
            n_points=config.decoder_n_points,
        )
        self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        # feedforward neural networks
        self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
        self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = nn.LayerNorm(self.embed_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Optional[torch.Tensor] = None,
        reference_points=None,
        spatial_shapes=None,
        level_start_index=None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        dn_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = False,
    ):
        residual = hidden_states

        # Self Attention
        if dn_mask is not None:
            dn_mask = dn_mask.unsqueeze(1)
        if position_embeddings is None:
            position_embeddings = torch.zeros_like(hidden_states)
        hidden_states, self_attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=dn_mask,
            position_embeddings=position_embeddings,
            output_attentions=output_attentions,
        )

        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        second_residual = hidden_states

        # Cross-Attention
        cross_attn_weights = None
        hidden_states, cross_attn_weights = self.encoder_attn(
            hidden_states=hidden_states,
            attention_mask=encoder_attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            position_embeddings=position_embeddings,
            reference_points=reference_points,
            spatial_shapes=spatial_shapes,
            level_start_index=level_start_index,
            output_attentions=output_attentions,
        )

        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = second_residual + hidden_states

        hidden_states = self.encoder_attn_layer_norm(hidden_states)

        # Fully Connected
        residual = hidden_states
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.final_layer_norm(hidden_states)

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (self_attn_weights, cross_attn_weights)

        return outputs


class DinoDecoder(OccupancyDetrPretrainedModel):
    def __init__(self, config: OccupancyDetrConfig):
        super().__init__(config)

        self.dropout = config.dropout
        self.layers = nn.ModuleList([DinoDecoderLayer(config) for _ in range(config.decoder_layers)])
        self.reference_points_head = MLP(2 * config.d_model, config.d_model, config.d_model, 2)

        # hack implementation for iterative bounding box refinement and two-stage Deformable DETR
        self.bbox_embed = None
        self.class_embed = None

        # Initialize weights and apply final processing
        self.post_init()

    def get_sine_pos_embed(self, proposals):
        """Get the position embedding of the proposals."""

        num_pos_feats = self.config.d_model // 2
        temperature = 10000
        scale = 2 * math.pi

        dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=proposals.device)
        dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats)
        # batch_size, num_queries, 4
        proposals = proposals.sigmoid() * scale
        # batch_size, num_queries, 4, 128
        pos = proposals[:, :, :, None] / dim_t
        # batch_size, num_queries, 4, 64, 2 -> batch_size, num_queries, 512
        pos = torch.stack((pos[:, :, :, 0::2].sin(), pos[:, :, :, 1::2].cos()), dim=4).flatten(2)
        return pos

    def forward(
        self,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        position_embeddings=None,
        reference_points=None,
        spatial_shapes=None,
        level_start_index=None,
        valid_ratios=None,
        output_attentions=None,
        output_hidden_states=None,
        dn_mask=None,
    ):
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states

        if inputs_embeds is not None:
            hidden_states = inputs_embeds

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
        intermediate = ()
        intermediate_reference_points = ()

        for idx, decoder_layer in enumerate(self.layers):
            reference_points_input = reference_points[:, :, None] * torch.cat([valid_ratios, valid_ratios], -1)[:, None]
            query_sine_embed = self.get_sine_pos_embed(reference_points_input[:, :, 0, :])
            query_position_embed = self.reference_points_head(query_sine_embed)

            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            layer_outputs = decoder_layer(
                hidden_states,
                position_embeddings=query_position_embed,
                encoder_hidden_states=encoder_hidden_states,
                reference_points=reference_points_input,
                spatial_shapes=spatial_shapes,
                level_start_index=level_start_index,
                encoder_attention_mask=encoder_attention_mask,
                output_attentions=output_attentions,
                dn_mask=dn_mask,
            )
            hidden_states = layer_outputs[0]

            # hack implementation for iterative bounding box refinement
            if self.bbox_embed is not None:
                tmp = self.bbox_embed[idx](hidden_states)
                new_reference_points = tmp + inverse_sigmoid(reference_points)
                new_reference_points = new_reference_points.sigmoid()
                reference_points = new_reference_points.detach()

            intermediate += (hidden_states,)
            intermediate_reference_points += (reference_points,)

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

                if encoder_hidden_states is not None:
                    all_cross_attentions += (layer_outputs[2],)

        # Keep batch_size as first dimension
        intermediate = torch.stack(intermediate, dim=1)
        intermediate_reference_points = torch.stack(intermediate_reference_points, dim=1)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        return DeformableDetrDecoderOutput(
            last_hidden_state=hidden_states,
            intermediate_hidden_states=intermediate,
            intermediate_reference_points=intermediate_reference_points,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            cross_attentions=all_cross_attentions,
        )

src/models/DinoDecoder.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.modeling_deformable_detr import *


def center_to_corners_format2d(box):
    boxes = torch.cat([box[..., :2] - box[..., 2:] / 2, box[..., :2] + box[..., 2:] / 2], -1)
    boxes = torch.clamp(boxes, 0, 1)
    return boxes


class HungarianMatcher(DeformableDetrHungarianMatcher):
    @torch.no_grad()
    def forward(self, outputs, targets, mask=None):
        batch_size, num_queries = outputs["pred_boxes"].shape[:2]

        # Compute the classification cost.
        class_cost = 0
        if "logits" in outputs:
            alpha = 0.25
            gamma = 2.0
            # [batch_size * num_queries, num_classes]
            out_prob = outputs["logits"].flatten(0, 1).sigmoid()
            target_ids = torch.cat([v["class_labels"] for v in targets])
            neg_cost_class = (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log())
            pos_cost_class = alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log())
            class_cost = pos_cost_class[:, target_ids] - neg_cost_class[:, target_ids]

        # Compute the L1 cost between boxes
        out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]
        target_bbox = torch.cat([v["boxes"] for v in targets])
        bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)

        # Compute the giou cost between boxes
        giou_cost = -generalized_box_iou(center_to_corners_format2d(out_bbox), center_to_corners_format2d(target_bbox))

        # Final cost matrix
        cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
        cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()
        if mask is not None:
            cost_matrix[mask.expand_as(cost_matrix).cpu()] = float("inf")

        sizes = [len(v["boxes"]) for v in targets]
        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]
        return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]

src/models/HungarianMatcher.py
jypjypjypjyp/OccupancyDETR
import torch
import torch.nn as nn
from transformers.models.deformable_detr.modeling_deformable_detr import *

from . import MLP
from .ContrastiveDeNoising import ContrastiveDeNoising
from .HungarianMatcher import HungarianMatcher
from .OccupancyDetrConfig import OccupancyDetrConfig
from .OccupancyDetrLoss import *
from .OccupancyDetrModel import OccupancyDetrModel, OccupancyDetrModelOutput
from .OccupancyDetrPretrainedModel import OccupancyDetrPretrainedModel
from .SpatialTransformer3D import SpatialTransformer3D
from .SpatialTransformerBEV import SpatialTransformerBEV


@dataclass
class OccupancyDetrOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    loss_dict: Optional[Dict] = None
    logits: torch.FloatTensor = None  # all queries logits
    pred_logits: torch.FloatTensor = None  # filtered by confidence_threshold
    boxes: torch.FloatTensor = None  # all queries boxes
    pred_boxes: torch.FloatTensor = None  # filtered by confidence_threshold
    boxes3d: torch.FloatTensor = None  # all queries boxes3d
    pred_boxes3d: torch.FloatTensor = None  # filtered by confidence_threshold
    pred_occ: torch.FloatTensor = None  # filtered by confidence_threshold
    num_preds: torch.LongTensor = None  # number of predictions for each image


class OccupancyDetr(OccupancyDetrPretrainedModel):
    def __init__(self, config: OccupancyDetrConfig, mode: str = "detection"):
        """mode: detection or occupancy"""
        super().__init__(config)
        self.config = config
        self.mode = mode
        self.model = OccupancyDetrModel(config)
        if config.query_mode == "bev":
            self.occ_transformer = SpatialTransformerBEV(config)
        elif config.query_mode == "3d":
            self.occ_transformer = SpatialTransformer3D(config)
        else:
            assert mode == "detection"
            self.occ_transformer = None
        self.cdn = ContrastiveDeNoising(config)
        self.matcher = HungarianMatcher(class_cost=config.class_cost, bbox_cost=config.bbox_cost, giou_cost=config.giou_cost)
        self.criterion = OccupancyDetrLoss(
            matcher=self.matcher,
            losses=["labels", "boxes", "boxes3d", "occupancy"],
            config=config,
        )

        # Detection heads on top
        self.class_embed = nn.Linear(config.d_model, config.num_labels)
        self.bbox_embed = MLP(config.d_model, config.d_model, 4, 3)
        self.bbox3d_head = MLP(config.d_model, config.d_model, 6, 3)

        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        self.class_embed.bias.data = torch.ones(config.num_labels) * bias_value
        nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
        nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)

        # if two-stage, the last class_embed and bbox_embed is for region proposal generation
        num_pred = config.decoder_layers + 1
        self.class_embed = nn.ModuleList([copy.deepcopy(self.class_embed) for i in range(num_pred)])
        self.bbox_embed = nn.ModuleList([copy.deepcopy(self.bbox_embed) for i in range(num_pred)])
        self.model.decoder.bbox_embed = self.bbox_embed
        self.model.decoder.class_embed = self.class_embed
        # hack implementation for two-stage
        for bbox_embed_layer in self.bbox_embed[:-1]:
            nn.init.constant_(bbox_embed_layer.layers[-1].bias.data[2:], 0.0)

        # Initialize weights and apply final processing
        self.post_init()

    def freeze_detection(self):
        for name, param in self.named_parameters():
            param.requires_grad_(False)
        for name, param in self.occ_transformer.named_parameters():
            param.requires_grad_(True)

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

    def forward(
        self,
        cam_Es,
        cam_Ks,
        pixel_values,
        pixel_mask=None,
        decoder_attention_mask=None,
        encoder_outputs=None,
        inputs_embeds=None,
        decoder_inputs_embeds=None,
        labels=None,
        label_boxes=None,
        label_boxes3d=None,
        label_obj3d=None,
        output_attentions=None,
        output_hidden_states=None,
        ids=None,
    ):
        labels_dict = None
        if self.training and labels is not None:
            labels_dict = []
            for cls, bbox, bbox3d, obj3d in zip(labels, label_boxes, label_boxes3d, label_obj3d):
                cls_mask = cls != -1
                labels_dict.append(
                    {
                        "class_labels": cls[cls_mask],
                        "boxes": bbox[cls_mask],
                        "boxes3d": bbox3d[cls_mask],
                        "obj3d": obj3d[cls_mask],
                    }
                )

            dn_query_label, dn_query_bbox, dn_mask, dn_meta = self.cdn(labels_dict)
        else:
            dn_query_label, dn_query_bbox, dn_mask, dn_meta = None, None, None, None

        # First, sent images through DETR base model to obtain encoder + decoder outputs
        early_matching = self.config.early_matching and self.training
        outputs: OccupancyDetrModelOutput = self.model(
            pixel_values,
            pixel_mask=pixel_mask,
            decoder_attention_mask=decoder_attention_mask,
            encoder_outputs=encoder_outputs,
            inputs_embeds=inputs_embeds,
            decoder_inputs_embeds=decoder_inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            dn_query_label=dn_query_label,
            dn_query_bbox=dn_query_bbox,
            dn_mask=dn_mask,
            matcher=self.matcher if early_matching else None,
            labels_dict=labels_dict if early_matching else None,
        )

        hidden_states = outputs.intermediate_hidden_states
        init_reference = outputs.init_reference_points
        inter_references = outputs.intermediate_reference_points
        encoder_last_hidden_state = outputs.encoder_last_hidden_state
        spatial_shapes = outputs.spatial_shapes
        level_start_index = outputs.level_start_index
        valid_ratios = outputs.valid_ratios
        indices = outputs.query_indices

        # class logits + predicted bounding boxes
        outputs_classes = []
        outputs_boxes = []
        for level in range(hidden_states.shape[1]):
            if level == 0:
                reference = init_reference
            else:
                reference = inter_references[:, level - 1]
            reference = inverse_sigmoid(reference)
            outputs_class = self.class_embed[level](hidden_states[:, level])
            delta_bbox = self.bbox_embed[level](hidden_states[:, level])
            if reference.shape[-1] == 4:
                outputs_bbox_logits = delta_bbox + reference
            elif reference.shape[-1] == 2:
                delta_bbox[..., :2] += reference
                outputs_bbox_logits = delta_bbox
            else:
                raise ValueError(f"reference.shape[-1] should be 4 or 2, but got {reference.shape[-1]}")
            outputs_bbox = outputs_bbox_logits.sigmoid()
            outputs_classes.append(outputs_class)
            outputs_boxes.append(outputs_bbox)
        outputs_class = torch.stack(outputs_classes)
        outputs_boxes = torch.stack(outputs_boxes)
        outputs_boxes3d = self.bbox3d_head(hidden_states[:, -1]).sigmoid()
        if self.config.dn_number > 0 and dn_meta is not None:
            outputs_class, outputs_boxes, outputs_boxes3d = self.cdn.dn_post_process(
                outputs_class,
                outputs_boxes,
                outputs_boxes3d,
                dn_meta,
                self.config.auxiliary_loss,
                self._set_aux_loss,
            )
        logits = outputs_class[-1]
        boxes = outputs_boxes[-1]
        boxes3d = outputs_boxes3d
        pred_occ = None
        pred_logits = None
        pred_boxes = None
        pred_boxes3d = None

        loss_dict = {}
        loss = torch.tensor(0.0, device=logits.device)
        num_preds = torch.zeros(logits.shape[0], dtype=torch.long, device=logits.device)
        if self.training and labels is not None:
            # First: extract outputs
            outputs_loss = {"ids": ids}
            outputs_loss["logits"] = logits
            outputs_loss["pred_boxes"] = boxes
            outputs_loss["pred_boxes3d"] = boxes3d
            outputs_loss["dn_meta"] = dn_meta
            outputs_loss["enc_outputs"] = {
                "logits": outputs.enc_outputs_class,
                "pred_boxes": outputs.enc_outputs_coord_logits.sigmoid(),
                "proposal_indices": outputs.proposal_indices,
            }
            if self.config.auxiliary_loss:
                auxiliary_outputs = self._set_aux_loss(outputs_class, outputs_boxes)
                outputs_loss["auxiliary_outputs"] = auxiliary_outputs

            # Second: match
            if indices is None:
                indices = self.matcher(outputs_loss, labels_dict)

            # Third: compute the losses, based on outputs and labels
            self.criterion.to(self.device)

            if self.mode == "occupancy":
                b_idx, idx = self.criterion._get_source_permutation_idx(indices)
                target_b_idx, target_idx = self.criterion._get_target_permutation_idx(indices)
                sample = torch.randperm(idx.nelement())[: self.config.occ_decoder_batch_size]
                b_idx, target_b_idx = b_idx[sample], target_b_idx[sample]
                idx, target_idx = idx[sample], target_idx[sample]
                obj_hidden_states = hidden_states[b_idx, -1, idx]
                obj_boxes3d = boxes3d[b_idx, idx].clone()
                obj_boxes3d[:, 3:] += self.config.larger_boxes
                obj_boxes3d = corners_to_center_format3d(center_to_corners_format3d(obj_boxes3d))
                features2d = encoder_last_hidden_state[b_idx]
                cam_Es = cam_Es[b_idx]
                cam_Ks = cam_Ks[b_idx]
                pred_occ = self.occ_transformer(
                    cam_Es,
                    cam_Ks,
                    obj_hidden_states,
                    obj_boxes3d,
                    features2d,
                    spatial_shapes=spatial_shapes,
                    level_start_index=level_start_index,
                    valid_ratios=valid_ratios[b_idx],
                )
                outputs_loss["pred_occ"] = pred_occ
                outputs_loss["enlarged_boxes3d"] = obj_boxes3d
                outputs_loss["occ_idx2target_idx"] = (target_b_idx, target_idx)

            loss_dict = self.criterion(outputs_loss, labels_dict, indices)

            # Fourth: compute total loss, as a weighted sum of the various losses
            loss = self.criterion.sum_loss(loss_dict)
        
        else:
            confidences = logits.sigmoid().max(dim=-1).values
            indices = (confidences > self.config.confidence_threshold).nonzero()
            b_idx, idx = indices[:, 0], indices[:, 1]
            for i in range(len(b_idx)):
                num_preds[b_idx[i]] += 1
            obj_hidden_states = hidden_states[:, -1][b_idx, idx]
            pred_logits = logits[b_idx, idx]
            pred_boxes = boxes[b_idx, idx]
            pred_boxes3d = boxes3d[b_idx, idx]
            obj_boxes3d = pred_boxes3d.clone()
            obj_boxes3d[:, 3:] += self.config.larger_boxes
            obj_boxes3d = corners_to_center_format3d(center_to_corners_format3d(obj_boxes3d))
            if self.mode == "occupancy":
                occs = []
                batch_size = self.config.occ_decoder_batch_size
                for i in range(0, len(b_idx), batch_size):
                    b_idx_batch = b_idx[i : i + batch_size]
                    obj_hidden_states_batch = obj_hidden_states[i : i + batch_size]
                    obj_boxes3d_batch = obj_boxes3d[i : i + batch_size]
                    features2d_batch = encoder_last_hidden_state[b_idx_batch]
                    cam_Es_batch = cam_Es[b_idx_batch]
                    cam_Ks_batch = cam_Ks[b_idx_batch]
                    occupancy_batch = self.occ_transformer(
                        cam_Es_batch,
                        cam_Ks_batch,
                        obj_hidden_states_batch,
                        obj_boxes3d_batch,
                        features2d_batch,
                        spatial_shapes=spatial_shapes,
                        level_start_index=level_start_index,
                        valid_ratios=valid_ratios[b_idx_batch],
                    )
                    occs_batch = occupancy_batch.detach().cpu()
                    del occupancy_batch
                    occs.append(occs_batch)
                pred_occ = torch.cat(occs, dim=0)

        return OccupancyDetrOutput(
            loss=loss,
            loss_dict=loss_dict,
            logits=logits,
            pred_logits=pred_logits,
            boxes=boxes,
            pred_boxes=pred_boxes,
            boxes3d=boxes3d,
            pred_boxes3d=pred_boxes3d,
            pred_occ=pred_occ,
            num_preds=num_preds,
        )

src/models/OccupancyDetr.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.configuration_deformable_detr import DeformableDetrConfig


class OccupancyDetrConfig(DeformableDetrConfig):
    def __init__(
        self,
        # Dataset parameters
        voxel_shape=(256, 256, 32),
        voxel_size=0.2,
        voxel_origin=(0, -25.6, -2),
        image_width=1220,
        image_height=370,
        # OccupancyDetr parameters
        mixed_query_selection=False,
        query_mode="bev",  # "bev" or "3d"
        query_shape=(32, 32, 8),
        larger_boxes=0.05,
        confidence_threshold=0.5,
        occupancy_threshold=0.5,
        occ_decoder_layers=6,
        occ_decoder_batch_size=32,
        # Training parameters
        early_matching=True,
        only_encoder_loss=False,
        dn_number=0,
        dn_box_noise_scale=0.4,
        dn_label_noise_ratio=0.5,
        weight_dict = {
            "ce": 1,
            "bbox": 5,
            # "bbox3d": 5,
            # "giou": 2,
            # "giou3d": 5,
            # "dice3d": 2,
        },
        # original DeformableDetrConfig parameters
        use_timm_backbone=True,
        backbone_config=None,
        num_channels=3,
        num_queries=900,
        encoder_layers=6,
        encoder_ffn_dim=1024,
        encoder_attention_heads=8,
        decoder_layers=6,
        decoder_ffn_dim=1024,
        decoder_attention_heads=8,
        is_encoder_decoder=True,
        activation_function="relu",
        d_model=256,
        position_embedding_type="sine",
        backbone="resnet50",
        use_pretrained_backbone=True,
        num_feature_levels=4,
        encoder_n_points=4,
        decoder_n_points=4,
        **kwargs,
    ):
        super().__init__(
            use_timm_backbone=use_timm_backbone,
            backbone_config=backbone_config,
            num_channels=num_channels,
            num_queries=num_queries,
            encoder_layers=encoder_layers,
            encoder_ffn_dim=encoder_ffn_dim,
            encoder_attention_heads=encoder_attention_heads,
            decoder_layers=decoder_layers,
            decoder_ffn_dim=decoder_ffn_dim,
            decoder_attention_heads=decoder_attention_heads,
            is_encoder_decoder=is_encoder_decoder,
            activation_function=activation_function,
            d_model=d_model,
            position_embedding_type=position_embedding_type,
            backbone=backbone,
            use_pretrained_backbone=use_pretrained_backbone,
            num_feature_levels=num_feature_levels,
            encoder_n_points=encoder_n_points,
            decoder_n_points=decoder_n_points,
            **kwargs,
        )

        self.voxel_shape = voxel_shape
        self.voxel_size = voxel_size
        self.voxel_origin = voxel_origin
        self.image_width = image_width
        self.image_height = image_height
        self.mixed_query_selection = mixed_query_selection
        self.query_mode = query_mode
        self.query_shape = query_shape
        self.larger_boxes = larger_boxes
        self.confidence_threshold = confidence_threshold
        self.occupancy_threshold = occupancy_threshold
        self.occ_decoder_layers = occ_decoder_layers
        self.occ_decoder_batch_size = occ_decoder_batch_size
        self.early_matching = early_matching
        self.only_encoder_loss = only_encoder_loss
        self.dn_number = dn_number
        self.dn_box_noise_scale = dn_box_noise_scale
        self.dn_label_noise_ratio = dn_label_noise_ratio
        self.weight_dict = weight_dict
        self.pad_token_id = 0
        self.dropout = 0
        self.encoder_layerdrop = 0
        self.attention_dropout = 0
        self.activation_dropout = 0
        self.two_stage = True
        self.with_box_refine = True
        self.auxiliary_loss = True

src/models/OccupancyDetrConfig.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.modeling_deformable_detr import *

from .OccupancyDetrConfig import OccupancyDetrConfig


def center_to_corners_format3d(box):
    boxes = torch.cat([box[..., :3] - box[..., 3:] / 2, box[..., :3] + box[..., 3:] / 2], -1)
    boxes = torch.clamp(boxes, 0, 1)
    return boxes


def corners_to_center_format3d(box):
    boxes = torch.cat([(box[..., :3] + box[..., 3:]) / 2, box[..., 3:] - box[..., :3]], -1)
    boxes = torch.clamp(boxes, 0, 1)
    return boxes


def generalized_box_iou3d(bbox_preds, bbox_targets):
    bbox_preds = bbox_preds.unsqueeze(1)
    bbox_targets = bbox_targets.unsqueeze(0)
    intersect_mins = torch.max(bbox_preds[..., :3], bbox_targets[..., :3])
    intersect_maxs = torch.min(bbox_preds[..., 3:], bbox_targets[..., 3:])
    intersect_whd = torch.clamp(intersect_maxs - intersect_mins, min=0)
    intersect_vol = intersect_whd[..., 0] * intersect_whd[..., 1] * intersect_whd[..., 2]
    bbox_preds_vol = (bbox_preds[..., 3] - bbox_preds[..., 0]) * (bbox_preds[..., 4] - bbox_preds[..., 1]) * (bbox_preds[..., 5] - bbox_preds[..., 2])
    bbox_targets_vol = (bbox_targets[..., 3] - bbox_targets[..., 0]) * (bbox_targets[..., 4] - bbox_targets[..., 1]) * (bbox_targets[..., 5] - bbox_targets[..., 2])
    union_vol = bbox_preds_vol + bbox_targets_vol - intersect_vol
    iou = intersect_vol / torch.clamp(union_vol, min=1e-6)
    enclose_mins = torch.min(bbox_preds[..., :3], bbox_targets[..., :3])
    enclose_maxs = torch.max(bbox_preds[..., 3:], bbox_targets[..., 3:])
    enclose_whd = torch.clamp(enclose_maxs - enclose_mins, min=0)
    enclose_vol = enclose_whd[..., 0] * enclose_whd[..., 1] * enclose_whd[..., 2]
    return iou - (enclose_vol - union_vol) / torch.clamp(enclose_vol, min=1e-6)


def dice_loss3d(pred, target):
    pred = pred.contiguous().flatten(1)
    target = target.contiguous().flatten(1)
    intersection = 2 * (pred * target).sum(1)
    union = pred.sum(1) + target.sum(1)
    loss = 1 - ((intersection + 1) / (union + 1))
    return loss


class OccupancyDetrLoss(DeformableDetrLoss):
    def __init__(self, matcher, losses, config: OccupancyDetrConfig):
        super().__init__(matcher, config.num_labels, config.focal_alpha, losses)
        self.config = config

    def prep_for_dn(self, dn_meta):
        output_known_lbs_bboxes = dn_meta["output_known_lbs_bboxes"]
        num_dn_groups, pad_size = dn_meta["num_dn_group"], dn_meta["pad_size"]
        assert pad_size % num_dn_groups == 0
        single_pad = pad_size // num_dn_groups
        return output_known_lbs_bboxes, single_pad, num_dn_groups

    def loss_boxes3d(self, outputs, targets, indices, num_boxes):
        if "pred_boxes3d" not in outputs:
            return {}
        idx = self._get_source_permutation_idx(indices)
        source_boxes = outputs["pred_boxes3d"][idx]
        target_boxes = torch.cat([t["boxes3d"][i] for t, (_, i) in zip(targets, indices)], dim=0)

        losses = {}
        loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="none")
        losses["loss_bbox3d"] = loss_bbox.sum() / num_boxes
        loss_giou = 1 - torch.diag(generalized_box_iou3d(center_to_corners_format3d(source_boxes), center_to_corners_format3d(target_boxes)))
        losses["loss_giou3d"] = loss_giou.sum() / num_boxes
        return losses

    def loss_occ(self, outputs, targets, indices, num_boxes):
        if "pred_occ" not in outputs:
            return {}
        tb, ti = outputs["occ_idx2target_idx"]
        source_occs_in_box = outputs["pred_occ"]
        source_boxes = outputs["enlarged_boxes3d"]
        target_occs_in_box = torch.stack([targets[b]["obj3d"][i] for b, i in zip(tb, ti)]).float()
        target_boxes = torch.stack([targets[b]["boxes3d"][i] for b, i in zip(tb, ti)])

        target_occs_in_pred_box = []
        for occ, box, tocc, tbox in zip(source_occs_in_box, source_boxes, target_occs_in_box, target_boxes):
            x, y, z, d, w, h = tbox
            x, y, z = int(x * 256), int(y * 256), int(z * 32)
            d, w, h = int(d * 256), int(w * 256), int(h * 32)
            minx = max(0, x - d // 2)
            miny = max(0, y - w // 2)
            minz = max(0, z - h // 2)
            w = max(1, min(w, 256 - minx))
            d = max(1, min(d, 256 - miny))
            h = max(1, min(h, 32 - minz))
            tworld = torch.zeros(self.config.voxel_shape, device=tocc.device)
            tocc_in_box = F.interpolate(tocc[None, None, :, :, :], size=(d, w, h))[0, 0]
            tworld[minx : minx + d, miny : miny + w, minz : minz + h] = tocc_in_box

            x, y, z, d, w, h = box
            x, y, z = int(x * 256), int(y * 256), int(z * 32)
            d, w, h = int(d * 256), int(w * 256), int(h * 32)
            minx = max(0, x - d // 2)
            miny = max(0, y - w // 2)
            minz = max(0, z - h // 2)
            w = max(1, min(w, 256 - minx))
            d = max(1, min(d, 256 - miny))
            h = max(1, min(h, 32 - minz))
            tocc_in_box = tworld[minx : minx + d, miny : miny + w, minz : minz + h]
            tocc_in_box = F.interpolate(tocc_in_box[None, None, :, :, :], size=occ.shape)[0, 0]
            target_occs_in_pred_box.append(tocc_in_box)
        target_occs_in_pred_box = torch.stack(target_occs_in_pred_box)

        losses = {}
        loss_dice3d = dice_loss3d(source_occs_in_box, target_occs_in_pred_box)
        losses["loss_dice3d"] = loss_dice3d.sum() / num_boxes
        return losses

    def get_loss(self, loss, outputs, targets, indices, num_boxes):
        loss_map = {
            "labels": self.loss_labels,
            "cardinality": self.loss_cardinality,
            "boxes": self.loss_boxes,
            "boxes3d": self.loss_boxes3d,
            "occupancy": self.loss_occ,
        }
        if loss not in loss_map:
            raise ValueError(f"Loss {loss} not supported")
        return loss_map[loss](outputs, targets, indices, num_boxes)

    def forward(self, outputs, targets, indices):
        device = outputs["logits"].device

        # Compute the average number of target boxes accross all nodes, for normalization purposes
        num_boxes = sum(len(t["class_labels"]) for t in targets)
        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=device)
        num_boxes = torch.clamp(num_boxes, min=1).item()

        losses = {}
        # For two stage, compute losses for proposals
        if "enc_outputs" in outputs:
            enc_outputs = outputs["enc_outputs"]
            proposal_indices = enc_outputs["proposal_indices"]
            if proposal_indices is None:
                proposal_indices = self.matcher(enc_outputs, targets)
            bin_targets = copy.deepcopy(targets)
            for bt in bin_targets:
                bt["class_labels"] = torch.zeros_like(bt["class_labels"])
            for loss in self.losses:
                l_dict = self.get_loss(loss, enc_outputs, bin_targets, proposal_indices, num_boxes)
                l_dict = {k + "_enc": v for k, v in l_dict.items()}
                losses.update(l_dict)
            if self.config.only_encoder_loss:
                return losses

        # Compute all the requested losses
        for loss in self.losses:
            losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))

        # Compute DN losses
        dn_meta = outputs["dn_meta"]
        if self.training and dn_meta and "output_known_lbs_bboxes" in dn_meta:
            output_known_lbs_bboxes, single_pad, scalar = self.prep_for_dn(dn_meta)

            dn_pos_idx = []
            dn_neg_idx = []
            for i in range(len(targets)):
                if len(targets[i]["class_labels"]) > 0:
                    t = torch.arange(0, len(targets[i]["class_labels"]) - 1, device=device)
                    t = t.unsqueeze(0).repeat(scalar, 1)
                    tgt_idx = t.flatten()
                    output_idx = torch.arange(0, scalar, device=device).unsqueeze(1) * single_pad + t
                    output_idx = output_idx.flatten()
                else:
                    output_idx = tgt_idx = torch.tensor([], dtype=torch.long, device=device)
                dn_pos_idx.append((output_idx, tgt_idx))
                dn_neg_idx.append((output_idx + single_pad // 2, tgt_idx))

            l_dict = {}
            for loss in self.losses:
                l_dict.update(self.get_loss(loss, output_known_lbs_bboxes, targets, dn_pos_idx, num_boxes * scalar))

            l_dict = {k + f"_dn": v for k, v in l_dict.items()}
            losses.update(l_dict)

        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
        if "auxiliary_outputs" in outputs:
            for i, auxiliary_outputs in enumerate(outputs["auxiliary_outputs"]):
                for loss in self.losses:
                    l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes)
                    l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
                    losses.update(l_dict)

                if self.training and dn_meta and "output_known_lbs_bboxes" in dn_meta:
                    aux_outputs_known = output_known_lbs_bboxes["aux_outputs"][i]
                    l_dict = {}
                    for loss in self.losses:
                        l_dict.update(self.get_loss(loss, aux_outputs_known, targets, dn_pos_idx, num_boxes * scalar))
                    l_dict = {k + f"_dn_{i}": v for k, v in l_dict.items()}
                    losses.update(l_dict)

        return losses

    def sum_loss(self, loss_dict):
        w_losses = []
        for k, v in loss_dict.items():
            tag, loss_k = k.split("_")[:2]
            if tag == "loss" and loss_k in self.config.weight_dict:
                w_losses.append(self.config.weight_dict[loss_k] * v)
        loss = sum(w_losses) 
        return loss
src/models/OccupancyDetrLoss.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.modeling_deformable_detr import *

from .DinoDecoder import DinoDecoder
from .OccupancyDetrConfig import OccupancyDetrConfig
from .OccupancyDetrPretrainedModel import OccupancyDetrPretrainedModel


@dataclass
class OccupancyDetrModelOutput(DeformableDetrModelOutput):
    spatial_shapes: Optional[torch.LongTensor] = None
    level_start_index: Optional[torch.LongTensor] = None
    valid_ratios: Optional[torch.FloatTensor] = None
    output_proposals: Optional[torch.FloatTensor] = None
    query_indices: Optional[List] = None
    proposal_indices: Optional[List] = None


class OccupancyDetrModel(OccupancyDetrPretrainedModel):
    def __init__(self, config: OccupancyDetrConfig):
        super().__init__(config)
        self.config = config

        # Create backbone + positional encoding
        backbone = DeformableDetrConvEncoder(config)
        position_embeddings = build_position_encoding(config)
        self.backbone = DeformableDetrConvModel(backbone, position_embeddings)

        # Create input projection layers
        if config.num_feature_levels > 1:
            num_backbone_outs = len(backbone.intermediate_channel_sizes)
            input_proj_list = []
            for _ in range(num_backbone_outs):
                in_channels = backbone.intermediate_channel_sizes[_]
                input_proj_list.append(
                    nn.Sequential(
                        nn.Conv2d(in_channels, config.d_model, kernel_size=1),
                        nn.GroupNorm(32, config.d_model),
                    )
                )
            for _ in range(config.num_feature_levels - num_backbone_outs):
                input_proj_list.append(
                    nn.Sequential(
                        nn.Conv2d(in_channels, config.d_model, kernel_size=3, stride=2, padding=1),
                        nn.GroupNorm(32, config.d_model),
                    )
                )
                in_channels = config.d_model
            self.input_proj = nn.ModuleList(input_proj_list)
        else:
            self.input_proj = nn.ModuleList(
                [
                    nn.Sequential(
                        nn.Conv2d(backbone.intermediate_channel_sizes[-1], config.d_model, kernel_size=1),
                        nn.GroupNorm(32, config.d_model),
                    )
                ]
            )

        self.encoder = DeformableDetrEncoder(config)
        self.decoder = DinoDecoder(config)

        self.level_embed = nn.Parameter(torch.Tensor(config.num_feature_levels, config.d_model))

        self.enc_output = nn.Linear(config.d_model, config.d_model)
        self.enc_output_norm = nn.LayerNorm(config.d_model)
        self.query_embeddings = nn.Embedding(config.num_queries, config.d_model)
        self.reference_points = nn.Linear(config.d_model, 4)

        self.post_init()

    def freeze_backbone(self):
        for name, param in self.backbone.conv_encoder.model.named_parameters():
            param.requires_grad_(False)

    def unfreeze_backbone(self):
        for name, param in self.backbone.conv_encoder.model.named_parameters():
            param.requires_grad_(True)

    def get_valid_ratio(self, mask):
        """Get the valid ratio of all feature maps."""

        _, height, width = mask.shape
        valid_height = torch.sum(mask[:, :, 0], 1)
        valid_width = torch.sum(mask[:, 0, :], 1)
        valid_ratio_heigth = valid_height.float() / height
        valid_ratio_width = valid_width.float() / width
        valid_ratio = torch.stack([valid_ratio_width, valid_ratio_heigth], -1)
        return valid_ratio

    def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes):
        """Generate the encoder output proposals from encoded enc_output."""

        batch_size = enc_output.shape[0]
        proposals = []
        _cur = 0
        for level, (height, width) in enumerate(spatial_shapes):
            mask_flatten_ = padding_mask[:, _cur : (_cur + height * width)].view(batch_size, height, width, 1)
            valid_height = torch.sum(~mask_flatten_[:, :, 0, 0], 1)
            valid_width = torch.sum(~mask_flatten_[:, 0, :, 0], 1)

            grid_y, grid_x = meshgrid(
                torch.linspace(0, height - 1, height, dtype=torch.float32, device=enc_output.device),
                torch.linspace(0, width - 1, width, dtype=torch.float32, device=enc_output.device),
                indexing="ij",
            )
            grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1)

            scale = torch.cat([valid_width.unsqueeze(-1), valid_height.unsqueeze(-1)], 1).view(batch_size, 1, 1, 2)
            grid = (grid.unsqueeze(0).expand(batch_size, -1, -1, -1) + 0.5) / scale
            width_heigth = torch.ones_like(grid) * 0.05 * (2.0**level)
            proposal = torch.cat((grid, width_heigth), -1).view(batch_size, -1, 4)
            proposals.append(proposal)
            _cur += height * width
        output_proposals = torch.cat(proposals, 1)
        output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True)
        output_proposals = torch.log(output_proposals / (1 - output_proposals))  # inverse sigmoid
        proposals_mask = padding_mask.unsqueeze(-1) | ~output_proposals_valid
        output_proposals = output_proposals.masked_fill(proposals_mask, float("-inf"))

        # assign each pixel as an object query
        object_query = enc_output
        object_query = object_query.masked_fill(proposals_mask, float(0))
        object_query = self.enc_output_norm(self.enc_output(object_query))
        return object_query, output_proposals, proposals_mask

    def forward(
        self,
        pixel_values,
        pixel_mask=None,
        decoder_attention_mask=None,
        encoder_outputs=None,
        inputs_embeds=None,
        decoder_inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        dn_query_label=None,
        dn_query_bbox=None,
        dn_mask=None,
        matcher=None,
        labels_dict=None,
    ):
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states

        batch_size, num_channels, height, width = pixel_values.shape
        device = pixel_values.device

        if pixel_mask is None:
            pixel_mask = torch.ones(((batch_size, height, width)), dtype=torch.long, device=device)

        # Extract multi-scale feature maps of same resolution `config.d_model` (cf Figure 4 in paper)
        # First, sent pixel_values + pixel_mask through Backbone to obtain the features
        features, position_embeddings_list = self.backbone(pixel_values, pixel_mask)

        # Then, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
        sources = []
        masks = []
        for level, (source, mask) in enumerate(features):
            sources.append(self.input_proj[level](source))
            masks.append(mask)
            if mask is None:
                raise ValueError("No attention mask was provided")

        # Lowest resolution feature maps are obtained via 3x3 stride 2 convolutions on the final stage
        if self.config.num_feature_levels > len(sources):
            _len_sources = len(sources)
            for level in range(_len_sources, self.config.num_feature_levels):
                if level == _len_sources:
                    source = self.input_proj[level](features[-1][0])
                else:
                    source = self.input_proj[level](sources[-1])
                mask = nn.functional.interpolate(pixel_mask[None].float(), size=source.shape[-2:]).to(torch.bool)[0]
                pos_l = self.backbone.position_embedding(source, mask).to(source.dtype)
                sources.append(source)
                masks.append(mask)
                position_embeddings_list.append(pos_l)

        # Prepare encoder inputs (by flattening)
        source_flatten = []
        mask_flatten = []
        lvl_pos_embed_flatten = []
        spatial_shapes = []
        for level, (source, mask, pos_embed) in enumerate(zip(sources, masks, position_embeddings_list)):
            batch_size, num_channels, height, width = source.shape
            spatial_shape = (height, width)
            spatial_shapes.append(spatial_shape)
            source = source.flatten(2).transpose(1, 2)
            mask = mask.flatten(1)
            pos_embed = pos_embed.flatten(2).transpose(1, 2)
            lvl_pos_embed = pos_embed + self.level_embed[level].view(1, 1, -1)
            lvl_pos_embed_flatten.append(lvl_pos_embed)
            source_flatten.append(source)
            mask_flatten.append(mask)
        source_flatten = torch.cat(source_flatten, 1)
        mask_flatten = torch.cat(mask_flatten, 1)
        lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
        spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=source_flatten.device)
        level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
        valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
        valid_ratios = valid_ratios.float()

        # Transformer encoder
        encoder_outputs: BaseModelOutput = self.encoder(
            inputs_embeds=source_flatten,
            attention_mask=mask_flatten,
            position_embeddings=lvl_pos_embed_flatten,
            spatial_shapes=spatial_shapes,
            level_start_index=level_start_index,
            valid_ratios=valid_ratios,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        # Prepare decoder inputs
        output_memory, output_proposals, proposals_mask = self.gen_encoder_output_proposals(encoder_outputs.last_hidden_state, ~mask_flatten, spatial_shapes)
        enc_outputs_class = self.decoder.class_embed[-1](output_memory)
        enc_outputs_coord_logits = self.decoder.bbox_embed[-1](output_memory) + output_proposals

        indices = None
        query_indices = None
        if labels_dict is not None and matcher is not None:
            outputs = {"pred_boxes": output_proposals.sigmoid()}
            indices = matcher(outputs, labels_dict, proposals_mask)
            query_indices = []
            topk_proposals = []
            enc_outputs_class_detach = enc_outputs_class.detach().clone()
            for b, (b_source_idx, b_target_idx) in enumerate(indices):
                left_topk = self.config.num_queries - b_source_idx.shape[0]
                enc_outputs_class_detach[b, b_source_idx, :] = float("-inf")
                left_topk_proposals = torch.topk(enc_outputs_class_detach[b, :, 0], left_topk, dim=0)[1]
                topk_proposals.append(torch.cat([b_source_idx.to(device), left_topk_proposals], dim=0))
                new_source_idx = torch.arange(b_source_idx.shape[0], device=device)
                query_indices.append((new_source_idx, b_target_idx))
            topk_proposals = torch.stack(topk_proposals, dim=0)
        else:
            topk = self.config.num_queries
            topk_proposals = torch.topk(enc_outputs_class[..., 0], topk, dim=1)[1]

        topk_coords_unact = torch.gather(enc_outputs_coord_logits, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4)).detach()
        reference_points = topk_coords_unact.sigmoid()
        if self.config.mixed_query_selection:
            query = self.query_embeddings.weight.unsqueeze(0).repeat(batch_size, 1, 1)
        else:
            query = torch.gather(output_memory, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, output_memory.shape[-1])).detach()
        if dn_query_bbox is not None and dn_query_label is not None:
            reference_points = torch.cat([dn_query_bbox.sigmoid(), reference_points], 1)
            query = torch.cat([dn_query_label, query], dim=1)
        init_reference_points = reference_points

        # Transformer decoder
        decoder_outputs: DeformableDetrDecoderOutput = self.decoder(
            inputs_embeds=query,
            encoder_hidden_states=encoder_outputs[0],
            encoder_attention_mask=mask_flatten,
            reference_points=reference_points,
            spatial_shapes=spatial_shapes,
            level_start_index=level_start_index,
            valid_ratios=valid_ratios,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            dn_mask=dn_mask,
        )

        return OccupancyDetrModelOutput(
            init_reference_points=init_reference_points,
            last_hidden_state=decoder_outputs.last_hidden_state,
            intermediate_hidden_states=decoder_outputs.intermediate_hidden_states,
            intermediate_reference_points=decoder_outputs.intermediate_reference_points,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
            enc_outputs_class=enc_outputs_class,
            enc_outputs_coord_logits=enc_outputs_coord_logits,
            output_proposals=output_proposals,
            spatial_shapes=spatial_shapes,
            level_start_index=level_start_index,
            valid_ratios=valid_ratios,
            query_indices=query_indices,
            proposal_indices=indices,
        )

src/models/OccupancyDetrModel.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.modeling_deformable_detr import *

from .OccupancyDetrConfig import OccupancyDetrConfig


class OccupancyDetrPretrainedModel(DeformableDetrPreTrainedModel):
    config_class = OccupancyDetrConfig

src/models/OccupancyDetrPretrainedModel.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.modeling_deformable_detr import *

from .DeformableAttention3D import DeformableAttention3D
from .OccupancyDetrConfig import OccupancyDetrConfig
from .OccupancyDetrPretrainedModel import OccupancyDetrPretrainedModel


class SpatialDecoderLayer(nn.Module):
    def __init__(self, config: OccupancyDetrConfig):
        super().__init__()

        d_model = config.d_model
        self_config = copy.deepcopy(config)
        self_config.num_feature_levels = 1
        self.self_attn = DeformableAttention3D(self_config, num_heads=config.encoder_attention_heads, n_points=config.encoder_n_points)
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout

        self.self_attn_layer_norm = nn.LayerNorm(d_model)
        self.encoder_attn = DeformableDetrMultiscaleDeformableAttention(config, num_heads=config.encoder_attention_heads, n_points=config.encoder_n_points)
        self.encoder_attn_layer_norm = nn.LayerNorm(d_model)
        self.fc1 = nn.Linear(d_model, config.decoder_ffn_dim)
        self.fc2 = nn.Linear(config.decoder_ffn_dim, d_model)
        self.final_layer_norm = nn.LayerNorm(d_model)

    def forward(
        self,
        hidden_states: torch.Tensor,
        query_position_embeddings: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        reference_points_3d: torch.Tensor,
        spatial_shapes_3d: torch.Tensor,
        level_start_index_3d: torch.Tensor,
        reference_points: torch.Tensor,
        spatial_shapes: torch.Tensor,
        level_start_index: torch.Tensor,
    ):
        # Self Attention
        residual = hidden_states

        hidden_states, _ = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=query_position_embeddings,
            encoder_hidden_states=hidden_states,
            reference_points=reference_points_3d,
            spatial_shapes=spatial_shapes_3d,
            level_start_index=level_start_index_3d,
        )

        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        # Cross Attention
        residual = hidden_states

        hidden_states, _ = self.encoder_attn(
            hidden_states=hidden_states,
            position_embeddings=query_position_embeddings,
            encoder_hidden_states=encoder_hidden_states,
            reference_points=reference_points,
            spatial_shapes=spatial_shapes,
            level_start_index=level_start_index,
        )

        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.encoder_attn_layer_norm(hidden_states)

        # Fully Connected
        residual = hidden_states
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.final_layer_norm(hidden_states)

        return hidden_states


class VoxelLearnedPositionEmbedding(nn.Module):
    def __init__(self, embedding_dim, voxel_shape, query_shape):
        super().__init__()
        self.voxel_shape = voxel_shape
        self.query_shape = query_shape  # d, w, h
        each_dim = embedding_dim // 3
        self.x_embeddings = nn.Embedding(voxel_shape[0], each_dim)
        self.y_embeddings = nn.Embedding(voxel_shape[1], each_dim)
        self.z_embeddings = nn.Embedding(voxel_shape[2], embedding_dim - 2 * each_dim)

    def forward(self, boxes3d):
        coords = torch.ones(self.query_shape, device=boxes3d.device).nonzero()
        coords = coords.unsqueeze(0).repeat(boxes3d.shape[0], 1, 1)
        xs = boxes3d[:, 0] * self.voxel_shape[0]
        ys = boxes3d[:, 1] * self.voxel_shape[1]
        zs = boxes3d[:, 2] * self.voxel_shape[2]
        ds = boxes3d[:, 3] * self.voxel_shape[0] * (1 - 1 / self.query_shape[0])
        ws = boxes3d[:, 4] * self.voxel_shape[1] * (1 - 1 / self.query_shape[1])
        hs = boxes3d[:, 5] * self.voxel_shape[2] * (1 - 1 / self.query_shape[2])
        minxs = torch.max(torch.zeros_like(xs), xs - ds / 2)
        minys = torch.max(torch.zeros_like(ys), ys - ws / 2)
        minzs = torch.max(torch.zeros_like(zs), zs - hs / 2)
        ds = torch.max(torch.ones_like(ds), torch.min(ds, self.voxel_shape[0] - minxs))
        ws = torch.max(torch.ones_like(ws), torch.min(ws, self.voxel_shape[1] - minys))
        hs = torch.max(torch.ones_like(hs), torch.min(hs, self.voxel_shape[2] - minzs))
        coords_x = coords[:, :, 0] * ds[:, None] / self.query_shape[0] + minxs[:, None]
        coords_y = coords[:, :, 1] * ws[:, None] / self.query_shape[1] + minys[:, None]
        coords_z = coords[:, :, 2] * hs[:, None] / self.query_shape[2] + minzs[:, None]
        coords = torch.stack([coords_x, coords_y, coords_z], dim=-1).floor()
        x_emb = self.x_embeddings(coords_x.long())
        y_emb = self.y_embeddings(coords_y.long())
        z_emb = self.z_embeddings(coords_z.long())
        emb = torch.cat([x_emb, y_emb, z_emb], dim=-1)
        return emb, coords


class Up(nn.Module):
    def __init__(self, in_channels, dim, out_channels):
        super().__init__()
        self.conv3d = nn.Sequential(
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=True),
            nn.Conv3d(in_channels, dim, 3, padding=1, bias=False),
            nn.BatchNorm3d(dim),
            nn.ReLU(inplace=True),
            nn.Conv3d(dim, out_channels, 1, padding=0, bias=False),
            nn.BatchNorm3d(out_channels),
        )

    def forward(self, x):
        x = x.permute(0, 4, 1, 2, 3)
        x = self.conv3d(x)
        return x.permute(0, 2, 3, 4, 1)


def voxel2world(coords, origin, scale, offsets=(0.5, 0.5, 0.5)):
    offsets = torch.tensor(offsets, device=coords.device)
    world_pts = origin + (scale * coords) + scale * offsets
    return world_pts


def world2cam(xyz, cam_Es):
    ones = torch.ones((*xyz.shape[:2], 1), device=xyz.device)
    xyz_h = torch.cat((xyz, ones), dim=2)
    xyz_t_h = torch.bmm(cam_Es, xyz_h.transpose(1, 2)).transpose(1, 2)
    return xyz_t_h[:, :, :3]


def cam2pix(cam_pts, cam_Ks):
    cam_pts = cam_pts.clone()
    fx, fy = cam_Ks[:, 0, 0], cam_Ks[:, 1, 1]
    cx, cy = cam_Ks[:, 0, 2], cam_Ks[:, 1, 2]
    cam_pts[:, :, :2] = cam_pts[:, :, :2] / cam_pts[:, :, 2, None]
    cam_pts[:, :, 0] = cam_pts[:, :, 0] * fx.unsqueeze(1) + cx.unsqueeze(1)
    cam_pts[:, :, 1] = cam_pts[:, :, 1] * fy.unsqueeze(1) + cy.unsqueeze(1)
    pix = torch.round(cam_pts).long()
    return pix[:, :, :2]


class SpatialTransformer3D(OccupancyDetrPretrainedModel):
    def __init__(self, config: OccupancyDetrConfig):
        super().__init__(config)
        self.voxel_shape = config.voxel_shape
        self.voxel_origin = config.voxel_origin
        self.voxel_size = config.voxel_size
        self.image_size = (config.image_width, config.image_height)
        self.query_shape = config.query_shape

        d_model = config.d_model
        self.query_position_embeddings = VoxelLearnedPositionEmbedding(d_model, self.voxel_shape, self.query_shape)
        self.decoder_layers = []
        for _ in range(config.occ_decoder_layers):
            self.decoder_layers.append(SpatialDecoderLayer(config))
        self.decoder_layers = nn.ModuleList(self.decoder_layers)

        self.up = Up(d_model, d_model, d_model)
        self.occupancy_head = nn.Linear(d_model, 1)

        # Initialize weights and apply final processing
        self.post_init()

    @staticmethod
    def get_reference_points(spatial_shapes, device):
        reference_points_list = []
        for level, (depth, width, height) in enumerate(spatial_shapes):
            # in query coordinate system
            ref_x, ref_y, ref_z = meshgrid(
                torch.linspace(0.5, depth - 0.5, depth, dtype=torch.float32, device=device),
                torch.linspace(0.5, width - 0.5, width, dtype=torch.float32, device=device),
                torch.linspace(0.5, height - 0.5, height, dtype=torch.float32, device=device),
                indexing="ij",
            )
            ref_x = ref_x.reshape(-1)[None] / depth
            ref_y = ref_y.reshape(-1)[None] / width
            ref_z = ref_z.reshape(-1)[None] / height
            ref = torch.stack((ref_x, ref_y, ref_z), -1)
            reference_points_list.append(ref)
        reference_points = torch.cat(reference_points_list, 1)
        reference_points = reference_points[:, :, None]
        return reference_points

    def forward(
        self,
        cam_Es,
        cam_Ks,
        context,
        boxes3d,
        features2d,
        spatial_shapes,
        level_start_index,
        valid_ratios,
    ):
        shape = (context.shape[0], *self.query_shape, context.shape[1])

        position_embeddings, coords = self.query_position_embeddings(boxes3d)
        hidden_states = context[:, None, None, None, :].expand(shape).detach()

        device = hidden_states.device
        spatial_shapes_3d = torch.tensor([self.query_shape], device=device)
        level_start_index_3d = torch.tensor([0], device=device)
        reference_points_3d = self.get_reference_points(spatial_shapes_3d, device)

        voxel_origin = torch.tensor(self.voxel_origin, device=device)
        world_pts = voxel2world(coords, voxel_origin, self.voxel_size)
        cam_pts = world2cam(world_pts, cam_Es)
        pixs = cam2pix(cam_pts, cam_Ks)
        reference_points = pixs / torch.tensor(self.image_size, device=device)[None, None, :]
        reference_points = torch.clamp(reference_points, 0.05, 0.95)
        reference_points_input = reference_points[:, :, None] * valid_ratios[:, None]
        for decoder_layer in self.decoder_layers:
            hidden_states = decoder_layer(
                hidden_states.flatten(1, 3),
                query_position_embeddings=position_embeddings,
                encoder_hidden_states=features2d,
                reference_points_3d=reference_points_3d,
                spatial_shapes_3d=spatial_shapes_3d,
                level_start_index_3d=level_start_index_3d,
                reference_points=reference_points_input,
                spatial_shapes=spatial_shapes,
                level_start_index=level_start_index,
            )
            hidden_states = hidden_states.view(shape)

        voxel_feature = self.up(hidden_states)
        occupancy = self.occupancy_head(voxel_feature.flatten(1, 3)).view(voxel_feature.shape[:-1])
        return occupancy.sigmoid()

src/models/SpatialTransformer3D.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.modeling_deformable_detr import *
from transformers.models.detr.modeling_detr import *

from .OccupancyDetrConfig import OccupancyDetrConfig
from .OccupancyDetrPretrainedModel import OccupancyDetrPretrainedModel


class SpatialDecoderLayer(nn.Module):
    def __init__(self, config: OccupancyDetrConfig):
        super().__init__()

        d_model = config.d_model
        self_config = copy.deepcopy(config)
        self_config.num_feature_levels = 1
        self.self_attn = DeformableDetrMultiscaleDeformableAttention(self_config, num_heads=config.encoder_attention_heads, n_points=config.encoder_n_points)
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout

        self.self_attn_layer_norm = nn.LayerNorm(d_model)
        self.encoder_attn = DeformableDetrMultiscaleDeformableAttention(config, num_heads=config.encoder_attention_heads, n_points=config.encoder_n_points)
        self.encoder_attn_layer_norm = nn.LayerNorm(d_model)
        self.fc1 = nn.Linear(d_model, config.decoder_ffn_dim)
        self.fc2 = nn.Linear(config.decoder_ffn_dim, d_model)
        self.final_layer_norm = nn.LayerNorm(d_model)

    def forward(
        self,
        hidden_states: torch.Tensor,
        query_position_embeddings: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        reference_points_bev: torch.Tensor,
        spatial_shapes_bev: torch.Tensor,
        level_start_index_bev: torch.Tensor,
        reference_points: torch.Tensor,
        spatial_shapes: torch.Tensor,
        level_start_index: torch.Tensor,
    ):
        # Self Attention
        residual = hidden_states

        hidden_states, _ = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=query_position_embeddings,
            encoder_hidden_states=hidden_states,
            reference_points=reference_points_bev,
            spatial_shapes=spatial_shapes_bev,
            level_start_index=level_start_index_bev,
        )

        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        # Cross Attention
        residual = hidden_states

        hidden_states, _ = self.encoder_attn(
            hidden_states=hidden_states,
            position_embeddings=query_position_embeddings,
            encoder_hidden_states=encoder_hidden_states,
            reference_points=reference_points,
            spatial_shapes=spatial_shapes,
            level_start_index=level_start_index,
        )

        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.encoder_attn_layer_norm(hidden_states)

        # Fully Connected
        residual = hidden_states
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.final_layer_norm(hidden_states)

        return hidden_states


class VoxelLearnedPositionEmbedding(nn.Module):
    def __init__(self, embedding_dim, voxel_shape, query_shape):
        super().__init__()
        self.voxel_shape = voxel_shape
        self.query_shape = query_shape
        each_dim = embedding_dim // 3
        self.x_embeddings = nn.Embedding(voxel_shape[0], each_dim)
        self.y_embeddings = nn.Embedding(voxel_shape[1], each_dim)
        self.z_embeddings = nn.Embedding(voxel_shape[2], embedding_dim - 2 * each_dim)

    def forward(self, boxes3d):
        coords = torch.ones((*self.query_shape[:2], 1), device=boxes3d.device).nonzero()
        coords = coords.unsqueeze(0).repeat(boxes3d.shape[0], 1, 1)
        xs = boxes3d[:, 0] * self.voxel_shape[0]
        ys = boxes3d[:, 1] * self.voxel_shape[1]
        zs = boxes3d[:, 2] * self.voxel_shape[2]
        ds = boxes3d[:, 3] * self.voxel_shape[0]
        ws = boxes3d[:, 4] * self.voxel_shape[1]
        minxs = torch.max(torch.zeros_like(xs), xs - ds / 2)
        minys = torch.max(torch.zeros_like(ys), ys - ws / 2)
        ds = torch.max(torch.ones_like(ds), torch.min(ds, self.voxel_shape[0] - minxs))
        ws = torch.max(torch.ones_like(ws), torch.min(ws, self.voxel_shape[1] - minys))
        coords_x = coords[:, :, 0] * ds[:, None] / self.query_shape[0] + minxs[:, None]
        coords_y = coords[:, :, 1] * ws[:, None] / self.query_shape[1] + minys[:, None]
        coords_z = zs.unsqueeze(1).repeat(1, self.query_shape[0] * self.query_shape[1])
        coords = torch.stack([coords_x, coords_y, coords_z], dim=-1).floor()
        x_emb = self.x_embeddings(coords_x.long())
        y_emb = self.y_embeddings(coords_y.long())
        z_emb = self.z_embeddings(coords_z.long())
        emb = torch.cat([x_emb, y_emb, z_emb], dim=-1)
        return emb, coords


def voxel2world(coords, origin, scale, offsets=(0.5, 0.5, 0.5)):
    offsets = torch.tensor(offsets, device=coords.device)
    world_pts = origin + (scale * coords) + scale * offsets
    return world_pts


def world2cam(xyz, cam_Es):
    ones = torch.ones((*xyz.shape[:2], 1), device=xyz.device)
    xyz_h = torch.cat((xyz, ones), dim=2)
    xyz_t_h = torch.bmm(cam_Es, xyz_h.transpose(1, 2)).transpose(1, 2)
    return xyz_t_h[:, :, :3]


def cam2pix(cam_pts, cam_Ks):
    cam_pts = cam_pts.clone()
    fx, fy = cam_Ks[:, 0, 0], cam_Ks[:, 1, 1]
    cx, cy = cam_Ks[:, 0, 2], cam_Ks[:, 1, 2]
    cam_pts[:, :, :2] = cam_pts[:, :, :2] / cam_pts[:, :, 2, None]
    cam_pts[:, :, 0] = cam_pts[:, :, 0] * fx.unsqueeze(1) + cx.unsqueeze(1)
    cam_pts[:, :, 1] = cam_pts[:, :, 1] * fy.unsqueeze(1) + cy.unsqueeze(1)
    pix = torch.round(cam_pts).long()
    return pix[:, :, :2]


class SpatialTransformerBEV(OccupancyDetrPretrainedModel):
    def __init__(self, config: OccupancyDetrConfig):
        super().__init__(config)
        self.voxel_shape = config.voxel_shape
        self.voxel_origin = config.voxel_origin
        self.voxel_size = config.voxel_size
        self.image_size = (config.image_width, config.image_height)
        self.query_shape = config.query_shape
        self.H = config.query_shape[2]

        d_model = config.d_model
        self.query_position_embeddings = VoxelLearnedPositionEmbedding(d_model, self.voxel_shape, self.query_shape)
        self.decoder_layers = []
        for _ in range(config.occ_decoder_layers):
            self.decoder_layers.append(SpatialDecoderLayer(config))
        self.decoder_layers = nn.ModuleList(self.decoder_layers)
        self.occupancy_head = nn.Linear(d_model, self.H)

        # Initialize weights and apply final processing
        self.post_init()

    @staticmethod
    def get_reference_points(spatial_shapes, device):
        reference_points_list = []
        for level, (depth, width) in enumerate(spatial_shapes):
            ref_x, ref_y = meshgrid(
                torch.linspace(0.5, depth - 0.5, depth, dtype=torch.float32, device=device),
                torch.linspace(0.5, width - 0.5, width, dtype=torch.float32, device=device),
                indexing="ij",
            )
            ref_x = ref_x.reshape(-1)[None] / depth
            ref_y = ref_y.reshape(-1)[None] / width
            ref = torch.stack((ref_x, ref_y), -1)
            reference_points_list.append(ref)
        reference_points = torch.cat(reference_points_list, 1)
        reference_points = reference_points[:, :, None]
        return reference_points

    def forward(
        self,
        cam_Es,
        cam_Ks,
        context,
        boxes3d,
        features2d,
        spatial_shapes,
        level_start_index,
        valid_ratios,
    ):
        shape = (context.shape[0], *self.query_shape[:2], context.shape[1])

        position_embeddings, coords = self.query_position_embeddings(boxes3d)
        hidden_states = context[:, None, None, :].expand(shape).detach()

        device = hidden_states.device
        spatial_shapes_bev = torch.tensor([self.query_shape[:2]], device=device)
        level_start_index_bev = torch.tensor([0], device=device)
        reference_points_bev = self.get_reference_points(spatial_shapes_bev, device)
        spatial_shapes_bev = spatial_shapes_bev[:, [1, 0]]  # adjust for image deformable attn

        voxel_origin = torch.tensor(self.voxel_origin, device=device)
        world_pts = voxel2world(coords, voxel_origin, self.voxel_size)
        cam_pts = world2cam(world_pts, cam_Es)
        pixs = cam2pix(cam_pts, cam_Ks)
        reference_points = pixs / torch.tensor(self.image_size, device=device)[None, None, :]
        reference_points = torch.clamp(reference_points, 0.05, 0.95)
        reference_points_input = reference_points[:, :, None] * valid_ratios[:, None]
        for decoder_layer in self.decoder_layers:
            hidden_states = decoder_layer(
                hidden_states.flatten(1, 2),
                query_position_embeddings=position_embeddings,
                encoder_hidden_states=features2d,
                reference_points_bev=reference_points_bev,
                spatial_shapes_bev=spatial_shapes_bev,
                level_start_index_bev=level_start_index_bev,
                reference_points=reference_points_input,
                spatial_shapes=spatial_shapes,
                level_start_index=level_start_index,
            )
            hidden_states = hidden_states.view(shape)

        occupancy = self.occupancy_head(hidden_states.flatten(1, 2)).view(*shape[:3], self.H).sigmoid()

        return occupancy

src/models/SpatialTransformerBEV.py
jypjypjypjyp/OccupancyDETR
from transformers.models.deformable_detr.modeling_deformable_detr import DeformableDetrMLPPredictionHead as MLP

src/models/__init__.py
jypjypjypjyp/OccupancyDETR
from dataclasses import dataclass
from typing import Any, Dict, List

import numpy as np
import torch
from PIL import Image
from transformers.data.data_collator import DataCollatorMixin
from transformers.models.deformable_detr.image_processing_deformable_detr import DeformableDetrImageProcessor

from tools.data.semantic_kitti import image_height, image_width, read_calib

from ..models.OccupancyDetrConfig import OccupancyDetrConfig


@dataclass
class DataCollatorForKitti(DataCollatorMixin):
    return_tensors: str = "pt"

    def __init__(self, processor: DeformableDetrImageProcessor):
        self.processor = processor
        self.all_calibs = {}

    def __call__(self, features: List[Dict[str, Any]], return_tensors=None) -> Dict[str, Any]:
        if return_tensors is None:
            return_tensors = self.return_tensors

        # loading
        batch_images = []
        batch_labels = []
        batch_cam_Es = []
        batch_cam_Ks = []
        ids = []
        max_length = 0
        for feature in features:
            image = Image.open(feature["image_path"])
            batch_images.append(image)
            ids.append(feature["image_path"])

            if feature["calib_path"] in self.all_calibs:
                calib = self.all_calibs[feature["calib_path"]]
            else:
                calib = read_calib(feature["calib_path"])
                self.all_calibs[feature["calib_path"]] = calib
            batch_cam_Es.append(calib["Tr"].tolist())
            batch_cam_Ks.append(calib["P2"].tolist())

            labels = {
                "class_labels": [],
                "boxes": [],
                "boxes3d": [],
            }
            for obj in feature["labels"][:300]:
                labels["class_labels"].append(obj["label"])
                minx, miny, maxx, maxy = obj["box2d"]
                x = ((minx + maxx) / 2) / image_width
                y = ((miny + maxy) / 2) / image_height
                w = (maxx - minx + 1) / image_width
                h = (maxy - miny + 1) / image_height
                labels["boxes"].append([x, y, w, h])
                minx3d, miny3d, minz3d, maxx3d, maxy3d, maxz3d = obj["box3d"]
                x3d = ((minx3d + maxx3d) / 2) / 256
                y3d = ((miny3d + maxy3d) / 2) / 256
                z3d = ((minz3d + maxz3d) / 2) / 32
                d3d = (maxx3d - minx3d + 1) / 256
                w3d = (maxy3d - miny3d + 1) / 256
                h3d = (maxz3d - minz3d + 1) / 32
                labels["boxes3d"].append([x3d, y3d, z3d, d3d, w3d, h3d])
            labels["class_labels"] = torch.tensor(labels["class_labels"])
            labels["boxes"] = torch.tensor(labels["boxes"])
            labels["boxes3d"] = torch.tensor(labels["boxes3d"])
            labels["obj3d"] = torch.from_numpy(np.load(feature["obj3ds_path"])).float()
            batch_labels.append(labels)
            max_length = max(max_length, len(labels["class_labels"]))

        # For multi-gpu training, convert to tensors
        batch_class_labels = []
        batch_boxes = []
        batch_boxes3d = []
        batch_obj3d = []
        for labels in batch_labels:
            padding_class_labels = torch.full((max_length,), -1, dtype=torch.long)
            padding_boxes = torch.zeros((max_length, 4), dtype=torch.float)
            padding_boxes3d = torch.zeros((max_length, 6), dtype=torch.float)
            padding_obj3d = torch.zeros((max_length, *labels["obj3d"].shape[1:]), dtype=torch.float)
            padding_class_labels[: len(labels["class_labels"])] = labels["class_labels"]
            padding_boxes[: len(labels["boxes"])] = labels["boxes"]
            padding_boxes3d[: len(labels["boxes3d"])] = labels["boxes3d"]
            padding_obj3d[: len(labels["obj3d"])] = labels["obj3d"]
            batch_class_labels.append(padding_class_labels)
            batch_boxes.append(padding_boxes)
            batch_boxes3d.append(padding_boxes3d)
            batch_obj3d.append(padding_obj3d)

        batch = self.processor(batch_images, return_tensors=return_tensors)
        batch["labels"] = torch.stack(batch_class_labels)
        batch["label_boxes"] = torch.stack(batch_boxes)
        batch["label_boxes3d"] = torch.stack(batch_boxes3d)
        batch["label_obj3d"] = torch.stack(batch_obj3d)
        batch["cam_Es"] = torch.tensor(batch_cam_Es)
        batch["cam_Ks"] = torch.tensor(batch_cam_Ks)
        batch["ids"] = ids
        return batch

src/readers/DataCollatorForKitti.py
jypjypjypjyp/OccupancyDETR
from typing import Dict

import numpy as np
import torch
import torch.nn.functional as F
from transformers.models.deformable_detr.image_processing_deformable_detr import DeformableDetrImageProcessor

from ..models.OccupancyDetr import OccupancyDetrOutput
from ..models.OccupancyDetrConfig import OccupancyDetrConfig


class OccupancyDetrProcessor(DeformableDetrImageProcessor):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.config: OccupancyDetrConfig = None

    def post_process_for_OccupancyDetr(
        self,
        output: OccupancyDetrOutput,
    ) -> Dict:
        n = output.num_preds[0]  # only support batch size 1
        pred_occ = output.pred_occ[:n].detach().cpu() if output.pred_occ is not None else None
        pred_logits = output.pred_logits[:n].detach().sigmoid().cpu().numpy()
        pred_boxes = output.pred_boxes[:n].detach().cpu().numpy()
        pred_boxes3d = output.pred_boxes3d[:n].detach().cpu().numpy()
        classes, scores = pred_logits.argmax(-1), pred_logits.max(-1)

        volumes = []
        for i, box3d in enumerate(pred_boxes3d):
            d, w, h = box3d[3:]
            volumes.append(d * w * h)
        idx = np.argsort(volumes)[::-1]

        voxel = torch.zeros((256, 256, 32), dtype=torch.int8)
        results = {
            "classes": [],
            "scores": [],
            "boxes": [],
            "boxes3d": [],
        }
        for i in idx:
            results["classes"].append(int(classes[i]))
            results["scores"].append(float(scores[i]))
            box = np.zeros((4))
            box[:2] = pred_boxes[i][:2] - pred_boxes[i][2:] / 2
            box[2:] = pred_boxes[i][:2] + pred_boxes[i][2:] / 2
            results["boxes"].append(box.tolist())
            box3d = np.zeros((6))
            box3d[:3] = pred_boxes3d[i][:3] - pred_boxes3d[i][3:6] / 2
            box3d[3:] = pred_boxes3d[i][:3] + pred_boxes3d[i][3:6] / 2
            results["boxes3d"].append(box3d.tolist())
            if pred_occ is not None:
                x, y, z, d, w, h = pred_boxes3d[i]
                x, y, z = int(x * 256), int(y * 256), int(z * 32)
                d, w, h = int(d * 256), int(w * 256), int(h * 32)
                minx = max(0, x - d // 2)
                miny = max(0, y - w // 2)
                minz = max(0, z - h // 2)
                d = max(1, min(d, 256 - minx))
                w = max(1, min(w, 256 - miny))
                h = max(1, min(h, 32 - minz))
                occ_in_world = torch.zeros((256, 256, 32))
                occ_in_box = F.interpolate(pred_occ[i][None, None, :, :, :], size=(d, w, h))[0, 0]
                occ_in_world[minx : minx + d, miny : miny + w, minz : minz + h] = occ_in_box
                voxel[occ_in_world > self.config.occupancy_threshold] = classes[i]
        return results, voxel.numpy()

src/utils/OccupancyDetrProcessor.py
jypjypjypjyp/OccupancyDETR
import json
import os


def write_json(path, data, default=None):
    with open(path, "w") as f:
        json.dump(data, f, default=default)


def read_json(path):
    with open(path) as f:
        return json.load(f)


tools/common.py
jypjypjypjyp/OccupancyDETR
"""
Most of the code in this file is taken from https://github.com/cv-rits/LMSCNet/blob/main/LMSCNet/data/io_data.py
"""

import numpy as np

id2label = {
    0: "unlabeled",
    1: "car",
    2: "bicycle",
    3: "motorcycle",
    4: "truck",
    5: "other-vehicle",
    6: "person",
    7: "bicyclist",
    8: "motorcyclist",
    9: "road",
    10: "parking",
    11: "sidewalk",
    12: "other-ground",
    13: "building",
    14: "fence",
    15: "vegetation",
    16: "trunk",
    17: "terrain",
    18: "pole",
    19: "traffic-sign",
}

label2id = {
    "unlabeled": 0,
    "car": 1,
    "bicycle": 2,
    "motorcycle": 3,
    "truck": 4,
    "other-vehicle": 5,
    "person": 6,
    "bicyclist": 7,
    "motorcyclist": 8,
    "road": 9,
    "parking": 1,
    "sidewalk": 1,
    "other-ground": 1,
    "building": 1,
    "fence": 1,
    "vegetation": 1,
    "trunk": 1,
    "terrain": 1,
    "pole": 1,
    "traffic-sign": 1,
}

learning_map = {
    0: 0,  # "unlabeled"
    1: 0,  # "outlier" mapped to "unlabeled" --------------------------mapped
    10: 1,  # "car"
    11: 2,  # "bicycle"
    13: 5,  # "bus" mapped to "other-vehicle" --------------------------mapped
    15: 3,  # "motorcycle"
    16: 5,  # "on-rails" mapped to "other-vehicle" ---------------------mapped
    18: 4,  # "truck"
    20: 5,  # "other-vehicle"
    30: 6,  # "person"
    31: 7,  # "bicyclist"
    32: 8,  # "motorcyclist"
    40: 9,  # "road"
    44: 10,  # "parking"
    48: 11,  # "sidewalk"
    49: 12,  # "other-ground"
    50: 13,  # "building"
    51: 14,  # "fence"
    52: 0,  # "other-structure" mapped to "unlabeled" ------------------mapped
    60: 9,  # "lane-marking" to "road" ---------------------------------mapped
    70: 15,  # "vegetation"
    71: 16,  # "trunk"
    72: 17,  # "terrain"
    80: 18,  # "pole"
    81: 19,  # "traffic-sign"
    99: 0,  # "other-object" to "unlabeled" ----------------------------mapped
    252: 1,  # "moving-car" to "car" ------------------------------------mapped
    253: 7,  # "moving-bicyclist" to "bicyclist" ------------------------mapped
    254: 6,  # "moving-person" to "person" ------------------------------mapped
    255: 8,  # "moving-motorcyclist" to "motorcyclist" ------------------mapped
    256: 5,  # "moving-on-rails" mapped to "other-vehicle" --------------mapped
    257: 5,  # "moving-bus" mapped to "other-vehicle" -------------------mapped
    258: 4,  # "moving-truck" to "truck" --------------------------------mapped
    259: 5,  # "moving-other"-vehicle to "other-vehicle" ----------------mapped
}

learning_map_inv = {  # inverse of previous map
    0: 0,  # "unlabeled", and others ignored
    1: 10,  # "car"
    2: 11,  # "bicycle"
    3: 15,  # "motorcycle"
    4: 18,  # "truck"
    5: 20,  # "other-vehicle"
    6: 30,  # "person"
    7: 31,  # "bicyclist"
    8: 32,  # "motorcyclist"
    9: 40,  # "road"
    10: 44,  # "parking"
    11: 48,  # "sidewalk"
    12: 49,  # "other-ground"
    13: 50,  # "building"
    14: 51,  # "fence"
    15: 70,  # "vegetation"
    16: 71,  # "trunk"
    17: 72,  # "terrain"
    18: 80,  # "pole"
    19: 81,  # "traffic-sign"
}


vox_origin = np.array([0, -25.6, -2])
voxel_size = 0.2
image_width = 1220
image_height = 370
scene_size = (51.2, 51.2, 6.4)


def unpack(compressed):
    """given a bit encoded voxel grid, make a normal voxel grid out of it."""
    uncompressed = np.zeros(compressed.shape[0] * 8, dtype=np.uint8)
    uncompressed[::8] = compressed[:] >> 7 & 1
    uncompressed[1::8] = compressed[:] >> 6 & 1
    uncompressed[2::8] = compressed[:] >> 5 & 1
    uncompressed[3::8] = compressed[:] >> 4 & 1
    uncompressed[4::8] = compressed[:] >> 3 & 1
    uncompressed[5::8] = compressed[:] >> 2 & 1
    uncompressed[6::8] = compressed[:] >> 1 & 1
    uncompressed[7::8] = compressed[:] & 1

    return uncompressed


def pack(array):
    """convert a boolean array into a bitwise array."""
    array = array.reshape((-1))

    # compressing bit flags.
    # yapf: disable
    compressed = array[::8] << 7 | array[1::8] << 6    | array[2::8] << 5 | array[3::8] << 4 | array[4::8] << 3 | array[5::8] << 2 | array[6::8] << 1 | array[7::8]
    # yapf: enable

    return np.array(compressed, dtype=np.uint8)


def get_remap_lut():
    # make lookup table for mapping
    maxkey = max(learning_map.keys())

    # +100 hack making lut bigger just in case there are unknown labels
    remap_lut = np.zeros((maxkey + 100), dtype=np.int32)
    remap_lut[list(learning_map.keys())] = list(learning_map.values())

    # in completion we have to distinguish empty and invalid voxels.
    # Important: For voxels 0 corresponds to "empty" and not "unlabeled".
    remap_lut[remap_lut == 0] = 255  # map 0 to 'invalid'
    remap_lut[0] = 0  # only 'empty' stays 'empty'.

    return remap_lut


def get_cmap():
    colors = np.array(
        [
            [128, 128, 128, 55],
            [100, 150, 245, 255],
            [100, 230, 245, 255],
            [30, 60, 150, 255],
            [80, 30, 180, 255],
            [100, 80, 250, 255],
            [255, 30, 30, 255],
            [255, 40, 200, 255],
            [150, 30, 90, 255],
            [255, 0, 255, 255],
            [255, 150, 255, 255],
            [75, 0, 75, 255],
            [175, 0, 75, 255],
            [255, 200, 0, 255],
            [255, 120, 50, 255],
            [0, 175, 0, 255],
            [135, 60, 0, 255],
            [150, 240, 80, 255],
            [255, 240, 150, 255],
            [255, 0, 0, 255],
        ]
    ).astype(np.uint8)
    return colors


def read_bin(path, dtype, do_unpack):
    bin = np.fromfile(path, dtype=dtype)  # Flattened array
    if do_unpack:
        bin = unpack(bin)
    return bin


def read_label(path):
    label = read_bin(path, dtype=np.uint16, do_unpack=False).astype(np.uint16)
    return label


def read_invalid(path):
    invalid = read_bin(path, dtype=np.uint8, do_unpack=True)
    return invalid


def read_occluded(path):
    occluded = read_bin(path, dtype=np.uint8, do_unpack=True)
    return occluded


def read_occupancy(path):
    occupancy = read_bin(path, dtype=np.uint8, do_unpack=True).astype(np.float32)
    return occupancy


def read_pointcloud(path):
    "Return pointcloud semantic kitti with remissions (x, y, z, intensity)"
    pointcloud = read_bin(path, dtype=np.float32, do_unpack=False)
    pointcloud = pointcloud.reshape((-1, 4))
    return pointcloud


def read_calib(calib_path):
    """
    :param calib_path: Path to a calibration text file.
    :return: dict with calibration matrices.
    """
    calib_all = {}
    with open(calib_path, "r") as f:
        for line in f.readlines():
            if line == "\n":
                break
            key, value = line.split(":", 1)
            calib_all[key] = np.array([float(x) for x in value.split()])

    # reshape matrices
    calib_out = {}
    calib_out["P2"] = calib_all["P2"].reshape(3, 4)  # 3x4 projection matrix for left camera
    calib_out["Tr"] = np.identity(4)  # 4x4 matrix
    calib_out["Tr"][:3, :4] = calib_all["Tr"].reshape(3, 4)
    return calib_out

tools/data/semantic_kitti.py
jypjypjypjyp/OccupancyDETR
import os

import numpy as np
from tqdm import tqdm

from tools.common import read_json
from tools.data.semantic_kitti import image_height, image_width


def eval_voxel(output_path, restore_path):
    pred = np.load(output_path)
    gt = np.load(restore_path)
    Is = []
    Us = []
    for i in range(20):
        cur_pred = pred == i
        cur_gt = gt == i
        Is.append(np.sum(cur_pred & cur_gt))
        Us.append(np.sum(cur_pred | cur_gt))
    return Is, Us


def eval_voxels(output_dir, restore_dir):
    Iss, Uss = [], []
    for file in tqdm(os.listdir(restore_dir)):
        output_path = os.path.join(output_dir, file)
        restore_path = os.path.join(restore_dir, file)
        Is, Us = eval_voxel(output_path, restore_path)
        Iss.append(Is)
        Uss.append(Us)
    Iss = np.array(Iss).sum(axis=0)
    Uss = np.array(Uss).sum(axis=0) + 1e-6
    ious = Iss / Uss
    return ious


def calculate_iou(box1, box2):
    min_x1, min_y1, max_x1, max_y1 = box1
    min_x2, min_y2, max_x2, max_y2 = box2
    w1, h1 = max_x1 - min_x1, max_y1 - min_y1
    w2, h2 = max_x2 - min_x2, max_y2 - min_y2

    x_inter_left = max(min_x1, min_x2)
    y_inter_top = max(min_y1, min_y2)
    x_inter_right = min(max_x1, max_x2)
    y_inter_bottom = min(max_y1, max_y2)

    I = max(0, x_inter_right - x_inter_left) * max(0, y_inter_bottom - y_inter_top)
    U = w1 * h1 + w2 * h2 - I
    iou = I / U
    return iou


def calculate_ap(gt_boxes, pred_boxes, scroes, iou_threshold=0.5):
    sorted_indices = np.argsort(scroes)[::-1]
    pred_boxes = pred_boxes[sorted_indices]
    matches = np.zeros(len(pred_boxes), dtype=bool)
    gt_matches = np.zeros(len(gt_boxes), dtype=bool)
    for i, predicted_box in enumerate(pred_boxes):
        best_iou = 0.0
        best_match_index = -1
        for j, gt_box in enumerate(gt_boxes[~gt_matches]):
            iou = calculate_iou(predicted_box, gt_box)
            if iou > best_iou:
                best_iou = iou
                best_match_index = j

        if best_iou >= iou_threshold:
            matches[i] = True
            gt_matches[best_match_index] = True

    tp = np.sum(matches)
    fp = len(matches) - tp
    fn = len(gt_boxes) - tp
    return tp, fp, fn


def eval_2d_obj(output_path, gt_path, iou_threshold=0.5):
    result = read_json(output_path)
    classes = np.array(result["classes"])
    boxes = np.array(result["boxes"])
    scores = np.array(result["scores"])
    boxes[:, [0, 2]] *= image_width
    boxes[:, [1, 3]] *= image_height
    gt = read_json(gt_path)
    classes_gt = []
    boxes_gt = []
    for l in gt["labels"]:
        classes_gt.append(l["label"])
        boxes_gt.append(l["box2d"])
    classes_gt = np.array(classes_gt)
    boxes_gt = np.array(boxes_gt)
    sum_tp, sum_fp, sum_fn = 0, 0, 0
    for i in range(1, 20):
        cls_mask = classes == i
        cls_mask_gt = classes_gt == i
        tp, fp, fn = calculate_ap(boxes_gt[cls_mask_gt], boxes[cls_mask], scores[cls_mask], iou_threshold)
        sum_tp += tp
        sum_fp += fp
        sum_fn += fn
    return sum_tp, sum_fp, sum_fn


def eval_2d_objs(output_dir, preprocessed_dir, iou_threshold=0.5):
    sum_tp, sum_fp, sum_fn = 0, 0, 0
    for file in tqdm(os.listdir(preprocessed_dir)):
        output_path = os.path.join(output_dir, file)
        gt_path = os.path.join(preprocessed_dir, file)
        tp, fp, fn = eval_2d_obj(output_path, gt_path, iou_threshold)
        sum_tp += tp
        sum_fp += fp
        sum_fn += fn
    precision = sum_tp / (sum_tp + sum_fp)
    recall = sum_tp / (sum_tp + sum_fn)
    return precision, recall


def eval_2d_objs_during_training(preds, labels, iou_threshold=0.5):
    sum_tp, sum_fp, sum_fn = 0, 0, 0
    pred_logits = preds[2]
    classes, scores = pred_logits.argmax(-1), pred_logits.max(-1)
    boxes = np.zeros_like(preds[3])
    boxes[:, :2] = preds[3][:, :2] - preds[3][:, 2:] / 2
    boxes[:, 2:] = preds[3][:, :2] + preds[3][:, 2:] / 2
    gt_classes = labels[0]
    gt_boxes = np.zeros_like(labels[1])
    gt_boxes[:, :, :2] = labels[1][:, :, :2] - labels[1][:, :, 2:] / 2
    gt_boxes[:, :, 2:] = labels[1][:, :, :2] + labels[1][:, :, 2:] / 2
    bs = labels[0].shape[0]
    num_preds = preds[5]
    cur_i = 0
    for b in range(bs):
        b_classes = classes[cur_i : cur_i + num_preds[b]]
        b_scores = scores[cur_i : cur_i + num_preds[b]]
        b_boxes = boxes[cur_i : cur_i + num_preds[b]]
        cur_i += num_preds[b]
        b_gt_classes = gt_classes[b]
        b_gt_boxes = gt_boxes[b]
        for i in range(1, 20):
            cls_mask = b_classes == i
            cls_mask_gt = b_gt_classes == i
            tp, fp, fn = calculate_ap(b_gt_boxes[cls_mask_gt], b_boxes[cls_mask], b_scores[cls_mask], iou_threshold)
            sum_tp += tp
            sum_fp += fp
            sum_fn += fn
    precision = sum_tp / (sum_tp + sum_fp + 1e-6)
    recall = sum_tp / (sum_tp + sum_fn)
    return {"accuracy": precision, "recall": recall}

tools/eval.py
jypjypjypjyp/OccupancyDETR
import glob
import multiprocessing
import shutil
from multiprocessing.pool import AsyncResult
from typing import List

from common import *
from tqdm import tqdm
from voxel import get_can_see
from voxel2obj import voxel2obj

from tools.data.semantic_kitti import *


def preprocess_frame(label_path, occluded_path, invalid_path, calib_path, output_path, obj3ds_path):
    calib = read_calib(calib_path)
    remap_lut = get_remap_lut()
    label = read_label(label_path)
    occluded = read_occluded(occluded_path)
    invalid = read_invalid(invalid_path)
    label = remap_lut[label.astype(np.uint16)]
    label[occluded == 0] = 255
    label[invalid == 1] = 255
    label = label.reshape([256, 256, 32])
    can_see = get_can_see(label, calib)
    label_dict, obj_voxels = voxel2obj(label, can_see, calib)
    np.save(obj3ds_path, obj_voxels)
    write_json(
        output_path,
        {
            "labels": label_dict,
            "calib_path": calib_path,
            "label_path": label_path,
            "occluded_path": occluded_path,
            "invalid_path": invalid_path,
            "obj3ds_path": obj3ds_path,
        },
    )
    print(len(label_dict))


def preprocess_sequence(seq_dir):
    pool = multiprocessing.Pool(processes=10)
    results: List[AsyncResult] = []

    calib_path = os.path.join(seq_dir, "calib.txt")
    voxel_dir = os.path.join(seq_dir, "voxels")
    image_dir = os.path.join(seq_dir, "image_2")
    preprocessed_dir = os.path.join(seq_dir, "preprocessed")
    obj3d_dir = os.path.join(seq_dir, "obj3d")
    os.makedirs(preprocessed_dir, exist_ok=True)
    os.makedirs(obj3d_dir, exist_ok=True)
    params = []
    glob_path = os.path.join(voxel_dir, "*.label")
    for label_path in glob.glob(glob_path):
        file_name = os.path.basename(label_path)
        frame_id = file_name.split(".")[0]
        occluded_path = os.path.join(voxel_dir, f"{frame_id}.occluded")
        invalid_path = os.path.join(voxel_dir, f"{frame_id}.invalid")
        preprocessed_path = os.path.join(preprocessed_dir, f"{frame_id}.json")
        obj3ds_path = os.path.join(obj3d_dir, f"{frame_id}.npy")
        if os.path.exists(preprocessed_path):
            continue
        params.append((label_path, occluded_path, invalid_path, calib_path, preprocessed_path, obj3ds_path))
    for label_path, occluded_path, invalid_path, calib_path, preprocessed_path, obj3ds_path in params:
        # preprocess_frame(label_path, occluded_path, invalid_path, calib_path, preprocessed_path, obj3ds_path)
        results.append(pool.apply_async(preprocess_frame, args=(label_path, occluded_path, invalid_path, calib_path, preprocessed_path, obj3ds_path)))
    for result in tqdm(results):
        result.wait()
    pool.close()
    pool.join()

    lines = []
    for file in os.listdir(preprocessed_dir):
        frame_id = file.split(".")[0]
        image_path = os.path.join(image_dir, f"{frame_id}.png")
        labels = read_json(os.path.join(preprocessed_dir, file))
        labels["image_path"] = image_path
        lines.append(json.dumps(labels))
    if len(lines) > 0:
        with open(os.path.join(seq_dir, "labels.json"), "w") as f:
            f.write("\n".join(lines))


def proprocess_all(dataset_dir):
    for seq_name in os.listdir(dataset_dir):
        seq_dir = os.path.join(dataset_dir, seq_name)
        if os.path.isdir(seq_dir):
            preprocess_sequence(seq_dir)


def clear(dataset_dir):
    for seq_name in os.listdir(dataset_dir):
        seq_dir = os.path.join(dataset_dir, seq_name, "preprocessed")
        if os.path.isdir(seq_dir):
            shutil.rmtree(seq_dir)


if __name__ == "__main__":
    # clear("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/00")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/01")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/02")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/03")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/04")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/05")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/06")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/07")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/08")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/09")
    preprocess_sequence("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/10")

tools/preprocess.py
jypjypjypjyp/OccupancyDETR
import numpy as np
import torch
import torch.nn.functional as F
from common import *
from tqdm import tqdm

from tools.data.semantic_kitti import *


def restore_label(label_path, voxel_path, output_path):
    labels = read_json(label_path)["labels"]
    volumes = []
    for i in range(len(labels)):
        box3d = np.array(labels[i]["box3d"], dtype=np.int32)
        d, w, h = box3d[:3] - box3d[3:]
        volumes.append(d * w * h)
    sorted_idx = np.argsort(volumes)[::-1]

    voxel = torch.zeros((256, 256, 32), dtype=torch.int8)
    obj3ds = np.load(voxel_path)
    for i in sorted_idx:
        box3d = np.array(labels[i]["box3d"], dtype=np.int32)
        cls = labels[i]["label"]
        occ_in_box = torch.from_numpy(obj3ds[i]).float()
        minx, miny, minz, maxx, maxy, maxz = box3d
        d, w, h = (box3d[3:] - box3d[:3]).tolist()
        d = max(1, min(d, 256 - minx))
        w = max(1, min(w, 256 - miny))
        h = max(1, min(h, 32 - minz))
        occ_in_world = torch.zeros((256, 256, 32))
        occ_in_box = F.interpolate(occ_in_box[None, None, :, :, :], size=(d, w, h))[0, 0]
        occ_in_world[minx : minx + d, miny : miny + w, minz : minz + h] = occ_in_box
        voxel[occ_in_world > 0.5] = cls
    voxel = voxel.numpy()
    np.save(output_path, voxel)


if __name__ == "__main__":
    seq_dir = "/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/08"
    obj3d_dir = f"{seq_dir}/obj3d"
    preprocessed_dir = f"{seq_dir}/preprocessed"
    restore_dir = f"{seq_dir}/restore"
    for file in tqdm(os.listdir(preprocessed_dir)):
        label_path = os.path.join(preprocessed_dir, file)
        obj3d_path = os.path.join(obj3d_dir, file.replace(".json", ".npy"))
        output_path = os.path.join(restore_dir, file.replace(".json", ".npy"))
        restore_label(label_path, obj3d_path, output_path)

tools/restore.py
jypjypjypjyp/OccupancyDETR
import cv2
import numpy as np
import pyvista as pv
import torch
import torch.nn.functional as F
from common import *
from tqdm import tqdm
from voxel import *

from data.semantic_kitti import *


def viz_raw_voxel(label_path, occluded_path=None, invalid_path=None):
    remap_lut = get_remap_lut()
    label = read_label(label_path)
    label = read_label(label_path)
    label = remap_lut[label.astype(np.uint16)]
    if occluded_path is not None:
        occluded = read_occluded(occluded_path)
        label[occluded == 0] = 255
    if invalid_path is not None:
        invalid = read_invalid(invalid_path)
        label[invalid == 0] = 255
    label = label.reshape([256, 256, 32])

    dims = [256, 256, 32]
    grid = pv.UniformGrid(dims)
    grid["Occupancy"] = label.flatten(order="F")
    cmap = get_cmap()
    plotter = pv.Plotter()
    for i in range(20):
        color = cmap[i]
        threshed_occupied = grid.threshold([i, i + 0.1])
        if threshed_occupied.n_points > 0:
            opacity = 0.5 if i == 0 else 1
            plotter.add_mesh(threshed_occupied, color=color, show_edges=False, opacity=opacity)
    plotter.show()


def viz_vox2pix(calib_path, invalid_path, label_path):
    calib = read_calib(calib_path)

    remap_lut = get_remap_lut()
    label = read_label(label_path)
    label = remap_lut[label.astype(np.uint16)]

    invalid = read_invalid(invalid_path)
    label[invalid == 1] = 0
    label = label.reshape([256, 256, 32])

    projected_pix, fov_mask, pix_z, vox_coords = vox2pix(calib)
    vox_coords = vox_coords[fov_mask]
    labe_mask = label[vox_coords[:, 0], vox_coords[:, 1], vox_coords[:, 2]] != 0
    vox_coords = vox_coords[labe_mask]
    pix_z = pix_z[fov_mask]
    pix_z = pix_z[labe_mask]
    projected_pix, _, _ = vox2pix2(calib, vox_coords)
    projected_vox = projected_pix.reshape((8, -1, 2))
    projected_vox = projected_vox.transpose((1, 0, 2))

    near2far = np.argsort(pix_z)
    depth = np.zeros([image_height, image_width], dtype=np.float32)
    can_see = []
    for i in tqdm(near2far):
        z = pix_z[i]
        points = projected_vox[i].astype(np.int32)
        hull = ConvexHull(points)
        vertex = points[hull.vertices]
        new_image = np.zeros_like(depth)
        cv2.fillPoly(new_image, [vertex], color=z)
        new_image = new_image * (depth == 0)
        if new_image.sum() > 0:
            depth = depth + new_image
            can_see.append(i)
    depth[depth == 0] = 100
    depth = (1 - depth / 100) * 255
    cv2.imwrite("viz.png", depth)

    can_see_coords = vox_coords[can_see]
    can_see_label = np.zeros_like(label)
    can_see_label[can_see_coords[:, 0], can_see_coords[:, 1], can_see_coords[:, 2]] = label[can_see_coords[:, 0], can_see_coords[:, 1], can_see_coords[:, 2]]
    np.save("viz.npy", can_see_label)


def viz_bbox2d(viz_path, label_path):
    img = cv2.imread(viz_path)
    labels = read_json(label_path)
    cmap = get_cmap()
    for data in labels["labels"]:
        label = data["label"]
        minx, miny, maxx, maxy = np.array(data["bbox"], dtype=np.int32)
        color = (int(cmap[label][0]), int(cmap[label][1]), int(cmap[label][2]))
        img = cv2.rectangle(img, (minx, miny), (maxx, maxy), color, 2)
    cv2.imwrite("viz2.png", img)


def viz_label(image_path, label_path, voxel_path):
    labels = read_json(label_path)["labels"]
    image = cv2.imread(image_path)
    cmap = get_cmap().tolist()
    h, w = image.shape[:2]
    for obj in labels:
        cls = obj["label"]
        minx, miny, maxx, maxy = obj["box2d"]
        minx = int(minx)
        miny = int(miny)
        maxx = int(maxx)
        maxy = int(maxy)
        color = cmap[cls]
        image = cv2.rectangle(image, (minx, miny), (maxx, maxy), color, 2)
        image = cv2.putText(image, f"{id2label[cls]}", (minx, miny), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
    cv2.imwrite("viz.png", image)

    voxel = torch.zeros((256, 256, 32))
    obj3ds = np.load(voxel_path)
    for i in range(len(obj3ds)):
        box3d = np.array(labels[i]["box3d"], dtype=np.int32)
        cls = labels[i]["label"]
        occ_in_box = torch.from_numpy(obj3ds[i]).float()
        minx, miny, minz, maxx, maxy, maxz = box3d
        d, w, h = box3d[3:] - box3d[:3]
        d = max(1, min(d, 256 - minx))
        w = max(1, min(w, 256 - miny))
        h = max(1, min(h, 32 - minz))
        occ_in_world = torch.zeros((256, 256, 32))
        occ_in_box = F.interpolate(occ_in_box[None, None, :, :, :], size=(d, w, h))[0, 0]
        occ_in_world[minx : minx + d, miny : miny + w, minz : minz + h] = occ_in_box
        voxel[occ_in_world > 0.5] = cls
    voxel = voxel.numpy()

    dims = [256, 256, 32]
    grid = pv.UniformGrid(dims)
    grid["Occupancy"] = voxel.flatten(order="F")
    cmap = get_cmap()
    plotter = pv.Plotter()
    for i in range(1, 20):
        color = cmap[i]
        threshed_occupied = grid.threshold([i, i + 0.1])
        if threshed_occupied.n_points > 0:
            plotter.add_mesh(threshed_occupied, color=color, show_edges=False)

    for obj in labels:
        cls = obj["label"]
        minx, miny, minz, maxx, maxy, maxz = obj["box3d"]
        color = cmap[cls]
        box = pv.Box([minx, maxx, miny, maxy, minz, maxz])
        edges = box.extract_feature_edges()
        plotter.add_mesh(edges, color=color)
        plotter.add_point_labels([[minx, miny, minz]], [id2label[cls]], font_size=8)
    plotter.show()


def viz_result(image_path, result_path, voxel_path):
    result = read_json(result_path)
    image = cv2.imread(image_path)
    cmap = get_cmap().tolist()
    h, w = image.shape[:2]
    for cls, s, bbox in zip(result["classes"], result["scores"], result["boxes"]):
        minx, miny, maxx, maxy = bbox
        minx = int(minx * w)
        miny = int(miny * h)
        maxx = int(maxx * w)
        maxy = int(maxy * h)
        color = cmap[cls]
        image = cv2.rectangle(image, (minx, miny), (maxx, maxy), color, 2)
        image = cv2.putText(image, f"{id2label[cls]} {s:.2f}", (minx, miny), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
    cv2.imwrite("viz.png", image)

    label = np.load(voxel_path)
    dims = [256, 256, 32]
    grid = pv.UniformGrid(dims)
    grid["Occupancy"] = label.flatten(order="F")
    cmap = get_cmap()
    plotter = pv.Plotter()
    for i in range(1, 20):
        color = cmap[i]
        threshed_occupied = grid.threshold([i, i + 0.1])
        if threshed_occupied.n_points > 0:
            plotter.add_mesh(threshed_occupied, color=color, show_edges=False)

    for cls, s, bbox3d in zip(result["classes"], result["scores"], result["boxes3d"]):
        minx, miny, minz, maxx, maxy, maxz = bbox3d
        minx = int(minx * 256)
        miny = int(miny * 256)
        minz = int(minz * 32)
        maxx = int(maxx * 256)
        maxy = int(maxy * 256)
        maxz = int(maxz * 32)
        color = cmap[cls]
        box = pv.Box([minx, maxx, miny, maxy, minz, maxz])
        edges = box.extract_feature_edges()
        plotter.add_mesh(edges, color=color)
        plotter.add_point_labels([[minx, miny, minz]], [id2label[cls]], font_size=8)
    plotter.show()


def viz_intersection(voxel_path, restore_path):
    label = np.load(restore_path).astype(np.float32)
    gt = np.load(voxel_path).astype(np.float32)
    inter = np.zeros_like(label)
    inter[gt != label] = 0.5
    label = label + inter
    dims = [256, 256, 32]
    grid = pv.UniformGrid(dims)
    grid["Occupancy"] = label.flatten(order="F")
    cmap = get_cmap()
    plotter = pv.Plotter()
    for i in range(1, 20):
        color = cmap[i]
        threshed_occupied = grid.threshold([i, i + 0.1])
        if threshed_occupied.n_points > 0:
            plotter.add_mesh(threshed_occupied, color=color, show_edges=False)
        threshed_occupied2 = grid.threshold([i+0.5, i + 0.6])
        if threshed_occupied2.n_points > 0:
            plotter.add_mesh(threshed_occupied2, color=color, show_edges=False, opacity=0.2)
    plotter.show()


def record_results(output_dir):
    dims = [256, 256, 32]
    cmap = get_cmap()
    plotter = pv.Plotter()
    plotter.camera_position = [
        (-170.04, 123.01, 276.53),
        (112.41, 126.32, -2.1660),
        (0.7023, -0.002, 0.71180),
    ]
    plotter.open_movie("viz.mp4", framerate=4)
    actors = []
    for f in tqdm(sorted(os.listdir(output_dir))):
        if not f.endswith(".npy"):
            continue
        for a in actors:
            plotter.remove_actor(a)
        actors.clear()
        voxel_path = os.path.join(output_dir, f)
        label = np.load(voxel_path)
        grid = pv.UniformGrid(dims)
        grid["Occupancy"] = label.flatten(order="F")
        for i in range(1, 20):
            color = cmap[i]
            threshed_occupied = grid.threshold([i, i + 0.1])
            if threshed_occupied.n_points > 0:
                actors.append(plotter.add_mesh(threshed_occupied, color=color, show_edges=False))
        plotter.write_frame()
    plotter.close()


if __name__ == "__main__":
    # restore_dir = "/home/jyp/Cloud/data2/Datasets/SemanticKITTI/dataset/sequences/08/restore"
    # output_dir = "/home/jyp/Cloud/data2/Datasets/SemanticKITTI/dataset/sequences/08/outputs"
    # for file in os.listdir(restore_dir):
    #     voxel_path = os.path.join(output_dir, file)
    #     restore_path = os.path.join(restore_dir, file)
    #     viz_intersection(voxel_path, restore_path)

    # viz_result(
    #     "/home/jyp/Cloud/data2/Datasets/SemanticKITTI/dataset/sequences/08/image_2/000755.png",
    #     "/home/jyp/Cloud/data2/Datasets/SemanticKITTI/dataset/sequences/08/outputs/000755.json",
    #     "/home/jyp/Cloud/data2/Datasets/SemanticKITTI/dataset/sequences/08/outputs/000755.npy",
    # )

    record_results("/mnt/data2/Datasets/SemanticKITTI/dataset/sequences/08/outputs")
tools/viz.py
jypjypjypjyp/OccupancyDETR
import cv2
import numpy as np
from numba import njit, prange
from scipy.interpolate import NearestNDInterpolator
from scipy.spatial import ConvexHull

voxel_origin = np.array([0, -25.6, -2])
voxel_size = 0.2
image_width = 1220
image_height = 370
scene_size = (51.2, 51.2, 6.4)


def set_meta_info(_voxel_origin, _voxel_size, _image_width, _image_height, _scene_size):
    global voxel_origin, voxel_size, image_width, image_height, scene_size
    voxel_origin = _voxel_origin
    voxel_size = _voxel_size
    image_width = _image_width
    image_height = _image_height
    scene_size = _scene_size


@njit(parallel=True)
def vox2world(vol_origin, vox_coords, vox_size, offsets=(0.5, 0.5, 0.5)):
    """Convert voxel grid coordinates to world coordinates."""
    vol_origin = vol_origin.astype(np.float32)
    vox_coords = vox_coords.astype(np.float32)
    cam_pts = np.empty_like(vox_coords, dtype=np.float32)

    for i in prange(vox_coords.shape[0]):
        for j in range(3):
            cam_pts[i, j] = vol_origin[j] + (vox_size * vox_coords[i, j]) + vox_size * offsets[j]
    return cam_pts


@njit(parallel=True)
def cam2pix(cam_pts, intr):
    intr = intr.astype(np.float32)
    fx, fy = intr[0, 0], intr[1, 1]
    cx, cy = intr[0, 2], intr[1, 2]
    pix = np.empty((cam_pts.shape[0], 2), dtype=np.int64)
    for i in prange(cam_pts.shape[0]):
        pix[i, 0] = int(np.round((cam_pts[i, 0] * fx / cam_pts[i, 2]) + cx))
        pix[i, 1] = int(np.round((cam_pts[i, 1] * fy / cam_pts[i, 2]) + cy))
    return pix


@njit(parallel=True)
def pix2cam(pix, depth, intr):
    intr = intr.astype(np.float32)
    fx, fy = intr[0, 0], intr[1, 1]
    cx, cy = intr[0, 2], intr[1, 2]
    cam_pts = np.empty((pix.shape[0], 3), dtype=np.float32)
    for i in prange(pix.shape[0]):
        cam_pts[i, 0] = (pix[i, 0] - cx) * depth[i] / fx
        cam_pts[i, 1] = (pix[i, 1] - cy) * depth[i] / fy
        cam_pts[i, 2] = depth[i]
    return cam_pts


def rigid_transform(xyz, transform):
    xyz_h = np.hstack([xyz, np.ones((len(xyz), 1), dtype=np.float32)])
    xyz_t_h = np.dot(transform, xyz_h.T).T
    return xyz_t_h[:, :3]


def vox2pix(calib, vox_coords=None):
    cam_E = calib["Tr"]
    cam_k = calib["P2"]

    if vox_coords is None:
        # Compute the x, y, z bounding of the scene in meter
        vol_bnds = np.zeros((3, 2))
        vol_bnds[:, 0] = voxel_origin
        vol_bnds[:, 1] = voxel_origin + np.array(scene_size)
        # Compute the voxels centroids in lidar cooridnates
        vol_dim = np.ceil((vol_bnds[:, 1] - vol_bnds[:, 0]) / voxel_size).copy(order="C").astype(int)
        xv, yv, zv = np.meshgrid(range(vol_dim[0]), range(vol_dim[1]), range(vol_dim[2]), indexing="ij")
        vox_coords = np.concatenate([xv.reshape(1, -1), yv.reshape(1, -1), zv.reshape(1, -1)], axis=0).astype(int).T

    # Project voxels'centroid from lidar coordinates to camera coordinates
    cam_pts = vox2world(voxel_origin, vox_coords, voxel_size)
    cam_pts = rigid_transform(cam_pts, cam_E)

    # Project camera coordinates to pixel positions
    projected_pix = cam2pix(cam_pts, cam_k)
    pix_x, pix_y = projected_pix[:, 0], projected_pix[:, 1]

    # Eliminate pixels outside view frustum
    pix_z = cam_pts[:, 2]
    fov_mask = np.logical_and(pix_x >= 0, np.logical_and(pix_x < image_width, np.logical_and(pix_y >= 0, np.logical_and(pix_y < image_height, pix_z > 0))))

    return projected_pix, fov_mask, pix_z, vox_coords


def vox2pix2(calib, vox_coords):
    cam_E = calib["Tr"]
    cam_k = calib["P2"]

    # Project voxels'centroid from lidar coordinates to camera coordinates
    cam_pts = []
    cam_pts.extend(vox2world(voxel_origin, vox_coords, voxel_size, (0, 0, 0)))
    cam_pts.extend(vox2world(voxel_origin, vox_coords, voxel_size, (0, 0, 1)))
    cam_pts.extend(vox2world(voxel_origin, vox_coords, voxel_size, (0, 1, 0)))
    cam_pts.extend(vox2world(voxel_origin, vox_coords, voxel_size, (0, 1, 1)))
    cam_pts.extend(vox2world(voxel_origin, vox_coords, voxel_size, (1, 0, 0)))
    cam_pts.extend(vox2world(voxel_origin, vox_coords, voxel_size, (1, 0, 1)))
    cam_pts.extend(vox2world(voxel_origin, vox_coords, voxel_size, (1, 1, 0)))
    cam_pts.extend(vox2world(voxel_origin, vox_coords, voxel_size, (1, 1, 1)))
    cam_pts = rigid_transform(cam_pts, cam_E)

    # Project camera coordinates to pixel positions
    projected_pix = cam2pix(cam_pts, cam_k)
    pix_x, pix_y = projected_pix[:, 0], projected_pix[:, 1]

    # Eliminate pixels outside view frustum
    pix_z = cam_pts[:, 2]
    fov_mask = np.logical_and(pix_x >= 0, np.logical_and(pix_x < image_width, np.logical_and(pix_y >= 0, np.logical_and(pix_y < image_height, pix_z > 0))))

    return projected_pix, fov_mask, pix_z


def pix2vol(calib, p2d, z):
    cam_E = calib["Tr"]
    cam_k = calib["P2"]
    p3d_cam = pix2cam(p2d, z, cam_k)
    p3d_world = rigid_transform(p3d_cam, cam_E.T)
    p3d_vox = (p3d_world - voxel_origin) / voxel_size
    return p3d_vox


def griddata(points, values, xi, rescale=False):
    ip = NearestNDInterpolator(points, values, rescale=rescale)
    return ip(xi)


def get_can_see(label, calib):
    projected_pix, fov_mask, pix_z, vox_coords = vox2pix(calib)
    vox_coords = vox_coords[fov_mask]
    fov_mask_label = label[vox_coords[:, 0], vox_coords[:, 1], vox_coords[:, 2]]
    label_mask = (fov_mask_label != 0) & (fov_mask_label != 255)
    vox_coords = vox_coords[label_mask]
    pix_z = pix_z[fov_mask]
    pix_z = pix_z[label_mask]
    projected_pix, _, _ = vox2pix2(calib, vox_coords)
    projected_vox = projected_pix.reshape((8, -1, 2))
    projected_vox = projected_vox.transpose((1, 0, 2))

    near2far = np.argsort(pix_z)
    depth = np.zeros([image_height, image_width], dtype=np.float32)
    can_see = []
    for i in near2far:
        z = pix_z[i]
        points = projected_vox[i].astype(np.int32)
        hull = ConvexHull(points)
        vertex = points[hull.vertices]
        new_image = np.zeros_like(depth)
        cv2.fillPoly(new_image, [vertex], color=z)
        new_image = new_image * (depth == 0)
        if new_image.sum() > 0:
            depth = depth + new_image
            can_see.append(i)

    can_see_coords = vox_coords[can_see]
    can_see_label = np.zeros_like(label)
    can_see_label[can_see_coords[:, 0], can_see_coords[:, 1], can_see_coords[:, 2]] = 1
    return can_see_label.astype(bool)

tools/voxel.py
jypjypjypjyp/OccupancyDETR
from copy import deepcopy

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from common import *
from scipy.spatial import cKDTree
from voxel import vox2pix


def filter_by_density(voxel_np, radius, threshold):
    voxel = torch.from_numpy(voxel_np).reshape([1, 1, 256, 256, 32])
    nn_avgpool = nn.AvgPool3d(kernel_size=radius * 2 + 1, stride=1, padding=radius)
    avgpool = nn_avgpool(voxel.float())
    sumpool = avgpool * (radius * 2 + 1) ** 3
    sumpool[sumpool < threshold] = 0
    sumpool[(sumpool > threshold) | (sumpool == threshold)] = 1
    return (sumpool * voxel).bool().numpy().reshape(voxel_np.shape)


cls2threshold = {
    0: 5,  # "unlabeled", and others ignored
    1: 50,  # "car"
    2: 5,  # "bicycle"
    3: 5,  # "motorcycle"
    4: 50,  # "truck"
    5: 50,  # "other-vehicle"
    6: 5,  # "person"
    7: 5,  # "bicyclist"
    8: 5,  # "motorcyclist"
    9: 100,  # "road"
    10: 50,  # "parking"
    11: 50,  # "sidewalk"
    12: 50,  # "other-ground"
    13: 50,  # "building"
    14: 5,  # "fence"
    15: 50,  # "vegetation"
    16: 20,  # "trunk"
    17: 100,  # "terrain"
    18: 5,  # "pole"
    19: 5,  # "traffic-sign"
}


def cluster_by_distance(label, points: np.ndarray, threshold=20):
    n = points.shape[0]
    kdtree = cKDTree(points)
    near_is = kdtree.query_ball_tree(kdtree, threshold)
    remainder = set(list(range(n)))
    clusters = []
    for i in range(n):
        if i not in remainder:
            continue
        cur_set = set()
        cur_queue = [i]
        while len(cur_queue) > 0:
            q = cur_queue.pop(0)
            cur_set.add(q)
            for j in near_is[q]:
                if j not in cur_set and j in remainder:
                    cur_queue.append(j)
                    cur_set.add(j)
        remainder -= cur_set
        if len(cur_set) > cls2threshold[label]:
            clusters.append(points[list(cur_set)])
    return clusters


def voxel2obj(voxel: np.ndarray, can_see: np.ndarray, calib):
    # filter by density
    unique_labels = set(np.unique(voxel))
    unique_labels -= set([0, 255])
    mask = np.zeros_like(voxel, dtype=bool)
    for label in unique_labels:
        label_voxel = deepcopy(voxel)
        label_voxel[~(label_voxel == label)] = 0
        label_voxel[(label_voxel == label)] = 1
        label_mask = filter_by_density(label_voxel, 2, 10)
        mask = mask | label_mask
    voxel[mask == 0] = 0

    obj_labels = []
    obj_voxels = []
    point_idxes = np.argwhere(~(voxel == 255) & ~(voxel == 0))
    labels = voxel[point_idxes[:, 0], point_idxes[:, 1], point_idxes[:, 2]].astype(np.int32)
    for label in unique_labels:
        label_clusters = point_idxes[np.isclose(labels, label)].astype(np.float32)
        clusters = cluster_by_distance(label, label_clusters, 5)
        for c_ps in clusters:
            c_ps = c_ps.astype(np.int32)
            c_ps[:, 0] = np.clip(c_ps[:, 0], 0, 255)
            c_ps[:, 1] = np.clip(c_ps[:, 1], 0, 255)
            c_ps[:, 2] = np.clip(c_ps[:, 2], 0, 255)
            c_can_see = can_see[c_ps[:, 0], c_ps[:, 1], c_ps[:, 2]]
            c_pix, _, _, _ = vox2pix(calib, c_ps[c_can_see])
            if c_can_see.sum() < 5:
                continue

            minx2d = np.min(c_pix[:, 0])
            maxx2d = np.max(c_pix[:, 0])
            miny2d = np.min(c_pix[:, 1])
            maxy2d = np.max(c_pix[:, 1])
            minx3d = np.min(c_ps[:, 0])
            maxx3d = np.max(c_ps[:, 0])
            miny3d = np.min(c_ps[:, 1])
            maxy3d = np.max(c_ps[:, 1])
            minz3d = np.min(c_ps[:, 2])
            maxz3d = np.max(c_ps[:, 2])
            obj_labels.append(
                {
                    "label": int(label),
                    "box2d": [float(minx2d), float(miny2d), float(maxx2d), float(maxy2d)],
                    "box3d": [
                        float(minx3d),
                        float(miny3d),
                        float(minz3d),
                        float(maxx3d),
                        float(maxy3d),
                        float(maxz3d),
                    ],
                    "id": len(obj_labels),
                }
            )

            voxel256 = np.zeros((256, 256, 32), dtype=np.int8)
            voxel256[c_ps[:, 0], c_ps[:, 1], c_ps[:, 2]] = 1
            box3d = torch.from_numpy(voxel256[minx3d : maxx3d + 1, miny3d : maxy3d + 1, minz3d : maxz3d + 1]).float()
            box3d_norm = F.interpolate(box3d[None, None, :, :, :], size=(32, 32, 32), mode="area")[0, 0]
            obj_voxels.append(box3d_norm.bool().numpy())

    return obj_labels, np.stack(obj_voxels)

tools/voxel2obj.py
jypjypjypjyp/OccupancyDETR