Code for only running the distilled model

[kuerbiskopf]

Hi, is there code to only run the distillation model on nuscenes for example and apply a PCA, like the images from the paper?
Thank you.

[YilmazKadir]

I found this code in our local codebase. I hope that it is helpful.

"""
Model Visualization Script

Author: Karim Abou Zeid
"""

import argparse
import logging
import os
from collections import OrderedDict

import rerun as rr
import torch
import torch.utils.data
from einops import rearrange
from sklearn.decomposition import PCA

from pointcept.datasets.builder import build_dataset
from pointcept.models import build_model
from pointcept.utils.config import Config, DictAction

logger = logging.getLogger(__name__)

MAX_SCENES = 500

EVERY_NTH = 1
PCA_USE_IGNORED = False
IGNORE_CLASSES = []
IGNORE_INVISIBLE = False


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config-file", default="", metavar="FILE", help="path to config file"
    )
    parser.add_argument(
        "--options", nargs="+", action=DictAction, help="custom options"
    )
    args = parser.parse_args()

    cfg = Config.fromfile(args.config_file)
    if args.options is not None:
        cfg.merge_from_dict(args.options)

    inferencer = Inferencer(cfg)

    rr.init(f"{cfg.data.val.type}", spawn=True, strict=True)
    rr.log("config", rr.TextLog(f"{args.config_file}\n{args.options}"), static=True)
    if "names" in cfg:
        rr.log(
            "",
            rr.AnnotationContext(list(enumerate(cfg.data.names))),
            static=True,
        )

    global EVERY_NTH, PCA_USE_IGNORED, IGNORE_CLASSES, IGNORE_INVISIBLE
    if "ScanNet" in cfg.data.val.type:
        EVERY_NTH = 1
        PCA_USE_IGNORED = False
        IGNORE_CLASSES = ["wall", "floor", "ceiling"]
        IGNORE_INVISIBLE = False
    elif "NuScenes" in cfg.data.val.type:
        EVERY_NTH = 30
        PCA_USE_IGNORED = False
        IGNORE_CLASSES = [
            # "driveable_surface",
            # "other_flat",
            # "sidewalk",
            # "terrain",
            # "manmade",
            # "vegetation",
        ]
        IGNORE_INVISIBLE = True
    elif "KITTI" in cfg.data.val.type:
        EVERY_NTH = 30
        PCA_USE_IGNORED = False
        IGNORE_CLASSES = [
            # "terrain",
            # "vegetation",
            # "building",
        ]
        IGNORE_INVISIBLE = True

    counter = 0
    for idx in range(len(inferencer.data_set))[::EVERY_NTH]:
        counter += 1
        if counter > MAX_SCENES:
            break
        input_dict = inferencer.data_set[idx]
        if "condition" in input_dict.keys():
            input_dict["condition"] = [input_dict["condition"]]
        input_dict["pca_ignore_mask"] = pca_ignore_mask(input_dict, cfg, IGNORE_CLASSES)
        input_dict["ignore_mask"] = input_dict["segment"] == cfg.data.ignore_index
        input_dict = dict2cuda(input_dict)

        with torch.no_grad():
            feat = inferencer.model(input_dict)["feat"].cpu()

        input_dict = dict2cpu(input_dict)
        log_sample(input_dict, feat, idx)


def dict2cpu(d):
    for key in d.keys():
        if isinstance(d[key], torch.Tensor):
            d[key] = d[key].cpu()
    return d


def dict2cuda(d):
    for key in d.keys():
        if isinstance(d[key], torch.Tensor):
            d[key] = d[key].cuda(non_blocking=True)
    return d


class Inferencer:
    def __init__(self, cfg):
        self.logger = logger
        self.cfg = cfg
        self.model = self.build_model()
        self.data_set = build_dataset(self.cfg.data.val)

    def build_model(self):
        model = build_model(self.cfg.model)
        model = model.cuda().eval()
        assert os.path.isfile(self.cfg.weight)
        checkpoint = torch.load(self.cfg.weight)
        weight = OrderedDict()
        for key, value in checkpoint["state_dict"].items():
            if key.startswith("module."):
                key = key[7:]  # module.xxx.xxx -> xxx.xxx
            weight[key] = value
        model.load_state_dict(weight, strict=True)
        self.logger.info(
            "=> Loaded weight '{}' (epoch {})".format(
                self.cfg.weight, checkpoint["epoch"]
            )
        )
        return model


def pca_ignore_mask(sample, cfg, ignore_names):
    ignore_mask = sample["segment"] == cfg.data.ignore_index
    for name in ignore_names:
        ignore_mask |= sample["segment"] == cfg.data.names.index(name)
    if IGNORE_INVISIBLE:
        ignore_mask |= ~sample["image_mask"].any(dim=1)
    if IGNORE_INVISIBLE:
        ignore_mask |= ~sample["image_mask"].any(dim=1)
    return ignore_mask


def log_sample(sample, feat, idx: int):
    rr.set_time_sequence("idx", idx)

    # visible_mask = sample["image_mask"].any(dim=1)
    ignore_mask = sample["ignore_mask"]
    # ignore_mask |= torch.linalg.norm(sample["coord"], dim=1) < 7  # KITTI

    pca = PCA(n_components=3)
    pca.fit(feat[~ignore_mask] if PCA_USE_IGNORED else feat[~sample["pca_ignore_mask"]])
    feat_pca = pca.transform(feat)

    # Normalize the PCA components to [0, 1] range for RGB mapping
    feat_min = feat_pca.min(axis=0, keepdims=True)
    feat_max = feat_pca.max(axis=0, keepdims=True)
    feat_pca_normalized = (feat_pca - feat_min) / (feat_max - feat_min)

    # colors = feat.new_zeros(feat.shape[0], 3)
    # colors[visible_mask] = torch.from_numpy(feat_pca_normalized).to(colors.dtype)
    pca_colors = torch.from_numpy(feat_pca_normalized).to(feat.dtype)

    colors = ((sample["feat"][:, :3] + 1) * 127.5).int()  # undo ColorNormalize

    rr.log(
        "color",
        rr.Points3D(sample["coord"][~ignore_mask], colors=colors[~ignore_mask]),
    )
    rr.log(
        "seg",
        rr.Points3D(
            sample["coord"][~ignore_mask], class_ids=sample["segment"][~ignore_mask]
        ),
    )
    rr.log(
        "feat/all",
        rr.Points3D(sample["coord"][~ignore_mask], colors=pca_colors[~ignore_mask]),
    )
    rr.log(
        "feat/unignored",
        rr.Points3D(
            sample["coord"][~sample["pca_ignore_mask"] & ~ignore_mask],
            colors=pca_colors[~sample["pca_ignore_mask"] & ~ignore_mask],
        ),
    )

    if "image" in sample:
        for i, (image, img_coord, img_mask) in enumerate(
            zip(
                sample["image"].squeeze(0),
                sample["image_coord"].transpose(0, 1),
                sample["image_mask"].transpose(0, 1),
                strict=True,
            )
        ):
            rr.log(f"image/cam_{i}", rr.Image(rearrange(image, "c h w -> h w c")))
            # rr.log(
            #     f"image/cam_{i}/coord/seg",
            #     rr.Points2D(img_coord[img_mask], class_ids=sample["segment"][img_mask]),
            # )
            rr.log(
                f"image/cam_{i}/coord/feat",
                rr.Points2D(
                    img_coord[img_mask & ~ignore_mask],
                    colors=pca_colors[img_mask & ~ignore_mask],
                ),
            )

    # rr.log(
    #     "coord/seg/visible",
    #     rr.Points3D(
    #         sample["coord"][all_masks], class_ids=sample["segment"][all_masks]
    #     ),
    # )
    # rr.log(
    #     "coord/feat/visible",
    #     rr.Points3D(sample["coord"][visible_mask], colors=colors[visible_mask]),
    # )


if __name__ == "__main__":
    main()