Codebase for the ICCV 2023 Paper "LoCUS: Learning Multiscale 3D-consistent Features from Posed Images"
from yacs.config import CfgNode as CN

cfg = CN()
############### DATA ###############
cfg.DATASET = CN()
BASE_PATH = "./matterport3d"

cfg.DATASET.TRAIN_DATA_SOURCE = "Matterport"
cfg.DATASET.VAL_DATA_SOURCE = "Matterport"
cfg.DATASET.TEST_DATA_SOURCE = "Matterport"

cfg.DATASET.TRAIN_DATA_ROOT = BASE_PATH
cfg.DATASET.TRAIN_POSE_ROOT = None
cfg.DATASET.TRAIN_INTRINSICS_PATH = None
cfg.DATASET.TRAIN_SCENE_LIST = f"{BASE_PATH}/train_scenes.txt"

cfg.DATASET.VAL_DATA_ROOT = BASE_PATH
cfg.DATASET.VAL_POSE_ROOT = None
cfg.DATASET.VAL_INTRINSICS_PATH = None
cfg.DATASET.VAL_SCENE_LIST = f"{BASE_PATH}/val_scenes.txt"

cfg.DATASET.TEST_DATA_ROOT = BASE_PATH
cfg.DATASET.TEST_POSE_ROOT = None
cfg.DATASET.TEST_INTRINSICS_PATH = None
cfg.DATASET.TEST_SCENE_LIST = f"{BASE_PATH}/test_scenes.txt"

cfg.DATASET.IMG_HEIGHT = 256
cfg.DATASET.IMG_WIDTH = 320
cfg.DATASET.AUGMENTATION_TYPE = None
cfg.DATASET.TRAIN_SHUFFLE = True
cfg.DATASET.USE_OVERLAPS = True

cfg.DATASET.MATTERPORT_HORIZONTAL_IMGS_ONLY = True

configs/data/matterport3d_config.py
dakloepfer/locus
from yacs.config import CfgNode as CN

cfg = CN()

############### MODEL ###############

cfg.MODEL = CN()
cfg.MODEL.TYPE = "dino"

cfg.MODEL.N_PRETRAINED_BLOCKS = 11  # number of pretrained DINO blocks
cfg.MODEL.VIT_BACKBONE = "vitb8"

cfg.MODEL.FEATURE_DIM = 64
cfg.MODEL.POS_LANDMARK_RADIUS = 0.2
cfg.MODEL.NEG_LANDMARK_RADIUS = 10.0
cfg.MODEL.LANDMARK_EMBEDDING_METHOD = "sampled_patch"


############### TRAINING ###############
cfg.TRAINER = CN()
cfg.TRAINER.SEED = 42
cfg.TRAINER.GRADIENT_CLIPPING = None
cfg.TRAINER.ACCUMULATE_GRAD_BATCHES = 1

cfg.TRAIN = CN()
cfg.TRAIN.N_LANDMARKS = 64
cfg.TRAIN.FRAC_PATCH_SUBSAMPLE = 0.1
cfg.TRAIN.KEEP_ALL_POSITIVE_PATCHES = True

############### OPTIMIZER ###############
cfg.OPTIMIZER = CN()
cfg.OPTIMIZER.TYPE = "Adam"
cfg.OPTIMIZER.LR = 1e-3
cfg.OPTIMIZER.WEIGHT_DECAY = 0.0


############## LOSS ###############
cfg.LOSS = CN()
cfg.LOSS.SIGMOID_TEMPERATURE = 0.01

configs/model/standard_feature_extractor_config.py
dakloepfer/locus
import os
import configargparse
from tqdm import tqdm

import numpy as np
import torch

from einops import rearrange

ALL_SCENES = [
    "17DRP5sb8fy",
    "1LXtFkjw3qL",
    "1pXnuDYAj8r",
    "29hnd4uzFmX",
    "2azQ1b91cZZ",
    "2n8kARJN3HM",
    "2t7WUuJeko7",
    "5LpN3gDmAk7",
    "5q7pvUzZiYa",
    "5ZKStnWn8Zo",
    "759xd9YjKW5",
    "7y3sRwLe3Va",
    "8194nk5LbLH",
    "82sE5b5pLXE",
    "8WUmhLawc2A",
    "aayBHfsNo7d",
    "ac26ZMwG7aT",
    "ARNzJeq3xxb",
    "B6ByNegPMKs",
    "b8cTxDM8gDG",
    "cV4RVeZvu5T",
    "D7G3Y4RVNrH",
    "D7N2EKCX4Sj",
    "dhjEzFoUFzH",
    "E9uDoFAP3SH",
    "e9zR4mvMWw7",
    "EDJbREhghzL",
    "EU6Fwq7SyZv",
    "fzynW3qQPVF",
    "GdvgFV5R1Z5",
    "gTV8FGcVJC9",
    "gxdoqLR6rwA",
    "gYvKGZ5eRqb",
    "gZ6f7yhEvPG",
    "HxpKQynjfin",
    "i5noydFURQK",
    "JeFG25nYj2p",
    "JF19kD82Mey",
    "jh4fc5c5qoQ",
    "JmbYfDe2QKZ",
    "jtcxE69GiFV",
    "kEZ7cmS4wCh",
    "mJXqzFtmKg4",
    "oLBMNvg9in8",
    "p5wJjkQkbXX",
    "pa4otMbVnkk",
    "pLe4wQe7qrG",
    "Pm6F8kyY3z2",
    "pRbA3pwrgk9",
    "PuKPg4mmafe",
    "PX4nDJXEHrG",
    "q9vSo1VnCiC",
    "qoiz87JEwZ2",
    "QUCTc6BB5sX",
    "r1Q1Z4BcV1o",
    "r47D5H71a5s",
    "rPc6DW4iMge",
    "RPmz2sHmrrY",
    "rqfALeAoiTq",
    "s8pcmisQ38h",
    "S9hNv5qa7GM",
    "sKLMLpTHeUy",
    "SN83YJsR3w2",
    "sT4fr6TAbpF",
    "TbHJrupSAjP",
    "ULsKaCPVFJR",
    "uNb9QFRL6hY",
    "ur6pFq6Qu1A",
    "UwV83HsGsw3",
    "Uxmj2M2itWa",
    "V2XKFyX4ASd",
    "VFuaQ6m2Qom",
    "VLzqgDo317F",
    "Vt2qJdWjCF2",
    "VVfe2KiqLaN",
    "Vvot9Ly1tCj",
    "vyrNrziPKCB",
    "VzqfbhrpDEA",
    "wc2JMjhGNzB",
    "WYY7iVyf5p8",
    "X7HyMhZNoso",
    "x8F5xyUWy9e",
    "XcA2TqTSSAj",
    "YFuZgdQ5vWj",
    "YmJkqBEsHnH",
    "yqstnuAEVhm",
    "YVUC4YcDtcY",
    "Z6MFQCViBuw",
    "ZMojNkEp431",
    "zsNo4HB9uLZ",
]


def invert_se3(se3_matrix):
    """Invert a batch of SE(3) transformations.

    Parameters
    ----------
    se3_matrix (batch_size x 4 x 4 tensor)
        The SE(3) transformations to invert.

    Returns
    -------
    batch_size x 4 x 4 tensor
        The inverted transformations.
    """

    R = se3_matrix[:, :3, :3]
    t = se3_matrix[:, :3, 3:]

    inverse = (
        torch.eye(4, device=se3_matrix.device)
        .unsqueeze(dim=0)
        .repeat(se3_matrix.shape[0], 1, 1)
    )

    inverse_R = rearrange(R, "b i j -> b j i")
    inverse[:, :3, :3] = inverse_R
    inverse[:, :3, 3:] = -inverse_R @ t

    return inverse


def invert_intrinsic_matrix(intrinsic_matrices):
    """Invert a batch of intrinsic matrices with zero skew.

    Parameters
    ----------
    intrinsic_matrices (batch_size x 3 x 3 tensor):
        The intrinsic matrices to invert

    Returns
    -------
    batch_size x 3 x 3 tensor:
        The inverted intrinsic matrices.
    """

    inverse = (
        torch.eye(3, device=intrinsic_matrices.device)
        .unsqueeze(dim=0)
        .repeat(intrinsic_matrices.shape[0], 1, 1)
    )
    inverse[:, 0, 0] = 1 / intrinsic_matrices[:, 0, 0]
    inverse[:, 1, 1] = 1 / intrinsic_matrices[:, 1, 1]
    inverse[:, 0, 2] = -intrinsic_matrices[:, 0, 2] / intrinsic_matrices[:, 0, 0]
    inverse[:, 1, 2] = -intrinsic_matrices[:, 1, 2] / intrinsic_matrices[:, 1, 1]

    return inverse


def calculate_pairwise_overlaps(depths, cam_to_world_matrices, intrinsic_matrices):
    """Calculate the overlaps between each pair of the images defined.

    Parameters
    ----------
    depths (n_imgs x height x width tensor):
        the depth maps in metres for the images.

    cam_to_world_matrices (n_imgs x 4 x 4 tensor):
        The camera to world matrices for the images.

    intrinsic_matrices (n_imgs x 3 x 3 tensor):
        The intrinsic matrices for the images.

    Returns
    -------
    n_imgs x n_imgs tensor:
        The overlaps between each pair of images. Element (i, j) is the fraction of pixels in image i that are visible (i.e. in the image and not occluded) in image j.
    """
    n_imgs, height, width = depths.shape

    overlaps = torch.eye(n_imgs, device=depths.device)  # the diagonals will be 1.0

    world_to_cam_matrices = invert_se3(cam_to_world_matrices)
    inverse_intrinsics = invert_intrinsic_matrix(intrinsic_matrices)

    # pixel coordinates for each pixel in the image; this is a (height x width x 2) array
    # element (i, j, 0) is the x coordinate of pixel (i, j), and element (i, j, 1) is the y coordinate of pixel (i, j)
    all_pixel_coordinates = np.stack(
        np.meshgrid(np.arange(width), height - 1 - np.arange(height), indexing="xy"),
        axis=2,
    )
    all_homogeneous_pixel_coordinates = np.concatenate(
        [all_pixel_coordinates, np.ones((height, width, 1))], axis=2
    )
    all_homogeneous_pixel_coordinates = torch.tensor(
        all_homogeneous_pixel_coordinates, device=depths.device, dtype=torch.float
    )

    for i in tqdm(range(n_imgs), desc="Calculating overlaps"):
        # Calculate the 3D points for each pixel in the image
        source_image_coordinates = torch.einsum(
            "tu,hwu->hwt", inverse_intrinsics[i], all_homogeneous_pixel_coordinates
        )
        source_image_coordinates[
            :, :, 2
        ] = -1.0  # in camera coordinates, z-coordinate points into camera
        source_camera_coordinates = torch.cat(
            [
                source_image_coordinates * depths[i, :, :, None],
                torch.ones((height, width, 1), device=depths.device),
            ],
            dim=2,
        )

        world_coordinates = torch.einsum(
            "fg,hwg->hwf", cam_to_world_matrices[i], source_camera_coordinates
        )

        # Project the 3D points onto the other images
        camera_coordinates = torch.einsum(
            "bfg,hwg->bhwf", world_to_cam_matrices, world_coordinates
        )
        camera_coordinates[:, :, :, 2][
            camera_coordinates[:, :, :, 2] == 0
        ] = 1e-6  # avoid divide by zero
        image_coordinates = -camera_coordinates[:, :, :, :3] / (
            camera_coordinates[:, :, :, 2:3]
        )
        image_coordinates[
            :, :, :, 2
        ] = 1.0  # in camera coordinaes, z-coordinate points into camera

        pixel_coordinates = torch.einsum(
            "btu,bhwu->bhwt", intrinsic_matrices, image_coordinates
        )

        # Check visibility
        rounded_pixel_idxs = rearrange(
            torch.clamp(
                torch.round(pixel_coordinates[:, :, :, 0]), min=0, max=width - 1
            )
            + (
                height
                - 1
                - torch.clamp(
                    torch.round(pixel_coordinates[:, :, :, 1]), min=0, max=height - 1
                )
            )
            * width,
            "b h w -> b (h w)",
        ).to(torch.int64)

        depths_at_pixel_coordinates = torch.gather(
            rearrange(depths, "b h w -> b (h w)"), dim=1, index=rounded_pixel_idxs
        )
        depths_at_pixel_coordinates = rearrange(
            depths_at_pixel_coordinates, "b (h w) -> b h w", h=height, w=width
        )
        visible = (
            (pixel_coordinates[:, :, :, 0] >= 0)
            & (pixel_coordinates[:, :, :, 0] < width)
            & (pixel_coordinates[:, :, :, 1] >= 0)
            & (pixel_coordinates[:, :, :, 1] < height)
            & (-camera_coordinates[:, :, :, 2] > 0)
            & (-camera_coordinates[:, :, :, 2] < depths_at_pixel_coordinates + 1e-2)
        )

        overlaps[i] = torch.sum(visible, dim=(1, 2)) / (height * width)

    return overlaps


def main(opts):
    for scene in ALL_SCENES:
        print(f"Processing scene {scene}...")

        scene_dir = os.path.join(opts.raw_dataset_dir, scene)
        depth_dir = os.path.join(scene_dir, "depth")

        depth_files = sorted(os.listdir(depth_dir))

        # Load depth images
        depths = []
        for depth_file in tqdm(depth_files, "Loading depth images"):
            depths.append(np.load(os.path.join(depth_dir, depth_file)))

        depths = np.stack(depths, axis=0)
        if depths.shape[1:] != (opts.img_height, opts.img_width):
            depths = torch.tensor(depths).unsqueeze(dim=1)
            depths = (
                torch.nn.functional.interpolate(
                    depths, size=(opts.img_height, opts.img_width), mode="nearest-exact"
                )
                .squeeze(dim=1)
                .numpy()
            )

        # Load camera matrices
        camera_parameter_file_path = os.path.join(
            scene_dir, "undistorted_camera_parameters", "{}.conf".format(scene)
        )
        with open(camera_parameter_file_path, "r") as f:
            parameter_file_lines = f.readlines()

        current_intrinsics = None
        intrinsic_matrices = np.zeros((len(depth_files), 3, 3))
        cam_to_world_matrices = np.zeros((len(depth_files), 4, 4))

        for line in tqdm(parameter_file_lines, "Loading camera parameters"):
            if line.startswith("intrinsics"):
                current_intrinsics = rearrange(
                    np.array(
                        [
                            float(x)
                            for x in line.split(" ")[1:]
                            if not (x.isspace() or len(x) == 0)
                        ]
                    ),
                    "(h w) -> h w",
                    h=3,
                    w=3,
                )
                current_intrinsics[:2, :] *= opts.img_height / 1024.0
            elif line.startswith("scan"):
                scan_line = line.split(" ")[1:]
                depth_file_name = scan_line[0].replace(".png", ".npy")

                if depth_file_name not in depth_files:
                    continue

                depth_file_index = depth_files.index(depth_file_name)

                if intrinsic_matrices[depth_file_index].any():
                    raise ValueError(
                        "Intrinsics already set for this image {}".format(
                            depth_file_name
                        )
                    )
                if cam_to_world_matrices[depth_file_index].any():
                    raise ValueError(
                        "Cam to world matrix already set for this image {}".format(
                            depth_file_name
                        )
                    )

                intrinsic_matrices[depth_file_index] = current_intrinsics
                cam_to_world_matrices[depth_file_index] = rearrange(
                    np.array([float(x) for x in scan_line[2:]]),
                    "(h w) -> h w",
                    h=4,
                    w=4,
                )
            else:
                continue

        depths = torch.tensor(depths, dtype=torch.float, device=opts.device)
        intrinsic_matrices = torch.tensor(
            intrinsic_matrices, dtype=torch.float, device=opts.device
        )
        cam_to_world_matrices = torch.tensor(
            cam_to_world_matrices, dtype=torch.float, device=opts.device
        )

        scene_overlaps = calculate_pairwise_overlaps(
            depths, cam_to_world_matrices, intrinsic_matrices
        )

        np.save(
            os.path.join(scene_dir, "img_overlaps.npy"), scene_overlaps.cpu().numpy()
        )


if __name__ == "__main__":
    parser = configargparse.ArgumentParser()
    parser.add_argument(
        "--raw_dataset_dir",
        type=str,
        default="data/matterport3d",
    )
    parser.add_argument(
        "--img_height",
        type=int,
        default=32,
        help="Image height to use for calculations",
    )
    parser.add_argument(
        "--img_width", type=int, default=40, help="Image width to use for calculations"
    )
    parser.add_argument(
        "--device", type=str, default="cuda:0", help="Device to use for calculations"
    )
    opts = parser.parse_args()
    opts.device = torch.device(opts.device)
    main(opts)

scripts/data_preparation_scripts/matterport3d_calc_img_overlap.py
dakloepfer/locus
"""Use the triangle meshes and region_segmentations information to for each image give the semantic segmentation labels for each pixel."""

import os
import configargparse
from tqdm import tqdm
import json
import csv

import numpy as np
from PIL import Image
import open3d

from einops import rearrange

ALL_SCENES = [
    "17DRP5sb8fy",
    "1LXtFkjw3qL",
    "1pXnuDYAj8r",
    "29hnd4uzFmX",
    "2azQ1b91cZZ",
    "2n8kARJN3HM",
    "2t7WUuJeko7",
    "5LpN3gDmAk7",
    "5q7pvUzZiYa",
    "5ZKStnWn8Zo",
    "759xd9YjKW5",
    "7y3sRwLe3Va",
    "8194nk5LbLH",
    "82sE5b5pLXE",
    "8WUmhLawc2A",
    "aayBHfsNo7d",
    "ac26ZMwG7aT",
    "ARNzJeq3xxb",
    "B6ByNegPMKs",
    "b8cTxDM8gDG",
    "cV4RVeZvu5T",
    "D7G3Y4RVNrH",
    "D7N2EKCX4Sj",
    "dhjEzFoUFzH",
    "E9uDoFAP3SH",
    "e9zR4mvMWw7",
    "EDJbREhghzL",
    "EU6Fwq7SyZv",
    "fzynW3qQPVF",
    "GdvgFV5R1Z5",
    "gTV8FGcVJC9",
    "gxdoqLR6rwA",
    "gYvKGZ5eRqb",
    "gZ6f7yhEvPG",
    "HxpKQynjfin",
    "i5noydFURQK",
    "JeFG25nYj2p",
    "JF19kD82Mey",
    "jh4fc5c5qoQ",
    "JmbYfDe2QKZ",
    "jtcxE69GiFV",
    "kEZ7cmS4wCh",
    "mJXqzFtmKg4",
    "oLBMNvg9in8",
    "p5wJjkQkbXX",
    "pa4otMbVnkk",
    "pLe4wQe7qrG",
    "Pm6F8kyY3z2",
    "pRbA3pwrgk9",
    "PuKPg4mmafe",
    "PX4nDJXEHrG",
    "q9vSo1VnCiC",
    "qoiz87JEwZ2",
    "QUCTc6BB5sX",
    "r1Q1Z4BcV1o",
    "r47D5H71a5s",
    "rPc6DW4iMge",
    "RPmz2sHmrrY",
    "rqfALeAoiTq",
    "s8pcmisQ38h",
    "S9hNv5qa7GM",
    "sKLMLpTHeUy",
    "SN83YJsR3w2",
    "sT4fr6TAbpF",
    "TbHJrupSAjP",
    "ULsKaCPVFJR",
    "uNb9QFRL6hY",
    "ur6pFq6Qu1A",
    "UwV83HsGsw3",
    "Uxmj2M2itWa",
    "V2XKFyX4ASd",
    "VFuaQ6m2Qom",
    "VLzqgDo317F",
    "Vt2qJdWjCF2",
    "VVfe2KiqLaN",
    "Vvot9Ly1tCj",
    "vyrNrziPKCB",
    "VzqfbhrpDEA",
    "wc2JMjhGNzB",
    "WYY7iVyf5p8",
    "X7HyMhZNoso",
    "x8F5xyUWy9e",
    "XcA2TqTSSAj",
    "YFuZgdQ5vWj",
    "YmJkqBEsHnH",
    "yqstnuAEVhm",
    "YVUC4YcDtcY",
    "Z6MFQCViBuw",
    "ZMojNkEp431",
    "zsNo4HB9uLZ",
]

# NOTE: had to clean up some of the labels; the only one that is maybe not completely obvious is replacing the "roof or floor / other room" with "ceiling"
MPCAT40_TO_INDEX = {
    "void": 0,
    "wall": 1,
    "floor": 2,
    "chair": 3,
    "door": 4,
    "table": 5,
    "picture": 6,
    "cabinet": 7,
    "cushion": 8,
    "window": 9,
    "sofa": 10,
    "bed": 11,
    "curtain": 12,
    "chest_of_drawers": 13,
    "plant": 14,
    "sink": 15,
    "stairs": 16,
    "ceiling": 17,
    "toilet": 18,
    "stool": 19,
    "towel": 20,
    "mirror": 21,
    "tv_monitor": 22,
    "shower": 23,
    "column": 24,
    "bathtub": 25,
    "counter": 26,
    "fireplace": 27,
    "lighting": 28,
    "beam": 29,
    "railing": 30,
    "shelving": 31,
    "blinds": 32,
    "gym_equipment": 33,
    "seating": 34,
    "board_panel": 35,
    "furniture": 36,
    "appliances": 37,
    "clothes": 38,
    "objects": 39,
    "misc": 40,
    "unlabeled": 41,
}
STUFF_CLASSES = {
    "void": 0,
    "wall": 1,
    "floor": 2,
    "curtain": 3,
    "stairs": 4,
    "ceiling": 5,
    "mirror": 6,
    "shower": 7,
    "column": 8,
    "beam": 9,
    "railing": 10,
    "shelving": 11,
    "blinds": 12,
    "board_panel": 13,
    "misc": 14,
    "unlabeled": 15,
}

STUFF_CLASSES_V1 = {
    "void": 0,
    "wall": 1,
    "floor": 2,
    "ceiling": 3,
    "column": 4,
    "beam": 5,
    "objects": 6,
    "misc": 7,
    "unlabeled": 8,
}

INVALID_ID = open3d.t.geometry.RaycastingScene.INVALID_ID  # 4294967295
INVALID_CATEGORY = 41  # the unlabeled index

CATEGORY_TO_HEX = {
    0: "#ffffff",
    1: "#aec7e8",
    2: "#708090",
    3: "#98df8a",
    4: "#c5b0d5",
    5: "#ff7f0e",
    6: "#d62728",
    7: "#1f77b4",
    8: "#bcbd22",
    9: "#ff9896",
    10: "#2ca02c",
    11: "#e377c2",
    12: "#de9ed6",
    13: "#9467bd",
    14: "#8ca252",
    15: "#843c39",
    16: "#9edae5",
    17: "#9c9ede",
    18: "#e7969c",
    19: "#637939",
    20: "#8c564b",
    21: "#dbdb8d",
    22: "#d6616b",
    23: "#cedb9c",
    24: "#e7ba52",
    25: "#393b79",
    26: "#a55194",
    27: "#ad494a",
    28: "#b5cf6b",
    29: "#5254a3",
    30: "#bd9e39",
    31: "#c49c94",
    32: "#f7b6d2",
    33: "#6b6ecf",
    34: "#ffbb78",
    35: "#c7c7c7",
    36: "#8c6d31",
    37: "#e7cb94",
    38: "#ce6dbd",
    39: "#17becf",
    40: "#7f7f7f",
    41: "#000000",
}


def hex_to_rgb(hex_string):
    """Convert a hex color string to an RGB color.

    Parameters
    ----------
    hex_string (str)
        The hex color string.

    Returns
    -------
    numpy array length 3
        The RGB color.
    """
    hex_string = hex_string.lstrip("#")
    return np.array(
        [int(hex_string[i : i + 2], base=16) for i in (0, 2, 4)], dtype=np.uint8
    )


def visualise_segmentation(segmentation_map, save_path):
    image = np.zeros(
        (segmentation_map.shape[0], segmentation_map.shape[1], 3), dtype=np.uint8
    )
    for i in range(segmentation_map.shape[0]):
        for j in range(segmentation_map.shape[1]):
            image[i, j, :] = hex_to_rgb(
                CATEGORY_TO_HEX[segmentation_map[i, j] % len(CATEGORY_TO_HEX)]
            )

    # save image
    img = Image.fromarray(image)
    img.save(save_path)


def invert_se3(se3_matrix):
    """Invert an SE(3) transformations.

    Parameters
    ----------
    se3_matrix (4 x 4 tensor)
        The SE(3) transformations to invert.

    Returns
    -------
    4 x 4 tensor
        The inverted transformation.
    """

    R = se3_matrix[:3, :3]
    t = se3_matrix[:3, 3:]

    inverse = np.eye(4, dtype=se3_matrix.dtype)

    inverse_R = R.T
    inverse[:3, :3] = inverse_R
    inverse[:3, 3:] = -inverse_R @ t

    return inverse


def calc_pixel_triangle_mapping(raycaster, camera_pose, intrinsic_matrix, opts):
    # the camera pose is the camera-to-world transformation, so invert_se3(camera_pose) is the world-to-camera transformation.
    # But Matterport3d uses a different camera-coordinate system than Open3D (Matterport uses y up, z into camera, x right, Open3D uses y down, z out of camera, x right)
    # so to get the world-to-Open3D camera transformation, we need to flip the y and z axes.
    extrinsic_matrix = np.array(
        [[1.0, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
    ) @ invert_se3(camera_pose)
    extrinsic_matrix = open3d.core.Tensor(extrinsic_matrix)

    # also, the principal point for y is now measured from the top of the image, not the bottom
    intrinsic_matrix[1, 2] = opts.img_height - intrinsic_matrix[1, 2]
    intrinsic_matrix = open3d.core.Tensor(intrinsic_matrix)

    rays = open3d.t.geometry.RaycastingScene.create_rays_pinhole(
        intrinsic_matrix, extrinsic_matrix, opts.img_width, opts.img_height
    )

    return raycaster.cast_rays(rays)["primitive_ids"].numpy()


def load_fsegs_and_segs_to_cat(
    opts, all_regions, scene_dir, raw_semantic_label_to_mpcat40
):
    """Creates maps from triangle indices to segment indices and segment indices to category indices.

    Returns:
    --------
    triangle_to_segment ((n_triangles, ) numpy array):
        the segment index for each triangle

    segment_to_category ((n_segments, ) numpy array):
        the category index for each segment
    """

    # load triangle -> segment mappings
    # this should be a numpy array of shape (num_triangles,), where each entry is the segment index of the triangle

    n_regions = len(all_regions)
    all_fsegs = []
    all_seggroups = []
    for region in all_regions:
        with open(
            os.path.join(scene_dir, "region_segmentations", region + ".fsegs.json")
        ) as f:
            all_fsegs.append(json.load(f)["segIndices"])

        with open(
            os.path.join(scene_dir, "region_segmentations", region + ".semseg.json")
        ) as f:
            all_seggroups.append(json.load(f)["segGroups"])

    # shift the segment indices so that they are unique across meshes
    for i in range(n_regions):
        all_fsegs[i] = np.array(all_fsegs[i], dtype=np.int32)

        if i > 0:
            index_shift = np.max(all_fsegs[i - 1]) + 1
            all_fsegs[i] += index_shift

            for j in range(len(all_seggroups[i])):
                # go through all object in the region
                all_seggroups[i][j]["segments"] = (
                    np.array(all_seggroups[i][j]["segments"]) + index_shift
                )

    # add a new segment for invalid ids -- basically all rays that don't hit a triangle hit this appended phantom triangle with a phantom segment
    # this might end up not being needed if they use the "void" category for all this, but I can't be sure of that
    all_fsegs.append(np.array([all_fsegs[-1].max() + 1]))
    triangle_to_segment = np.concatenate(all_fsegs, axis=0)

    n_segments = triangle_to_segment.max() + 1

    # a bunch of segments have no object (don't appear in the segGroups file), so start out with putting them all as INVALID_CATEGORY
    if opts.instance_segmentation:
        segment_to_category = np.ones(n_segments, dtype=np.int32) * (-1)
    else:
        segment_to_category = np.ones(n_segments, dtype=np.int32) * INVALID_CATEGORY

    for i in range(n_regions):
        if opts.instance_segmentation and i > 0:
            index_shift = np.max(segment_to_category) + 1
        else:
            index_shift = len(STUFF_CLASSES)

        n_stuff_objects = 0
        for j in range(len(all_seggroups[i])):
            if opts.instance_segmentation:
                # different categories for individual instances of objects of the same category
                # add the max category index to the object id to get a unique category index for each instance even across regions

                if (
                    raw_semantic_label_to_mpcat40[all_seggroups[i][j]["label"]]
                    in STUFF_CLASSES.keys()
                ):  # stuff objects are all mapped to their semantic class
                    category_idx = STUFF_CLASSES[
                        raw_semantic_label_to_mpcat40[all_seggroups[i][j]["label"]]
                    ]
                    n_stuff_objects += 1
                else:  # not stuff, so every object is its own category
                    category_idx = (
                        all_seggroups[i][j]["objectId"] + index_shift - n_stuff_objects
                    )

            else:
                category_idx = MPCAT40_TO_INDEX[
                    raw_semantic_label_to_mpcat40[all_seggroups[i][j]["label"]]
                ]
            segment_to_category[all_seggroups[i][j]["segments"]] = category_idx

    if opts.instance_segmentation:
        # change the category for invalid objects (the ones that are mapped to -1) to a positive number, ie the max category index + 1
        segment_to_category[segment_to_category < 0] = np.max(segment_to_category) + 1

    return triangle_to_segment, segment_to_category


def load_meshes(all_regions, scene_dir):
    # load the meshes
    all_meshes = []
    for region in all_regions:
        mesh = open3d.io.read_triangle_mesh(
            os.path.join(scene_dir, "region_segmentations", region + ".ply")
        )
        all_meshes.append(mesh)

    # combine meshes, taking care of segment and category mappings
    # (create RayMeshIntersector here)
    scene_mesh = all_meshes[0]
    for mesh in all_meshes[1:]:
        scene_mesh += mesh

    return scene_mesh


def main(opts):
    raw_semantic_label_to_mpcat40 = {}

    with open(
        os.path.join(opts.category_mapping_dir, "category_mapping.tsv"), "r"
    ) as tsvfile:
        reader = csv.reader(tsvfile, delimiter="\t")

        # Skip header row
        next(reader)
        for row in reader:
            # Map second column to last column
            raw_semantic_label_to_mpcat40[row[1]] = row[-1]

    for scene in ALL_SCENES:
        print(f"Processing scene {scene}")
        scene_dir = os.path.join(opts.raw_dataset_dir, scene)

        if opts.instance_segmentation:
            if not os.path.exists(os.path.join(scene_dir, "object_maps")):
                os.makedirs(os.path.join(scene_dir, "object_maps"))

        else:
            if not os.path.exists(os.path.join(scene_dir, "semantic_maps")):
                os.makedirs(os.path.join(scene_dir, "semantic_maps"))

        all_files = os.listdir(os.path.join(scene_dir, "region_segmentations"))
        all_regions = list(
            set([f.split(".")[0] for f in all_files if f.startswith("region")])
        )
        all_regions.sort()

        triangle_to_segment, segment_to_category = load_fsegs_and_segs_to_cat(
            opts, all_regions, scene_dir, raw_semantic_label_to_mpcat40
        )
        scene_mesh = load_meshes(all_regions, scene_dir)

        # For some reason open3d has more than one TriangleMesh class
        raycaster_mesh = open3d.t.geometry.TriangleMesh(
            open3d.core.Tensor(np.asarray(scene_mesh.vertices).astype(np.float32)),
            open3d.core.Tensor(np.asarray(scene_mesh.triangles).astype(np.float32)),
        )

        raycaster = open3d.t.geometry.RaycastingScene()
        raycaster.add_triangles(raycaster_mesh)

        # check that triangle_to_segment has length num_triangles + 1 (for the phantom triangle)
        assert (
            triangle_to_segment.shape[0]
            == np.asarray(scene_mesh.triangles).shape[0] + 1
        )
        # check that segment_to_category has length num_segments (including the phantom segment, which is already included in triangle_to_segment)
        assert segment_to_category.shape[0] == np.max(triangle_to_segment) + 1

        # load camera poses
        camera_parameter_file_path = os.path.join(
            scene_dir, "undistorted_camera_parameters", "{}.conf".format(scene)
        )

        with open(camera_parameter_file_path, "r") as f:
            params_file_lines = f.readlines()

        current_intrinsics = None
        image_names = []
        camera_poses = []
        intrinsic_matrices = []

        all_used_img_files = os.listdir(
            os.path.join(scene_dir, "rgb_highres" if opts.img_height == 448 else "rgb")
        )

        for line in params_file_lines:
            if line.startswith("intrinsics"):
                intrinsics_line = [
                    i
                    for i in line.strip().split(" ")
                    if not (i.isspace() or len(i) == 0)
                ]

                current_intrinsics = np.array([float(i) for i in intrinsics_line[1:]])
                current_intrinsics = rearrange(
                    current_intrinsics, "(h w) -> h w", h=3, w=3
                )
                # adjust for larger pixels
                current_intrinsics[0] *= opts.img_width / 1280
                current_intrinsics[1] *= opts.img_height / 1024

            elif line.startswith("scan"):
                scan_line = line.split(" ")[1:]

                img_file_name = scan_line[1].replace(".jpg", ".npy")

                # I filter out some images that had too few good depth values to inpaint the depth map
                if img_file_name not in all_used_img_files:
                    continue

                cam_pose = np.array([float(i) for i in scan_line[2:]])
                cam_pose = rearrange(cam_pose, "(h w) -> h w", h=4, w=4)

                image_names.append(img_file_name[:-4])
                intrinsic_matrices.append(current_intrinsics)
                camera_poses.append(cam_pose)

            else:
                continue

        # do this for each image, as I need to save each image's category indices separately
        for i in tqdm(
            range(len(camera_poses)),
            desc="Calculating segmentation maps for scene {}".format(scene),
        ):
            camera_pose = camera_poses[i]
            intrinsic_matrix = intrinsic_matrices[i]
            image_name = image_names[i]

            triangle_indices = calc_pixel_triangle_mapping(
                raycaster, camera_pose, intrinsic_matrix, opts
            )
            triangle_indices[triangle_indices == INVALID_ID] = triangle_to_segment.max()

            segment_indices = triangle_to_segment[triangle_indices]
            category_indices = segment_to_category[segment_indices]

            if opts.debug:
                if not os.path.exists("./test_segm/{}".format(scene)):
                    os.makedirs("./test_segm/{}".format(scene))
                visualise_segmentation(
                    category_indices,
                    "./test_segm/{}/".format(scene) + image_name + ".png",
                )

                command = input(
                    'Press "Enter" to continue to next image, "ns" to go to the next scene...'
                )
                if command == "ns":
                    break
            else:
                if opts.instance_segmentation:
                    np.save(
                        os.path.join(scene_dir, "object_maps", image_name),
                        category_indices,
                    )
                else:
                    np.save(
                        os.path.join(scene_dir, "semantic_maps", image_name),
                        category_indices,
                    )

    print("Done!")


if __name__ == "__main__":
    parser = configargparse.ArgumentParser()
    parser.add("-d", "--debug", action="store_true", help="debug mode", default=False)
    parser.add(
        "-i",
        "--instance_segmentation",
        action="store_true",
        help="segment individual instances of objects, not categories",
        default=True,
    )
    parser.add_argument(
        "--raw_dataset_dir",
        type=str,
        default="/data/matterport3d",
    )
    parser.add_argument(
        "--category_mapping_dir",
        type=str,
        default="./scripts/data_preparation_scripts",
    )
    parser.add_argument("--img_width", type=int, default=320)
    parser.add_argument("--img_height", type=int, default=256)
    opts = parser.parse_args()
    main(opts)

scripts/data_preparation_scripts/matterport3d_compute_segmentations.py
dakloepfer/locus
"""A script to generate better depth maps than by inpainting; this first uses the 3D meshes to fill in missing values, and only then uses inpainting to fill in the rest."""

import os
import configargparse
import zipfile
from tqdm import tqdm

import open3d
import numpy as np
from scipy import sparse
import cv2
from PIL import Image
from einops import rearrange

ALL_SCENES = [
    "17DRP5sb8fy",
    "1LXtFkjw3qL",
    "1pXnuDYAj8r",
    "29hnd4uzFmX",
    "2azQ1b91cZZ",
    "2n8kARJN3HM",
    "2t7WUuJeko7",
    "5LpN3gDmAk7",
    "5q7pvUzZiYa",
    "5ZKStnWn8Zo",
    "759xd9YjKW5",
    "7y3sRwLe3Va",
    "8194nk5LbLH",
    "82sE5b5pLXE",
    "8WUmhLawc2A",
    "aayBHfsNo7d",
    "ac26ZMwG7aT",
    "ARNzJeq3xxb",
    "B6ByNegPMKs",
    "b8cTxDM8gDG",
    "cV4RVeZvu5T",
    "D7G3Y4RVNrH",
    "D7N2EKCX4Sj",
    "dhjEzFoUFzH",
    "E9uDoFAP3SH",
    "e9zR4mvMWw7",
    "EDJbREhghzL",
    "EU6Fwq7SyZv",
    "fzynW3qQPVF",
    "GdvgFV5R1Z5",
    "gTV8FGcVJC9",
    "gxdoqLR6rwA",
    "gYvKGZ5eRqb",
    "gZ6f7yhEvPG",
    "HxpKQynjfin",
    "i5noydFURQK",
    "JeFG25nYj2p",
    "JF19kD82Mey",
    "jh4fc5c5qoQ",
    "JmbYfDe2QKZ",
    "jtcxE69GiFV",
    "kEZ7cmS4wCh",
    "mJXqzFtmKg4",
    "oLBMNvg9in8",
    "p5wJjkQkbXX",
    "pa4otMbVnkk",
    "pLe4wQe7qrG",
    "Pm6F8kyY3z2",
    "pRbA3pwrgk9",
    "PuKPg4mmafe",
    "PX4nDJXEHrG",
    "q9vSo1VnCiC",
    "qoiz87JEwZ2",
    "QUCTc6BB5sX",
    "r1Q1Z4BcV1o",
    "r47D5H71a5s",
    "rPc6DW4iMge",
    "RPmz2sHmrrY",
    "rqfALeAoiTq",
    "s8pcmisQ38h",
    "S9hNv5qa7GM",
    "sKLMLpTHeUy",
    "SN83YJsR3w2",
    "sT4fr6TAbpF",
    "TbHJrupSAjP",
    "ULsKaCPVFJR",
    "uNb9QFRL6hY",
    "ur6pFq6Qu1A",
    "UwV83HsGsw3",
    "Uxmj2M2itWa",
    "V2XKFyX4ASd",
    "VFuaQ6m2Qom",
    "VLzqgDo317F",
    "Vt2qJdWjCF2",
    "VVfe2KiqLaN",
    "Vvot9Ly1tCj",
    "vyrNrziPKCB",
    "VzqfbhrpDEA",
    "wc2JMjhGNzB",
    "WYY7iVyf5p8",
    "X7HyMhZNoso",
    "x8F5xyUWy9e",
    "XcA2TqTSSAj",
    "YFuZgdQ5vWj",
    "YmJkqBEsHnH",
    "yqstnuAEVhm",
    "YVUC4YcDtcY",
    "Z6MFQCViBuw",
    "ZMojNkEp431",
    "zsNo4HB9uLZ",
]


def invert_se3(se3_matrix):
    """Invert an SE(3) transformations.

    Parameters
    ----------
    se3_matrix (4 x 4 tensor)
        The SE(3) transformations to invert.

    Returns
    -------
    4 x 4 tensor
        The inverted transformation.
    """

    R = se3_matrix[:3, :3]
    t = se3_matrix[:3, 3:]

    inverse = np.eye(4, dtype=se3_matrix.dtype)

    inverse_R = R.T
    inverse[:3, :3] = inverse_R
    inverse[:3, 3:] = -inverse_R @ t

    return inverse


def bad_depth_values(depth, min_depth=0.01, max_depth=20):
    return (depth < min_depth) | (depth > max_depth) | np.isnan(depth) | np.isinf(depth)


def inpaint_depth(depth, image, nyu_depth_dataset=False, alpha=1):
    """In-paint depth values using method from "Colorization Using Optimization" by Levin et al. (https://www.cs.huji.ac.il/w~yweiss/Colorization/), which was also used / adapted for depth inpainting in the NYU Depth Dataset (https://cs.nyu.edu/~silberman/datasets/nyu_depth_v2.html).

    Parameters
    ----------
    depth (height x width numpy array):
        the depth values with some missing / noisy values

    image (height x width x channels numpy array):
        the RGB image, whose intensities are used as a guide for the inpainting

    nyu_depth_dataset (bool, optional):
        Whether I should use the method from the NYU Depth Dataset code or the one from the original paper, by default False

    alpha (int, optional):
        Only used if nyu_depth_dataset=True, this is "a penalty value betwwen 0 and 1 for the current depth values", by default 1

    Returns
    -------
    height x width numpy array:
        the inpainted depth map
    """
    # NOTE: the Matlab code that I based this on used sparse matrices, but this should be equivalent and not too much slower

    window_radius = 1
    window_size = (2 * window_radius + 1) ** 2

    height, width = depth.shape
    n_pixels = height * width

    grayscale = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY).squeeze()

    depth_is_noise = bad_depth_values(depth)

    if depth_is_noise.sum() == 0:
        return depth

    if depth_is_noise.sum() >= 0.25 * n_pixels:
        raise ValueError("Not enough valid depth values to inpaint")

    # normalise depth
    max_depth = np.max(depth[~depth_is_noise])
    depth = depth / max_depth
    depth[depth > 1] = 1

    col_inds = np.zeros((n_pixels * window_size,))
    row_inds = np.zeros((n_pixels * window_size,))
    current_window_weight_vals = np.zeros((window_size,))
    vals = np.zeros(n_pixels * window_size)  # the values for the (1 - weights) matrix

    # NOTE if this takes too long (which I doubt), I could try to vectorise it; that would involve a bit more trickery with the outer_window_mask though to account for the fact that the window is not always the same size (at the edges of the image)
    # for now I just copy the for-loops from the Matlab code
    pixel_index = 0
    all_pixel_indices = np.arange(n_pixels).reshape(height, width)
    n_nonzero_vals = 0
    for i in range(height):
        for j in range(width):
            if ~depth_is_noise[i, j]:
                row_inds[n_nonzero_vals] = pixel_index
                col_inds[n_nonzero_vals] = all_pixel_indices[i, j]
                if nyu_depth_dataset:
                    vals[n_nonzero_vals] = alpha
                else:
                    vals[n_nonzero_vals] = 1

                n_nonzero_vals += 1
                pixel_index += 1
                continue

            current_window_index = 0
            for ii in range(
                max(0, i - window_radius), min(height, i + window_radius + 1)
            ):
                for jj in range(
                    max(0, j - window_radius), min(width, j + window_radius + 1)
                ):
                    if ii == i and jj == j:
                        continue

                    row_inds[n_nonzero_vals] = pixel_index
                    col_inds[n_nonzero_vals] = all_pixel_indices[ii, jj]
                    current_window_weight_vals[current_window_index] = grayscale[ii, jj]

                    n_nonzero_vals += 1
                    current_window_index += 1

            pixel_gs_value = grayscale[i, j]

            # keep that around to include it in the variance calculation
            current_window_weight_vals[current_window_index] = pixel_gs_value

            window_gs_variance = np.var(current_window_weight_vals)

            csig = 0.6 * window_gs_variance
            mgv = np.min(
                (current_window_weight_vals[:current_window_index] - pixel_gs_value)
                ** 2
            )

            if csig < (-mgv / np.log(0.01)):
                csig = -mgv / np.log(0.01)
            if csig < 2.2e-6:
                csig = 2.2e-6

            # this is not quite the formula that they use in the paper, but it's the one they use in their code
            current_window_weight_vals[:current_window_index] = np.exp(
                -(
                    (current_window_weight_vals[:current_window_index] - pixel_gs_value)
                    ** 2
                )
                / csig
            )
            current_window_weight_vals[:current_window_index] /= np.sum(
                current_window_weight_vals[:current_window_index]
            )

            # the -weights part in the (1-weights) matrix
            vals[
                n_nonzero_vals - current_window_index : n_nonzero_vals
            ] = -current_window_weight_vals[:current_window_index]

            # add the identity part of the (1-weights) matrix
            row_inds[n_nonzero_vals] = pixel_index
            col_inds[n_nonzero_vals] = all_pixel_indices[i, j]
            vals[n_nonzero_vals] = 1
            n_nonzero_vals += 1

            pixel_index += 1

    # NOTE: I would have thought that the equation should be slightly different, namely if the weights matrix is filled properly for all, then for the depth_is_noise values (1 - weights.T) @ (1 - weights) @ new_depths = 0, (with for the non-noise values still new_depths = old depths)
    # Here I solve (1 - weights) @ new_depths = 0 (with for the non-noise values still new_depths = old depths), any solution of which also solves the above equation, but maybe there is no solution to this equation but to the correct one. If I run into any errors to that effect, I might want to change this. Since this is a linear system of equations which seem pretty well-behaved, I don't expect something like this though.

    # for the non-noise pixels, if nyu_depth_dataset=False, only the diagonal is non-zero
    A = sparse.coo_matrix(
        (vals[:n_nonzero_vals], (row_inds[:n_nonzero_vals], col_inds[:n_nonzero_vals])),
        shape=(n_pixels, n_pixels),
    )
    b = np.zeros(n_pixels)
    b[(~depth_is_noise).flatten()] = depth[~depth_is_noise].flatten()

    if nyu_depth_dataset:
        # This way of handling alpha is not what the code provided by the NYU Depths authors is doing, but what they are doing seems very weird; I think their code is missing a line to skip non-noise pixels in the loop above
        # If I am correct, then I think alpha=1 is simply the original code, while alpha=0 allows the known depths to be ignored completely
        b *= alpha

    new_depths = sparse.linalg.spsolve(A.tocsc(), b)

    return rearrange(new_depths * max_depth, "(h w) -> h w", h=height, w=width)


def calc_depths_from_raycaster(raycaster, camera_pose, intrinsic_matrix, opts):
    # the camera pose is the camera-to-world transformation, so invert_se3(camera_pose) is the world-to-camera transformation.
    # But Matterport3d uses a different camera-coordinate system than Open3D (Matterport uses y up, z into camera, x right, Open3D uses y down, z out of camera, x right)
    # so to get the world-to-Open3D camera transformation, we need to flip the y and z axes.
    extrinsic_matrix = np.array(
        [[1.0, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
    ) @ invert_se3(camera_pose)
    extrinsic_matrix = open3d.core.Tensor(extrinsic_matrix)

    # also, the principal point for y is now measured from the top of the image, not the bottom
    intrinsic_matrix[1, 2] = opts.img_height - intrinsic_matrix[1, 2]
    px = round(intrinsic_matrix[0, 2])
    py = round(intrinsic_matrix[1, 2])
    intrinsic_matrix = open3d.core.Tensor(intrinsic_matrix)

    rays = open3d.t.geometry.RaycastingScene.create_rays_pinhole(
        intrinsic_matrix, extrinsic_matrix, opts.img_width, opts.img_height
    ).numpy()

    # normalise ray direction vectors to unit length
    rays[:, :, 3:] = rays[:, :, 3:] / np.linalg.norm(
        rays[:, :, 3:], axis=-1, keepdims=True
    )

    ray_directions = rays[:, :, 3:]
    center_ray_direction = ray_directions[py, px]

    tensor_rays = open3d.core.Tensor(rays)

    distances_to_hit = raycaster.cast_rays(tensor_rays)["t_hit"].numpy()
    depth_values = distances_to_hit * np.einsum(
        "i,hwi->hw", center_ray_direction, ray_directions
    )  # project the ray direction onto the center ray direction to get the depth values

    # set depth values to 0 for pixels that are not hit
    depth_values[depth_values > 100] = 0

    return depth_values


def load_raycaster(scene_dir):
    all_mesh_files = os.listdir(os.path.join(scene_dir, "region_segmentations"))
    all_regions = list(
        set([f.split(".")[0] for f in all_mesh_files if f.startswith("region")])
    )
    all_regions.sort()

    # load the meshes
    all_meshes = []
    for region in all_regions:
        mesh = open3d.io.read_triangle_mesh(
            os.path.join(scene_dir, "region_segmentations", region + ".ply")
        )
        all_meshes.append(mesh)

    # combine meshes, taking care of segment and category mappings
    # (create RayMeshIntersector here)
    scene_mesh = all_meshes[0]
    for mesh in all_meshes[1:]:
        scene_mesh += mesh

    # For some reason open3d has more than one TriangleMesh shape
    raycaster_mesh = open3d.t.geometry.TriangleMesh(
        open3d.core.Tensor(np.asarray(scene_mesh.vertices).astype(np.float32)),
        open3d.core.Tensor(np.asarray(scene_mesh.triangles).astype(np.float32)),
    )

    raycaster = open3d.t.geometry.RaycastingScene()
    raycaster.add_triangles(raycaster_mesh)

    return raycaster


def read_cam_parameters(camera_parameter_file_path):
    with open(camera_parameter_file_path, "r") as f:
        params_file_lines = f.readlines()

    current_intrinsics = None
    image_files = []
    depth_files = []
    camera_poses = []
    intrinsic_matrices = []

    for line in params_file_lines:
        if line.startswith("intrinsics"):
            intrinsics_line = [
                i for i in line.strip().split(" ") if not (i.isspace() or len(i) == 0)
            ]

            current_intrinsics = np.array([float(i) for i in intrinsics_line[1:]])
            current_intrinsics = rearrange(current_intrinsics, "(h w) -> h w", h=3, w=3)
            # adjust for larger pixels
            current_intrinsics[0] *= opts.img_width / 1280
            current_intrinsics[1] *= opts.img_height / 1024

        elif line.startswith("scan"):
            scan_line = line.split(" ")[1:]

            img_file_name = scan_line[1]
            depth_file_name = scan_line[0]
            cam_pose = np.array([float(i) for i in scan_line[2:]])
            cam_pose = rearrange(cam_pose, "(h w) -> h w", h=4, w=4)

            image_files.append(img_file_name)
            depth_files.append(depth_file_name)
            camera_poses.append(cam_pose)
            intrinsic_matrices.append(current_intrinsics)

        else:
            continue

    return image_files, depth_files, camera_poses, intrinsic_matrices


def precompute_color_and_depth(scene, scene_dir, opts):
    old_color_imgs_path = os.path.join(scene_dir, "undistorted_color_images")
    old_depth_path = os.path.join(scene_dir, "undistorted_depth_images")

    if not os.path.exists(os.path.join(scene_dir, opts.color_dir_name)):
        os.makedirs(os.path.join(scene_dir, opts.color_dir_name))
    if not os.path.exists(os.path.join(scene_dir, opts.depth_dir_name)):
        os.makedirs(os.path.join(scene_dir, opts.depth_dir_name))

    # Load raycaster for scene
    raycaster = load_raycaster(scene_dir)

    # load camera parameters
    camera_parameter_file_path = os.path.join(
        scene_dir, "undistorted_camera_parameters", "{}.conf".format(scene)
    )
    image_files, depth_files, camera_poses, intrinsic_matrices = read_cam_parameters(
        camera_parameter_file_path
    )

    n_times_meshes_were_used = 0
    n_times_inpainting_was_used = 0
    n_images_skipped = 0

    for i in tqdm(
        range(len(camera_poses)),
        desc="Precomputing images and depths for scene {}".format(scene),
    ):
        camera_pose = camera_poses[i]
        intrinsic_matrix = intrinsic_matrices[i]
        image_file_name = image_files[i]
        depth_file_name = depth_files[i]

        # Load ground-truth color image
        img = cv2.imread(os.path.join(old_color_imgs_path, image_file_name))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (opts.img_width, opts.img_height))

        # Load ground-truth depth image
        depths = cv2.imread(
            os.path.join(old_depth_path, depth_file_name), cv2.IMREAD_ANYDEPTH
        ).astype(np.float32)
        depths = depths / 4000.0  # scale to metres; 0.25mm per unit
        depths = cv2.resize(
            depths,
            (opts.img_width, opts.img_height),
            interpolation=cv2.INTER_NEAREST_EXACT,
        )

        # fill in bad or missing depth values
        bad_depths = bad_depth_values(depths)
        n_bad_gt_depths = np.sum(bad_depths)

        if n_bad_gt_depths > 0:  # fill in bad depth values using 3D mesh
            depths_from_mesh = calc_depths_from_raycaster(
                raycaster, camera_pose, intrinsic_matrix, opts
            )
            depths[bad_depths] = depths_from_mesh[bad_depths]

            bad_depths = bad_depth_values(depths)
            n_bad_depths_v1 = np.sum(bad_depths)
            print(
                "Reduced number of bad depths from {} to {} using the 3D Mesh".format(
                    n_bad_gt_depths, n_bad_depths_v1
                )
            )
            n_times_meshes_were_used += 1
            if n_bad_depths_v1 >= 0.2 * opts.img_width * opts.img_height:
                n_images_skipped += 1
                # too many bad depths to fill in with inpainting
                print(
                    "Skipping image {} in scene {} due to too many bad depths".format(
                        image_file_name, scene
                    )
                )
                continue

            elif n_bad_depths_v1 > 0:  # fill in bad depth values using inpainting
                n_times_inpainting_was_used += 1

                depths = inpaint_depth(depths, img)

                bad_depths = bad_depth_values(depths)
                n_bad_depths_final = np.sum(bad_depths)
                print(
                    "Reduced number of bad depths from {} to {} using inpainting".format(
                        n_bad_depths_v1, n_bad_depths_final
                    )
                )

        assert np.sum(bad_depth_values(depths)) == 0

        np.save(os.path.join(scene_dir, opts.color_dir_name, image_file_name[:-4]), img)
        np.save(
            os.path.join(scene_dir, opts.depth_dir_name, depth_file_name[:-4]), depths
        )

    print("Finished precomputing scene {}".format(scene))
    print(
        "Meshes were used to fill in bad depths {} times".format(
            n_times_meshes_were_used
        )
    )
    print(
        "Inpainting was used to fill in bad depths {} times".format(
            n_times_inpainting_was_used
        )
    )
    print("{} images were skipped due to too many bad depths".format(n_images_skipped))


def main(opts):
    for scene in ALL_SCENES:
        print(f"Processing scene {scene}")
        scene_dir = os.path.join(opts.raw_dataset_dir, scene)

        # unzip the files
        if "region_segmentations.zip" in os.listdir(scene_dir):
            with zipfile.ZipFile(
                os.path.join(scene_dir, "region_segmentations.zip"), "r"
            ) as zip_ref:
                zip_ref.extractall(opts.raw_dataset_dir)

            os.remove(os.path.join(scene_dir, "region_segmentations.zip"))

        if "undistorted_camera_parameters.zip" in os.listdir(scene_dir):
            with zipfile.ZipFile(
                os.path.join(scene_dir, "undistorted_camera_parameters.zip"), "r"
            ) as zip_ref:
                zip_ref.extractall(opts.raw_dataset_dir)

            os.remove(os.path.join(scene_dir, "undistorted_camera_parameters.zip"))

        if "undistorted_color_images.zip" in os.listdir(scene_dir):
            with zipfile.ZipFile(
                os.path.join(scene_dir, "undistorted_color_images.zip"), "r"
            ) as zip_ref:
                zip_ref.extractall(opts.raw_dataset_dir)

            os.remove(os.path.join(scene_dir, "undistorted_color_images.zip"))

        if "undistorted_depth_images.zip" in os.listdir(scene_dir):
            with zipfile.ZipFile(
                os.path.join(scene_dir, "undistorted_depth_images.zip"), "r"
            ) as zip_ref:
                zip_ref.extractall(opts.raw_dataset_dir)

            os.remove(os.path.join(scene_dir, "undistorted_depth_images.zip"))

        precompute_color_and_depth(scene, scene_dir, opts)

    print("Done!")


if __name__ == "__main__":
    parser = configargparse.ArgumentParser()
    parser.add_argument(
        "--raw_dataset_dir",
        type=str,
        default="data/matterport3d",
    )
    parser.add_argument("--img_width", type=int, default=320)
    parser.add_argument("--img_height", type=int, default=256)
    parser.add_argument("--color_dir_name", type=str, default="rgb")
    parser.add_argument("--depth_dir_name", type=str, default="depth")
    opts = parser.parse_args()
    main(opts)

scripts/data_preparation_scripts/matterport3d_precompute.py
dakloepfer/locus
import os

import numpy as np
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from torchvision.transforms.functional import to_tensor
from einops import rearrange

from src.utils.misc import invert_se3


class MatterportDataset(torch.utils.data.Dataset):
    """Dataset for a single Matterport3D scene."""

    def __init__(
        self,
        data_root,
        scene_name,
        intrinsics_path,
        mode="train",
        pose_dir=None,
        augment_fn=None,
        **kwargs,
    ):
        super().__init__()

        self.data_root = data_root
        self.scene_name = scene_name
        self.intrinsics_path = intrinsics_path
        self.augment_fn = augment_fn if mode == "train" else None

        self.horizontal_only = kwargs.get("horizontal_only", False)

        if kwargs.get("normalize", "imagenet") == "imagenet":
            self.normalize = transforms.Normalize(
                (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
            )  # for backbones pretrained on ImageNet
        else:
            raise NotImplementedError(
                f"Image normalization {kwargs['normalize']} not implemented."
            )
        self.img_height = kwargs.get("img_height", 1024)
        self.img_width = kwargs.get("img_width", 1024)
        self.downsampling_factor = self.img_height / 1024

        self.img_dir = None
        self.depth_dir = None

        self.data_dict = self._create_data_dict()

    def filter_img_files(self, img_dir):
        """Filter out images that aren't used for various reasons."""
        all_img_files = os.listdir(img_dir)
        used_img_files = []
        used_file_mask = torch.ones(len(all_img_files), dtype=torch.bool)

        for i, file_name in enumerate(all_img_files):
            if self.horizontal_only:
                # remove all the files that aren't looking horizontally
                if not file_name[-7] == "1":
                    used_file_mask[i] = False
                    continue

            used_img_files.append(file_name)

        return used_img_files, used_file_mask

    def _create_data_dict(self):
        data_dict = {}
        scene_dir = os.path.join(self.data_root, self.scene_name)
        camera_parameter_file_path = os.path.join(
            scene_dir,
            "undistorted_camera_parameters",
            "{}.conf".format(self.scene_name),
        )

        with open(camera_parameter_file_path, "r") as f:
            params_file_lines = f.readlines()

        current_intrinsics = None
        if self.img_height >= 448:
            self.depth_dir = os.path.join(scene_dir, "depth_highres")
            self.img_dir = os.path.join(scene_dir, "rgb_highres")
        else:
            self.depth_dir = os.path.join(scene_dir, "depth")
            self.img_dir = os.path.join(scene_dir, "rgb")

        used_img_files, _ = self.filter_img_files(self.img_dir)

        sample_idx = 0
        for line in params_file_lines:
            if line.startswith("intrinsics"):
                intrinsics_line = [
                    i
                    for i in line.strip().split(" ")
                    if not (i.isspace() or len(i) == 0)
                ]

                current_intrinsics = torch.eye(4)
                current_intrinsics[0, 0] = (
                    float(intrinsics_line[1]) * self.downsampling_factor
                )
                current_intrinsics[1, 1] = (
                    float(intrinsics_line[5]) * self.downsampling_factor
                )
                current_intrinsics[0, 2] = (
                    float(intrinsics_line[3]) * self.downsampling_factor
                )
                current_intrinsics[1, 2] = (
                    float(intrinsics_line[6]) * self.downsampling_factor
                )

            elif line.startswith("scan"):
                scan_line = line.split(" ")[1:]
                depth_file_name = scan_line[0].replace(".png", ".npy")
                img_file_name = scan_line[1].replace(".jpg", ".npy")

                # I filter out some images that had too few good depth values to inpaint the depth map
                if img_file_name not in used_img_files:
                    continue

                cam_to_world_pose = torch.tensor(
                    [float(t) for t in scan_line[2:]], dtype=torch.float
                )
                cam_to_world_pose = rearrange(
                    cam_to_world_pose, "(h w) -> h w", h=4, w=4
                )
                world_to_cam_pose = invert_se3(cam_to_world_pose)

                data_dict[sample_idx] = {
                    "img_path": os.path.join(self.img_dir, img_file_name),
                    "depth_path": os.path.join(self.depth_dir, depth_file_name),
                    "world_to_cam_pose": world_to_cam_pose,
                    "intrinsics": current_intrinsics,
                }
                sample_idx += 1

            else:
                continue

        return data_dict

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

    def __getitem__(self, index):
        img_path = self.data_dict[index]["img_path"]
        depth_path = self.data_dict[index]["depth_path"]
        world_to_cam_pose = self.data_dict[index]["world_to_cam_pose"]
        intrinsics = self.data_dict[index]["intrinsics"]

        img = np.load(img_path)
        img = F.interpolate(
            to_tensor(img).unsqueeze(0),
            size=(self.img_height, self.img_width),
        ).squeeze(0)
        img = self.normalize(img)

        if self.augment_fn is not None:
            img = self.augment_fn(img)

        depth = torch.from_numpy(np.load(depth_path)).to(dtype=torch.float)
        depth = F.interpolate(
            depth.unsqueeze(0).unsqueeze(0), size=img.shape[-2:]
        ).squeeze()

        sample = {
            "img": img,
            "gt_depth": depth,
            "intrinsics": intrinsics,
            "world_to_cam_pose": world_to_cam_pose,
        }

        return sample

src/data/matterport_dataset.py
dakloepfer/locus
import os

import numpy as np

import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from torchvision.transforms.functional import to_tensor

# number of objects, with all "stuff" classes grouped into one object per class
N_OBJECTS_PER_SCENE = {
    "17DRP5sb8fy": 120,
    "1LXtFkjw3qL": 279,
    "1pXnuDYAj8r": 588,
    "29hnd4uzFmX": 119,
    "2azQ1b91cZZ": 522,
    "2n8kARJN3HM": 512,
    "2t7WUuJeko7": 121,
    "5LpN3gDmAk7": 273,
    "5q7pvUzZiYa": 237,
    "5ZKStnWn8Zo": 598,
    "759xd9YjKW5": 432,
    "7y3sRwLe3Va": 418,
    "8194nk5LbLH": 109,
    "82sE5b5pLXE": 297,
    "8WUmhLawc2A": 356,
    "aayBHfsNo7d": 298,
    "ac26ZMwG7aT": 643,
    "ARNzJeq3xxb": 265,
    "B6ByNegPMKs": 1194,
    "b8cTxDM8gDG": 436,
    "cV4RVeZvu5T": 209,
    "D7G3Y4RVNrH": 82,
    "D7N2EKCX4Sj": 736,
    "dhjEzFoUFzH": 60,
    "E9uDoFAP3SH": 622,
    "e9zR4mvMWw7": 303,
    "EDJbREhghzL": 313,
    "EU6Fwq7SyZv": 247,
    "fzynW3qQPVF": 221,
    "GdvgFV5R1Z5": 66,
    "gTV8FGcVJC9": 710,
    "gxdoqLR6rwA": 37,
    "gYvKGZ5eRqb": 49,
    "gZ6f7yhEvPG": 24,
    "HxpKQynjfin": 79,
    "i5noydFURQK": 209,
    "JeFG25nYj2p": 245,
    "JF19kD82Mey": 134,
    "jh4fc5c5qoQ": 155,
    "JmbYfDe2QKZ": 329,
    "jtcxE69GiFV": 578,
    "kEZ7cmS4wCh": 344,
    "mJXqzFtmKg4": 381,
    "oLBMNvg9in8": 390,
    "p5wJjkQkbXX": 406,
    "pa4otMbVnkk": 773,
    "pLe4wQe7qrG": 22,
    "Pm6F8kyY3z2": 61,
    "pRbA3pwrgk9": 260,
    "PuKPg4mmafe": 271,
    "PX4nDJXEHrG": 604,
    "q9vSo1VnCiC": 262,
    "qoiz87JEwZ2": 332,
    "QUCTc6BB5sX": 527,
    "r1Q1Z4BcV1o": 345,
    "r47D5H71a5s": 139,
    "rPc6DW4iMge": 306,
    "RPmz2sHmrrY": 150,
    "rqfALeAoiTq": 188,
    "s8pcmisQ38h": 143,
    "S9hNv5qa7GM": 288,
    "sKLMLpTHeUy": 241,
    "SN83YJsR3w2": 666,
    "sT4fr6TAbpF": 206,
    "TbHJrupSAjP": 380,
    "ULsKaCPVFJR": 265,
    "uNb9QFRL6hY": 492,
    "ur6pFq6Qu1A": 298,
    "UwV83HsGsw3": 236,
    "Uxmj2M2itWa": 167,
    "V2XKFyX4ASd": 254,
    "VFuaQ6m2Qom": 440,
    "VLzqgDo317F": 387,
    "Vt2qJdWjCF2": 206,
    "VVfe2KiqLaN": 151,
    "Vvot9Ly1tCj": 362,
    "vyrNrziPKCB": 846,
    "VzqfbhrpDEA": 810,
    "wc2JMjhGNzB": 582,
    "WYY7iVyf5p8": 226,
    "X7HyMhZNoso": 246,
    "x8F5xyUWy9e": 124,
    "XcA2TqTSSAj": 235,
    "YFuZgdQ5vWj": 257,
    "YmJkqBEsHnH": 67,
    "yqstnuAEVhm": 527,
    "YVUC4YcDtcY": 156,
    "Z6MFQCViBuw": 351,
    "ZMojNkEp431": 219,
    "zsNo4HB9uLZ": 279,
}

STUFF_CLASSES = {
    "void": 0,
    "wall": 1,
    "floor": 2,
    "ceiling": 3,
    "column": 4,
    "beam": 5,
    "objects": 6,
    "misc": 7,
    "unlabeled": 8,
}

N_STUFF_CLASS_OBJECTS = {
    "17DRP5sb8fy": [15, 10, 12, 10, 0, 0, 14, 4, 10],
    "1LXtFkjw3qL": [10, 104, 31, 33, 6, 6, 54, 4, 31],
    "1pXnuDYAj8r": [18, 125, 36, 34, 10, 3, 108, 20, 29],
    "29hnd4uzFmX": [3, 54, 15, 8, 0, 7, 34, 10, 14],
    "2azQ1b91cZZ": [18, 126, 35, 37, 2, 4, 77, 22, 28],
    "2n8kARJN3HM": [15, 56, 34, 32, 0, 1, 72, 40, 28],
    "2t7WUuJeko7": [2, 6, 6, 6, 0, 0, 36, 5, 6],
    "5LpN3gDmAk7": [19, 62, 33, 26, 4, 0, 30, 2, 28],
    "5q7pvUzZiYa": [16, 47, 22, 19, 1, 0, 51, 7, 19],
    "5ZKStnWn8Zo": [28, 142, 32, 33, 5, 0, 142, 21, 28],
    "759xd9YjKW5": [2, 92, 29, 28, 0, 12, 45, 14, 19],
    "7y3sRwLe3Va": [13, 156, 40, 50, 0, 6, 61, 31, 35],
    "8194nk5LbLH": [1, 35, 9, 18, 3, 2, 17, 6, 8],
    "82sE5b5pLXE": [3, 86, 20, 21, 7, 0, 91, 9, 16],
    "8WUmhLawc2A": [6, 94, 27, 27, 0, 4, 86, 9, 18],
    "aayBHfsNo7d": [12, 86, 28, 21, 1, 0, 55, 14, 21],
    "ac26ZMwG7aT": [10, 193, 40, 37, 0, 4, 295, 47, 31],
    "ARNzJeq3xxb": [23, 68, 23, 18, 5, 1, 69, 19, 19],
    "B6ByNegPMKs": [9, 341, 83, 241, 65, 1, 180, 77, 81],
    "b8cTxDM8gDG": [13, 116, 38, 35, 1, 1, 69, 7, 24],
    "cV4RVeZvu5T": [15, 92, 16, 18, 0, 0, 39, 21, 11],
    "D7G3Y4RVNrH": [4, 32, 6, 6, 0, 0, 5, 9, 5],
    "D7N2EKCX4Sj": [37, 232, 43, 63, 26, 12, 328, 50, 47],
    "dhjEzFoUFzH": [0, 24, 9, 3, 0, 0, 2, 4, 4],
    "E9uDoFAP3SH": [20, 226, 46, 40, 3, 0, 56, 17, 40],
    "e9zR4mvMWw7": [25, 56, 24, 26, 1, 0, 29, 4, 24],
    "EDJbREhghzL": [15, 78, 15, 15, 1, 0, 73, 22, 15],
    "EU6Fwq7SyZv": [7, 50, 21, 11, 7, 8, 65, 14, 17],
    "fzynW3qQPVF": [9, 129, 33, 31, 0, 0, 25, 7, 31],
    "GdvgFV5R1Z5": [8, 9, 6, 6, 0, 0, 4, 3, 6],
    "gTV8FGcVJC9": [18, 191, 52, 65, 0, 2, 111, 34, 51],
    "gxdoqLR6rwA": [0, 29, 13, 3, 24, 0, 12, 3, 4],
    "gYvKGZ5eRqb": [0, 9, 2, 0, 14, 0, 21, 0, 1],
    "gZ6f7yhEvPG": [0, 5, 2, 1, 0, 0, 6, 0, 1],
    "HxpKQynjfin": [11, 16, 10, 11, 0, 0, 10, 1, 8],
    "i5noydFURQK": [12, 33, 14, 15, 0, 0, 23, 4, 14],
    "JeFG25nYj2p": [21, 68, 26, 23, 0, 0, 41, 4, 25],
    "JF19kD82Mey": [10, 61, 17, 13, 0, 0, 32, 0, 15],
    "jh4fc5c5qoQ": [4, 28, 20, 16, 1, 0, 37, 1, 16],
    "JmbYfDe2QKZ": [10, 60, 23, 26, 3, 0, 46, 8, 22],
    "jtcxE69GiFV": [26, 122, 37, 53, 0, 1, 54, 29, 44],
    "kEZ7cmS4wCh": [48, 94, 37, 35, 2, 1, 42, 14, 41],
    "mJXqzFtmKg4": [32, 61, 23, 30, 11, 0, 63, 25, 25],
    "oLBMNvg9in8": [24, 88, 48, 45, 4, 1, 91, 57, 31],
    "p5wJjkQkbXX": [41, 84, 39, 50, 1, 0, 46, 31, 38],
    "pa4otMbVnkk": [52, 181, 46, 66, 7, 2, 112, 35, 51],
    "pLe4wQe7qrG": [0, 6, 2, 3, 6, 0, 2, 0, 2],
    "Pm6F8kyY3z2": [0, 9, 5, 6, 4, 0, 3, 1, 4],
    "pRbA3pwrgk9": [2, 88, 31, 42, 0, 0, 50, 13, 28],
    "PuKPg4mmafe": [0, 27, 7, 8, 4, 0, 11, 0, 7],
    "PX4nDJXEHrG": [4, 224, 54, 57, 3, 0, 74, 8, 50],
    "q9vSo1VnCiC": [0, 105, 25, 25, 0, 2, 29, 6, 21],
    "qoiz87JEwZ2": [2, 66, 23, 21, 16, 0, 53, 10, 21],
    "QUCTc6BB5sX": [1, 94, 40, 37, 0, 0, 45, 18, 36],
    "r1Q1Z4BcV1o": [3, 62, 12, 16, 11, 5, 11, 8, 14],
    "r47D5H71a5s": [1, 62, 22, 21, 0, 0, 28, 4, 21],
    "rPc6DW4iMge": [8, 73, 33, 33, 4, 0, 23, 8, 27],
    "RPmz2sHmrrY": [5, 26, 8, 9, 0, 0, 48, 6, 8],
    "rqfALeAoiTq": [3, 98, 31, 30, 3, 0, 23, 17, 26],
    "s8pcmisQ38h": [2, 53, 18, 19, 0, 0, 15, 2, 17],
    "S9hNv5qa7GM": [9, 76, 21, 24, 0, 0, 16, 9, 19],
    "sKLMLpTHeUy": [6, 102, 27, 30, 0, 0, 28, 9, 27],
    "SN83YJsR3w2": [12, 226, 56, 68, 15, 2, 77, 22, 52],
    "sT4fr6TAbpF": [11, 43, 22, 26, 1, 1, 21, 10, 20],
    "TbHJrupSAjP": [28, 90, 34, 40, 0, 1, 34, 15, 28],
    "ULsKaCPVFJR": [6, 54, 26, 28, 0, 0, 55, 8, 25],
    "uNb9QFRL6hY": [13, 139, 54, 64, 8, 13, 65, 10, 50],
    "ur6pFq6Qu1A": [4, 98, 34, 34, 1, 0, 63, 15, 30],
    "UwV83HsGsw3": [21, 78, 30, 33, 0, 0, 37, 7, 27],
    "Uxmj2M2itWa": [3, 48, 27, 28, 1, 0, 18, 9, 27],
    "V2XKFyX4ASd": [9, 84, 33, 26, 8, 6, 17, 16, 29],
    "VFuaQ6m2Qom": [4, 106, 39, 35, 0, 0, 39, 12, 31],
    "VLzqgDo317F": [5, 95, 21, 21, 10, 2, 50, 34, 19],
    "Vt2qJdWjCF2": [3, 105, 29, 28, 10, 0, 93, 28, 41],
    "VVfe2KiqLaN": [7, 44, 12, 12, 19, 8, 30, 20, 12],
    "Vvot9Ly1tCj": [30, 137, 31, 46, 3, 0, 55, 29, 29],
    "vyrNrziPKCB": [111, 272, 71, 70, 28, 2, 230, 108, 84],
    "VzqfbhrpDEA": [101, 320, 59, 78, 0, 1, 247, 65, 63],
    "wc2JMjhGNzB": [62, 159, 46, 58, 6, 0, 76, 54, 42],
    "WYY7iVyf5p8": [8, 72, 31, 33, 0, 0, 10, 6, 23],
    "X7HyMhZNoso": [15, 153, 31, 27, 0, 0, 28, 24, 21],
    "x8F5xyUWy9e": [7, 62, 14, 14, 3, 0, 35, 13, 13],
    "XcA2TqTSSAj": [5, 110, 36, 37, 0, 2, 24, 4, 33],
    "YFuZgdQ5vWj": [15, 57, 20, 21, 2, 0, 32, 7, 16],
    "YmJkqBEsHnH": [1, 8, 2, 3, 12, 3, 9, 6, 6],
    "yqstnuAEVhm": [17, 107, 36, 42, 1, 0, 208, 39, 53],
    "YVUC4YcDtcY": [9, 28, 10, 17, 0, 0, 23, 46, 27],
    "Z6MFQCViBuw": [8, 32, 21, 20, 0, 0, 79, 26, 34],
    "ZMojNkEp431": [13, 44, 17, 28, 0, 0, 63, 26, 36],
    "zsNo4HB9uLZ": [24, 29, 22, 21, 0, 1, 83, 6, 22],
}


class MatterportSegmentationDataset(torch.utils.data.Dataset):
    """Dataset for a single Matterport3D scene, returning the image and the segmentation mask."""

    def __init__(self, data_root, scene_name, augment_fn=None, **kwargs):
        super().__init__()
        self.data_root = data_root
        self.scene_name = scene_name

        self.augment_fn = augment_fn

        self.horizontal_only = kwargs.get("horizontal_only", False)

        if kwargs.get("normalize", "imagenet") == "imagenet":
            self.normalize = transforms.Normalize(
                (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
            )  # for backbones pretrained on ImageNet
        else:
            raise NotImplementedError(
                f"Image normalization {kwargs['normalize']} not implemented."
            )

        self.img_height = kwargs.get("img_height", 1024)
        self.img_width = kwargs.get("img_width", 1024)
        self.downsampling_factor = self.img_height / 1024

        self.img_dir = None
        self.depth_dir = None

        self.data_dict = self._create_data_dict()

    def filter_img_files(self, img_dir):
        """Filter out images that aren't used for various reasons."""
        all_img_files = os.listdir(img_dir)
        used_img_files = []
        used_file_mask = torch.ones(len(all_img_files), dtype=torch.bool)

        for i, file_name in enumerate(all_img_files):
            if self.horizontal_only:
                # remove all the files that aren't looking horizontally
                if not file_name[-7] == "1":
                    used_file_mask[i] = False
                    continue

            used_img_files.append(file_name)

        return used_img_files, used_file_mask

    def _create_data_dict(self):
        data_dict = {}
        scene_dir = os.path.join(self.data_root, self.scene_name)

        if self.img_height >= 448:
            img_folder = os.path.join(scene_dir, "rgb_highres")
        else:
            img_folder = os.path.join(scene_dir, "rgb")
        segmentation_folder = os.path.join(
            scene_dir,
            "object_maps",
        )
        file_names, _ = self.filter_img_files(img_folder)

        sample_idx = 0
        for file_name in file_names:
            img_path = os.path.join(img_folder, file_name)
            segmentation_path = os.path.join(segmentation_folder, file_name)
            if not os.path.exists(segmentation_path):
                continue
            data_dict[sample_idx] = {
                "img_path": img_path,
                "segmentation_path": segmentation_path,
            }
            sample_idx += 1

        return data_dict

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

    def __getitem__(self, index):
        img_path = self.data_dict[index]["img_path"]
        segmentation_path = self.data_dict[index]["segmentation_path"]

        img = np.load(img_path)  # already RGB
        img = F.interpolate(
            to_tensor(img).unsqueeze(0),
            size=(self.img_height, self.img_width),
        ).squeeze(0)
        img = self.normalize(img)

        # should be a tensor that contains the segmentation categories as ints
        segmentation = torch.from_numpy(np.load(segmentation_path)).long()
        segmentation = (
            F.interpolate(
                segmentation[None, None].float(),
                size=(self.img_height, self.img_width),
                mode="nearest-exact",
            )
            .squeeze()
            .long()
        )
        if self.augment_fn is not None:
            img = self.transforms(img)

        sample = {"img": img, "segmentation": segmentation}
        return sample

src/data/matterport_segmentation_dataset.py
dakloepfer/locus
import os
import random
import numpy as np
import torch


class SingleSceneBatchSampler(torch.utils.data.Sampler):
    """Samples batches from a ConcatDataset so that each batch contains only samples from a single scene (sub-dataset)."""

    def __init__(self, datasource, batch_size, shuffle=True, use_overlaps=True):
        super().__init__(datasource)

        self.datasource = datasource
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.use_overlaps = use_overlaps

        self.n_scenes = len(datasource.datasets)
        self.scene_lengths = [len(dataset) for dataset in datasource.datasets]
        self.batches_per_scene = [length // batch_size for length in self.scene_lengths]

        if use_overlaps:
            self.overlap_matrices = []

            for scene in datasource.datasets:
                scene_name = scene.scene_name
                overlap_matrix = torch.from_numpy(
                    np.load(
                        os.path.join(
                            scene.data_root,
                            scene_name,
                            "img_overlaps.npy",
                        )
                    )
                )

                # filter the overlaps so they only contain images that are used
                _, used_img_mask = scene.filter_img_files(scene.img_dir)
                overlap_matrix = overlap_matrix[used_img_mask, :][:, used_img_mask]

                self.overlap_matrices.append(overlap_matrix)

    def _make_epoch_with_overlaps(self):
        """Create a list of batches, with each batch containing just images from a single scene.
        First randomly sample a scene, then randomly sample an seed image from that scene, then sample the rest of the batch weighted by the relative overlaps with the seed image (i.e. images that overlap more with the seed image are more likely to be sampled).
        """

        batches = []
        batches_sampled_per_scene = [0 for _ in range(self.n_scenes)]
        remaining_scene_idxs = list(range(self.n_scenes))
        all_unused_images = [list(range(length)) for length in self.scene_lengths]

        for _ in range(len(self)):
            # sample a scene for which there are batches left to sample
            if self.shuffle:
                scene_idx = random.choice(remaining_scene_idxs)
            else:
                scene_idx = remaining_scene_idxs[0]
            batches_sampled_per_scene[scene_idx] += 1
            if (
                batches_sampled_per_scene[scene_idx]
                >= self.batches_per_scene[scene_idx]
            ):
                remaining_scene_idxs.remove(scene_idx)

            scene_length = self.scene_lengths[scene_idx]
            overlap_matrix = self.overlap_matrices[scene_idx]
            unused_images = all_unused_images[scene_idx]

            scene_idx_offset = (
                self.datasource.cumulative_sizes[scene_idx] - scene_length
            )

            # sample a seed image
            if self.shuffle:
                seed_idx = random.choice(unused_images)
            else:
                seed_idx = unused_images[0]
            unused_images.remove(seed_idx)

            # sample the rest of the batch
            batch = [seed_idx + scene_idx_offset]
            if self.shuffle:
                # sample the rest of the batch from the overlap matrix
                indices = torch.multinomial(
                    overlap_matrix[seed_idx],
                    self.batch_size - 1,
                    replacement=False,
                )
                # NOTE I am actually not removing the images that are sampled, so some are used multiple times and some not at all (for a certain epoch). This is because otherwise the overlap tends to decrease throughout the epoch as the number of unused images decreases, which tends to artificially increase the AP.
            else:
                # take the images with the largest overlaps with the seed image
                indices = torch.topk(
                    overlap_matrix[seed_idx],
                    self.batch_size - 1,
                )[1]

            batch.extend([idx.item() for idx in indices + scene_idx_offset])
            batches.append(batch)

        return batches

    def _make_epoch_without_overlaps(self):
        """create list of batches sampled uniformly, with each batch coming from only a single scene"""
        batches = []

        for scene_idx in range(self.n_scenes):
            scene_length = self.scene_lengths[scene_idx]
            n_batches = self.batches_per_scene[scene_idx]

            indices = list(range(scene_length))
            indices = [
                idx + self.datasource.cumulative_sizes[scene_idx] - scene_length
                for idx in indices
            ]
            if self.shuffle:
                random.shuffle(indices)

            for batch_idx in range(n_batches):
                batch_start = batch_idx * self.batch_size
                batch_end = (batch_idx + 1) * self.batch_size
                batches.append(indices[batch_start:batch_end])

        if self.shuffle:
            random.shuffle(batches)

        return batches

    def __len__(self):
        return sum(self.batches_per_scene)

    def __iter__(self):
        # create batches first, then yield from the tensor of indices
        if self.use_overlaps:
            batches = self._make_epoch_with_overlaps()
        else:
            batches = self._make_epoch_without_overlaps()

        for batch in batches:
            yield batch

src/data/single_scene_batchsampler.py
dakloepfer/locus
from yacs.config import CfgNode as CN

_CN = CN()

############### MODEL ###############

_CN.MODEL = CN()

_CN.MODEL.TYPE = "dino"
_CN.MODEL.N_PRETRAINED_BLOCKS = 11  # number of pretrained DINO blocks
_CN.MODEL.VIT_BACKBONE = "vitb8"

_CN.MODEL.FEATURE_DIM = 64
_CN.MODEL.POS_LANDMARK_RADIUS = 0.2
_CN.MODEL.NEG_LANDMARK_RADIUS = 10.0
_CN.MODEL.LANDMARK_EMBEDDING_METHOD = "sampled_patch"


############### DATA ###############
_CN.DATASET = CN()
_CN.DATASET.TRAIN_DATA_SOURCE = "Matterport"
_CN.DATASET.VAL_DATA_SOURCE = "Matterport"
_CN.DATASET.TEST_DATA_SOURCE = "Matterport"

_CN.DATASET.TRAIN_DATA_ROOT = None
_CN.DATASET.TRAIN_POSE_ROOT = None
_CN.DATASET.TRAIN_INTRINSICS_PATH = None
_CN.DATASET.TRAIN_SCENE_LIST = None  # file path for list of scene names

_CN.DATASET.VAL_DATA_ROOT = None
_CN.DATASET.VAL_POSE_ROOT = None
_CN.DATASET.VAL_INTRINSICS_PATH = None
_CN.DATASET.VAL_SCENE_LIST = None  # file path for list of scene names

_CN.DATASET.TEST_DATA_ROOT = None
_CN.DATASET.TEST_POSE_ROOT = None
_CN.DATASET.TEST_INTRINSICS_PATH = None
_CN.DATASET.TEST_SCENE_LIST = None  # file path for list of scene names

_CN.DATASET.IMG_HEIGHT = 256
_CN.DATASET.IMG_WIDTH = 320
_CN.DATASET.AUGMENTATION_TYPE = None
_CN.DATASET.TRAIN_SHUFFLE = True
_CN.DATASET.USE_OVERLAPS = True

_CN.DATASET.MATTERPORT_HORIZONTAL_IMGS_ONLY = True
_CN.DATASET.MATTERPORT_NORMALIZE = "imagenet"


############### TRAINING ###############
_CN.TRAINER = CN()
_CN.TRAINER.SEED = 42
_CN.TRAINER.GRADIENT_CLIPPING = None
_CN.TRAINER.ACCUMULATE_GRAD_BATCHES = 1

_CN.TRAIN = CN()
_CN.TRAIN.N_LANDMARKS = 64
_CN.TRAIN.FRAC_PATCH_SUBSAMPLE = 0.2
_CN.TRAIN.KEEP_ALL_POSITIVE_PATCHES = True


############### OPTIMIZER ###############
_CN.OPTIMIZER = CN()
_CN.OPTIMIZER.TYPE = "Adam"
_CN.OPTIMIZER.LR = 1e-3
_CN.OPTIMIZER.WEIGHT_DECAY = 0.0


############## LOSS ###############
_CN.LOSS = CN()
_CN.LOSS.SIGMOID_TEMPERATURE = 0.01


def get_cfg_defaults():
    """Get a yacs CfgNode object with default values for my_project."""
    # Return a clone so that the defaults will not be altered
    # This is for the "local variable" use pattern
    return _CN.clone()

src/default_config/default.py
dakloepfer/locus
from tqdm import tqdm

import torch
import lightning.pytorch as pl
from loguru import logger
from torch.utils.data import ConcatDataset, DataLoader

from src.utils.augment import build_augmentor
from src.data.matterport_dataset import MatterportDataset
from src.data.single_scene_batchsampler import SingleSceneBatchSampler


class MultiSceneDataModule(pl.LightningDataModule):
    def __init__(self, args, config):
        super().__init__()

        # Collect all necessary config parameters
        self.train_data_source = config.DATASET.TRAIN_DATA_SOURCE
        self.val_data_source = config.DATASET.VAL_DATA_SOURCE
        self.test_data_source = config.DATASET.TEST_DATA_SOURCE

        self.train_data_root = config.DATASET.TRAIN_DATA_ROOT
        self.train_pose_root = config.DATASET.TRAIN_POSE_ROOT
        self.train_intrinsics_path = config.DATASET.TRAIN_INTRINSICS_PATH
        self.train_scene_list_path = config.DATASET.TRAIN_SCENE_LIST

        self.val_data_root = config.DATASET.VAL_DATA_ROOT
        self.val_pose_root = config.DATASET.VAL_POSE_ROOT
        self.val_intrinsics_path = config.DATASET.VAL_INTRINSICS_PATH
        self.val_scene_list_path = config.DATASET.VAL_SCENE_LIST

        self.test_data_root = config.DATASET.TEST_DATA_ROOT
        self.test_pose_root = config.DATASET.TEST_POSE_ROOT
        self.test_intrinsics_path = config.DATASET.TEST_INTRINSICS_PATH
        self.test_scene_list_path = config.DATASET.TEST_SCENE_LIST

        self.augment_type = config.DATASET.AUGMENTATION_TYPE
        self.shuffle = config.DATASET.TRAIN_SHUFFLE

        self.matterport_config = {
            "horizontal_only": config.DATASET.MATTERPORT_HORIZONTAL_IMGS_ONLY,
            "normalize": config.DATASET.MATTERPORT_NORMALIZE,
            "img_height": config.DATASET.IMG_HEIGHT,
            "img_width": config.DATASET.IMG_WIDTH,
        }
        self.batch_sampler_config = {"use_overlaps": config.DATASET.USE_OVERLAPS}

        # Collect all necessary args parameters
        self.batch_size = args.batch_size
        self.train_loader_params = {
            "num_workers": args.num_workers,
            "pin_memory": getattr(args, "pin_memory", True),
        }
        self.val_loader_params = {
            "num_workers": args.num_workers,
            "pin_memory": getattr(args, "pin_memory", True),
        }
        self.test_loader_params = {
            "num_workers": args.num_workers,
            "pin_memory": True,
        }
        self.disable_tqdm = args.tqdm_refresh_rate != 1

    def setup(self, stage: str = None):
        """
        Setup train / val / test dataset. This method will be called by PL automatically.
        Args:
            stage (str): 'fit' in training phase, and 'test' in testing phase.
        """

        assert stage in ["fit", "test"], "stage must be either fit or test"

        if stage == "fit":
            self.augment_fn = build_augmentor(
                self.augment_type, data_source=self.train_data_source
            )

            self.train_dataset = self._setup_dataset(
                self.train_data_source,
                self.train_data_root,
                self.train_scene_list_path,
                self.train_intrinsics_path,
                mode="train",
                pose_dir=self.train_pose_root,
                augment_fn=self.augment_fn,
            )

            self.val_dataset = self._setup_dataset(
                self.val_data_source,
                self.val_data_root,
                self.val_scene_list_path,
                self.val_intrinsics_path,
                mode="val",
                pose_dir=self.val_pose_root,
                augment_fn=None,
            )
            logger.info("Train & Val Dataset loaded!")

        elif stage == "test":
            self.test_dataset = self._setup_dataset(
                self.test_data_source,
                self.test_data_root,
                self.test_scene_list_path,
                self.test_intrinsics_path,
                mode="test",
                pose_dir=self.test_pose_root,
                augment_fn=None,
            )
            logger.info("Test Dataset loaded!")

        else:
            raise NotImplementedError(f"Stage must be either fit or test, was {stage}")

    def _setup_dataset(
        self,
        data_source,
        data_root,
        scene_list_path,
        intrinsics_path,
        mode="train",
        pose_dir=None,
        augment_fn=None,
    ):
        """To make it a bit easier to set up different datasets"""

        with open(scene_list_path, "r") as f:
            scene_names = [name.split()[0] for name in f.readlines()]

        logger.info(f"{len(scene_names)} scene(s) assigned.")

        if self.disable_tqdm and int(self.rank) == 0:
            print(f"Loading {mode} datasets")

        datasets = []
        for scene_name in tqdm(
            scene_names,
            desc=f"Loading {mode} datasets",
            disable=self.disable_tqdm,
        ):
            if data_source == "Matterport":
                datasets.append(
                    MatterportDataset(
                        data_root,
                        scene_name,
                        intrinsics_path,
                        mode=mode,
                        pose_dir=pose_dir,
                        augment_fn=augment_fn,
                        **self.matterport_config,
                    )
                )
            else:
                raise NotImplementedError(f"Data source {data_source} not implemented")

        if self.disable_tqdm and int(self.rank) == 0:
            print(f"Finished loading {mode} datasets.")

        return ConcatDataset(datasets)

    def train_dataloader(self):
        batch_sampler = SingleSceneBatchSampler(
            self.train_dataset,
            batch_size=self.batch_size,
            shuffle=self.shuffle,
            **self.batch_sampler_config,
        )
        dataloader = DataLoader(
            self.train_dataset, batch_sampler=batch_sampler, **self.train_loader_params
        )

        return dataloader

    def val_dataloader(self):
        batch_sampler = SingleSceneBatchSampler(
            self.val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            **self.batch_sampler_config,
        )
        dataloader = DataLoader(
            self.val_dataset, batch_sampler=batch_sampler, **self.val_loader_params
        )

        return dataloader

    def test_dataloader(self):
        batch_sampler = SingleSceneBatchSampler(
            self.test_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            **self.batch_sampler_config,
        )
        dataloader = DataLoader(
            self.test_dataset, batch_sampler=batch_sampler, **self.test_loader_params
        )
        return dataloader

src/lightning/data.py
dakloepfer/locus
from tqdm import tqdm
import torch
from torch import nn
from torch.nn import functional as F
from lightning import pytorch as pl
import torchmetrics
from src.data.matterport_segmentation_dataset import (
    N_OBJECTS_PER_SCENE,
    N_STUFF_CLASS_OBJECTS,
)


class LinearProbeModule(pl.LightningModule):
    """LightningModule for linear probe training."""

    def __init__(self, config, scene):
        super().__init__()
        self.config = config
        self.scene = scene

        self.feature_dim = config.MODEL.FEATURE_DIM
        # +1 for invalid object
        self.n_objects = N_OBJECTS_PER_SCENE[scene] + 1
        # have outputs classes for individual "stuff class objects", but don't require the linear probe to re-ID them (see forward() method)
        # (the object segmentations for eg different parts of wall often don't make sense)

        self.n_stuff_objects = N_STUFF_CLASS_OBJECTS[scene]

        # one output per object, not grouping stuff class 'objects'
        self.n_outputs = (
            self.n_objects + sum(self.n_stuff_objects) - len(self.n_stuff_objects)
        )

        self.linear_layer = nn.Conv2d(
            self.feature_dim,
            self.n_outputs,
            kernel_size=1,
            bias=True,
        )

        self.test_results = {}

    def configure_optimizers(self):
        return torch.optim.Adam(self.linear_layer.parameters(), lr=0.01)

    def forward(self, x):
        raw_logits = self.linear_layer(x)

        # the logits for the "stuff class objects" are the maximum of the logits for the individual stuff class 'objects'

        logits = torch.zeros(
            raw_logits.shape[0],
            self.n_objects,
            *raw_logits.shape[2:],
            device=self.device,
        )
        for i in range(len(self.n_stuff_objects)):
            start_idx = sum(self.n_stuff_objects[:i])
            end_idx = sum(self.n_stuff_objects[: i + 1])
            if end_idx == start_idx:
                continue
            stuff_logits, used_idxs = torch.max(raw_logits[:, start_idx:end_idx], dim=1)

            logits[:, i] = stuff_logits

        logits[:, len(self.n_stuff_objects) :] = raw_logits[
            :, sum(self.n_stuff_objects) :
        ]

        return logits

    def training_step(self, batch, batch_idx):
        features, gt_segmentation = batch
        gt_segmentation = (
            F.interpolate(
                gt_segmentation.float().unsqueeze(dim=1),
                size=features.shape[-2:],
                mode="nearest-exact",
            )
            .long()
            .squeeze(dim=1)
        )

        logits = self(features)

        loss = F.cross_entropy(logits, gt_segmentation)

        self.log("linear_probe_train_loss", loss)
        return loss

    def on_test_epoch_start(self) -> None:
        self.test_results.update(
            {
                "all_logits": [],
                "all_gt_labels": [],
                "true_positives": torch.zeros(self.n_objects, dtype=torch.long),
                "false_positives": torch.zeros(self.n_objects, dtype=torch.long),
                "false_negatives": torch.zeros(self.n_objects, dtype=torch.long),
                "true_negatives": torch.zeros(self.n_objects, dtype=torch.long),
            }
        )

    def test_step(self, batch, batch_idx):
        features, gt_segmentation_class = batch

        logits = self(features)

        # upsample logits to match gt_segmentation
        upsampled_logits = F.interpolate(
            logits, size=gt_segmentation_class.shape[-2:], mode="nearest-exact"
        )
        predicted_segmentation = torch.argmax(upsampled_logits, dim=1)
        upsampled_logits = None  # free memory
        predicted_segmentation = predicted_segmentation.unsqueeze(-1) == torch.arange(
            self.n_objects, device=predicted_segmentation.device
        ).view(1, 1, 1, self.n_objects)

        gt_segmentation = gt_segmentation_class.to(
            predicted_segmentation.device
        ).unsqueeze(dim=-1) == torch.arange(
            self.n_objects, device=predicted_segmentation.device
        ).view(
            1, 1, 1, self.n_objects
        )

        for i in range(self.n_objects):
            self.test_results["true_positives"][i] += torch.sum(
                predicted_segmentation[:, :, :, i] & gt_segmentation[:, :, :, i]
            ).cpu()
            self.test_results["false_positives"][i] += torch.sum(
                predicted_segmentation[:, :, :, i] & ~gt_segmentation[:, :, :, i]
            ).cpu()
            self.test_results["false_negatives"][i] += torch.sum(
                ~predicted_segmentation[:, :, :, i] & gt_segmentation[:, :, :, i]
            ).cpu()
            self.test_results["true_negatives"][i] += torch.sum(
                ~predicted_segmentation[:, :, :, i] & ~gt_segmentation[:, :, :, i]
            ).cpu()
        predicted_segmentation = None
        gt_segmentation = None

        # to compute AP over the entire epoch, I need to store the actual logits and grond truth labels -- do that in the downsampled version to save memory
        downsampled_gt_segmentation = (
            F.interpolate(
                gt_segmentation_class.unsqueeze(dim=1).float(),
                size=logits.shape[-2:],
                mode="nearest-exact"
                if not torch.__version__ in ["1.10.2", "1.9.0"]
                else "nearest",
            )
            .squeeze(dim=1)
            .long()
        ).cpu()

        # scale to between 0 and 1 so the AP calculation is reasonably accurate with 100 bins
        # only thing that matters for average precision is that the ordering is preserved for each class individually
        logits = logits - torch.min(logits, dim=1, keepdim=True)[0]
        logits = logits / torch.max(logits, dim=1, keepdim=True)[0]
        logits = logits.cpu()

        self.test_results["all_logits"].append(logits)
        self.test_results["all_gt_labels"].append(downsampled_gt_segmentation)

    def on_test_epoch_end(self) -> None:
        """compute metrics across entire epoch"""
        self.test_results["mIoU"] = torch.mean(
            self.test_results["true_positives"]
            / (
                self.test_results["true_positives"]
                + self.test_results["false_positives"]
                + self.test_results["false_negatives"]
                + 1e-8
            )
        ).item()
        self.test_results["Jaccard"] = (
            torch.sum(self.test_results["true_positives"])
            / (
                torch.sum(self.test_results["true_positives"])
                + torch.sum(self.test_results["false_positives"])
                + torch.sum(self.test_results["false_negatives"])
            ).item()
        )
        self.test_results["Accuracy"] = (
            torch.sum(self.test_results["true_positives"])
            + torch.sum(self.test_results["true_negatives"])
        ) / (
            torch.sum(self.test_results["true_positives"])
            + torch.sum(self.test_results["false_positives"])
            + torch.sum(self.test_results["false_negatives"])
            + torch.sum(self.test_results["true_negatives"])
        ).item()

        self.test_results["object_mIoU"] = torch.mean(
            self.test_results["true_positives"][len(self.n_stuff_objects) :]
            / (
                self.test_results["true_positives"][len(self.n_stuff_objects) :]
                + self.test_results["false_positives"][len(self.n_stuff_objects) :]
                + self.test_results["false_negatives"][len(self.n_stuff_objects) :]
                + 1e-8
            )
        ).item()
        self.test_results["object_Jaccard"] = (
            self.test_results["true_positives"][len(self.n_stuff_objects) :].sum()
            / (
                self.test_results["true_positives"][len(self.n_stuff_objects) :].sum()
                + self.test_results["false_positives"][
                    len(self.n_stuff_objects) :
                ].sum()
                + self.test_results["false_negatives"][
                    len(self.n_stuff_objects) :
                ].sum()
            ).item()
        )
        self.test_results["object_Accuracy"] = (
            self.test_results["true_positives"][len(self.n_stuff_objects) :].sum()
            + self.test_results["true_negatives"][len(self.n_stuff_objects) :].sum()
        ) / (
            self.test_results["true_positives"][len(self.n_stuff_objects) :].sum()
            + self.test_results["false_positives"][len(self.n_stuff_objects) :].sum()
            + self.test_results["false_negatives"][len(self.n_stuff_objects) :].sum()
            + self.test_results["true_negatives"][len(self.n_stuff_objects) :].sum()
        ).item()

        self.test_results["stuff_mIoU"] = torch.mean(
            self.test_results["true_positives"][: len(self.n_stuff_objects)]
            / (
                self.test_results["true_positives"][: len(self.n_stuff_objects)]
                + self.test_results["false_positives"][: len(self.n_stuff_objects)]
                + self.test_results["false_negatives"][: len(self.n_stuff_objects)]
                + 1e-8
            )
        ).item()
        self.test_results["stuff_Jaccard"] = (
            self.test_results["true_positives"][: len(self.n_stuff_objects)].sum()
            / (
                self.test_results["true_positives"][: len(self.n_stuff_objects)].sum()
                + self.test_results["false_positives"][
                    : len(self.n_stuff_objects)
                ].sum()
                + self.test_results["false_negatives"][
                    : len(self.n_stuff_objects)
                ].sum()
            ).item()
        )
        self.test_results["stuff_Accuracy"] = (
            self.test_results["true_positives"][: len(self.n_stuff_objects)].sum()
            + self.test_results["true_negatives"][: len(self.n_stuff_objects)].sum()
        ) / (
            self.test_results["true_positives"][: len(self.n_stuff_objects)].sum()
            + self.test_results["false_positives"][: len(self.n_stuff_objects)].sum()
            + self.test_results["false_negatives"][: len(self.n_stuff_objects)].sum()
            + self.test_results["true_negatives"][: len(self.n_stuff_objects)].sum()
        ).item()

        all_logits = torch.cat(self.test_results["all_logits"], dim=0)
        all_gt_labels = torch.cat(self.test_results["all_gt_labels"], dim=0)

        AP_metric = torchmetrics.classification.BinaryAveragePrecision()
        class_APs = torch.zeros(self.n_objects, device=all_logits.device)
        for i in tqdm(range(self.n_objects), "calculating mAP"):
            class_APs[i] = AP_metric(all_logits[:, i], (all_gt_labels == i))

        self.test_results["mAP"] = torch.mean(class_APs[~torch.isnan(class_APs)]).item()
        self.test_results["object_mAP"] = torch.mean(
            class_APs[len(self.n_stuff_objects) :][
                ~torch.isnan(class_APs[len(self.n_stuff_objects) :])
            ]
        ).item()
        self.test_results["stuff_mAP"] = torch.mean(
            class_APs[: len(self.n_stuff_objects)][
                ~torch.isnan(class_APs[: len(self.n_stuff_objects)])
            ]
        ).item()

    def get_results(self):
        return self.test_results

src/lightning/linear_probe.py
dakloepfer/locus
from typing import Any, Optional
from lightning.pytorch.utilities.types import STEP_OUTPUT
from loguru import logger
from math import floor

import torch
from torch.nn import functional as F
import lightning.pytorch as pl
from kornia.geometry.camera import PinholeCamera
from kornia.geometry.linalg import transform_points
from kornia.geometry.conversions import convert_points_from_homogeneous
from kornia.utils import create_meshgrid
from einops import rearrange, repeat, reduce

from src.models.dino_feature_extractor import DINO_FeatureExtractor
from src.models.baselines import DINO_Baseline, ResNet50Baseline
from src.losses.vec_smooth_ap import VecSmoothAP
from src.utils.misc import batched_2d_index_select
from src.utils.profiler import PassThroughProfiler


class PLModule(pl.LightningModule):
    def __init__(self, config, ckpt_path=None, devices=[], profiler=None, **kwargs):
        super().__init__()
        self.config = config
        self.profiler = profiler or PassThroughProfiler()

        if self.config.MODEL.TYPE == "dino":
            self.model = DINO_FeatureExtractor(self.config.MODEL)
        elif self.config.MODEL.TYPE == "dino_baseline":
            self.model = DINO_Baseline(self.config)
        elif self.config.MODEL.TYPE == "resnet50_baseline":
            self.model = ResNet50Baseline(self.config)
        else:
            raise NotImplementedError(f"Unknown matcher: {self.config.MODEL.TYPE}")

        self.feature_subsample = self.model.output_subsample

        self.loss = VecSmoothAP(self.config.LOSS)

        if ckpt_path is not None:
            state_dict = torch.load(ckpt_path, map_location="cpu")
            if "model" not in state_dict and "model_state_dict" in state_dict:
                state_dict["model"] = state_dict["model_state_dict"]
            self.model.load_state_dict(state_dict["model"])
            logger.info(f"Loaded pretrained weights from {ckpt_path}")

        if len(devices) > 1:
            self.model = self.model.to(devices[0])
            self.model = torch.nn.DataParallel(self.model, device_ids=devices)

        self.n_landmarks = config.TRAIN.N_LANDMARKS  # number of landmarks to sample
        self.frac_patch_subsample = config.TRAIN.FRAC_PATCH_SUBSAMPLE
        self.keep_all_positive_patches = config.TRAIN.KEEP_ALL_POSITIVE_PATCHES
        self.pos_radius = config.MODEL.POS_LANDMARK_RADIUS
        self.neg_radius = config.MODEL.NEG_LANDMARK_RADIUS
        self.landmark_embedding_method = config.MODEL.LANDMARK_EMBEDDING_METHOD

        # stuff for testing
        self.test_task = kwargs.get("test_task", None)

    def configure_optimizers(self):
        optim_conf = self.config.OPTIMIZER
        if optim_conf.TYPE == "Adam":
            return torch.optim.Adam(
                self.model.parameters(),
                lr=optim_conf.LR,
                weight_decay=optim_conf.WEIGHT_DECAY,
            )
        else:
            raise NotImplementedError(f"Unknown optimizer: {optim_conf.TYPE}")

    def forward(self, imgs):
        """Extract features from the images (BxCxHxW)."""
        features = self.model(imgs)
        features = F.normalize(features, p=2, dim=1)
        return features

    def training_step(self, batch, batch_idx):
        return self._trainval_step(batch, batch_idx)

    def validation_step(self, batch, batch_idx):
        return self._trainval_step(batch, batch_idx)

    def test_step(self, batch, batch_idx):
        if self.test_task == "patch_retrieval":
            return self._trainval_step(batch, batch_idx)
        else:
            raise NotImplementedError(f"Unknown test task: {self.test_task}")

    def _trainval_step(self, batch, batch_idx):
        imgs = batch["img"]
        gt_depths = batch["gt_depth"].unsqueeze(1)
        cam_matrices = batch["intrinsics"]
        world_to_cam_poses = batch["world_to_cam_pose"]
        depths = F.interpolate(
            gt_depths, scale_factor=1 / self.feature_subsample, mode="nearest-exact"
        )

        if imgs.shape[0] <= 1:
            raise ValueError("Batch size must be greater than 1")
        img_height = imgs.shape[2]
        img_width = imgs.shape[3]
        cameras = PinholeCamera(
            cam_matrices,
            world_to_cam_poses,
            torch.ones_like(cam_matrices[:, 0, 0]) * img_height,
            torch.ones_like(cam_matrices[:, 0, 0]) * img_width,
        )

        # cannot use PyTorch Lightning's DDP because the loss function requires us to accumulate all the features from all the GPUs
        with self.profiler.profile("Feature Extraction"):
            features = self(imgs)

        with self.profiler.profile("Computing landmarks"):
            landmarks = self._sample_landmarks(cameras, depths, features)
            self._project_landmarks(landmarks, cameras, gt_depths)
            self._filter_landmarks(landmarks)
            self._compute_landmark_embeddings(landmarks, features)

        with self.profiler.profile("Computing positive / negative patches"):
            pos_patches, pos_neg_patches = self._calculate_contrastive_masks(
                landmarks, cameras, depths
            )
            (
                selected_patch_features,
                pos_patches,
                pos_neg_patches,
            ) = self._subsample_patches(pos_patches, pos_neg_patches, features)

        with self.profiler.profile("Computing loss"):
            loss, log_dict = self.loss(
                landmarks, pos_patches, pos_neg_patches, selected_patch_features
            )

        self.log_dict(log_dict)
        return loss

    def _sample_landmarks(self, cameras, depths, features):
        """Sample landmarks from the scene using the network outputs and ground truths.

        Parameters
        ----------
        cameras (PinholeCamera):
            A batch of Kornia PinholeCameras with the intrinsics and exrtrinsics for the batch.

        depths (batch_size x 1 x feature_height x feature_width tensor):
            ground truth depth in meters for each pixel patch

        features (batch_size x n_channels x feature_height x feature_width tensor):
            the features extracted.

        Returns
        -------
        landmarks (dict):
            keys:
                landmark_pos (n_landmarks x 3 tensor):
                    the positions of the selected (sampled) landmarks

                landmark_embeddings (n_landmarks x n_channels tensor):
                    the features of the selected patches
        """
        batch_size = depths.shape[0]
        height = depths.shape[2]
        width = depths.shape[3]

        patch_idxs = torch.randperm(height * width, device=self.device)[
            : self.n_landmarks
        ]
        batch_idxs = repeat(
            torch.arange(batch_size, device=self.device),
            "b -> (b h w)",
            h=height,
            w=width,
        )[patch_idxs]

        landmark_cameras = PinholeCamera(
            cameras.intrinsics[batch_idxs],
            cameras.extrinsics[batch_idxs],
            cameras.height[batch_idxs],
            cameras.width[batch_idxs],
        )

        landmark_pixel_coords = (
            create_meshgrid(
                height, width, normalized_coordinates=False, device=self.device
            )
            + 0.5
        ) * self.feature_subsample
        landmark_pixel_coords = repeat(
            landmark_pixel_coords,
            "() h w c -> (b h w) c",
            b=batch_size,
            h=height,
            w=width,
            c=2,
        )[patch_idxs]

        landmark_depths = rearrange(
            depths, "b () h w -> (b h w) ()", b=batch_size, h=height, w=width
        )[patch_idxs]
        landmark_pos = landmark_cameras.unproject(
            landmark_pixel_coords.unsqueeze(1), landmark_depths.unsqueeze(1)
        ).squeeze(1)

        landmark_embeddings = rearrange(
            features, "b c h w -> (b h w) c", b=batch_size, h=height, w=width
        )[patch_idxs]

        return {
            "landmark_pos": landmark_pos,
            "landmark_embeddings": landmark_embeddings,
        }

    def _project_landmarks(self, landmarks, cameras, gt_depths):
        """Modifies the landmarks dictionary in-place, updating it with the pixel locations, the visibility mask, and the depths for each landmark-image pair.

        Parameters
        ----------
        landmarks (dictionary):):
            minimum keys:
                landmark_pos (n_landmarks x 3 tensor):
                    the positions of the selected (sampled) landmarks in 3D space.

        cameras (PinholeCamera):
            A Kornia PinholeCamera object with the intrinsics and extrinsics for the batch.

        gt_depths (batch_size x 1 x img_height x img_width tensor):
            ground truth depth in meters for each pixel

        Updates
        ------
        landmarks (dictionary):
            new keys:
                landmark_pixel_coords (batch_size x n_landmarks x 2 tensor):
                    the pixel coordinates of the selected (sampled) landmarks in each image

                landmark_visibility_mask (batch_size x n_landmarks x 1 bool tensor):
                    a mask indicating which landmarks are visible whether the landmark is visible in a given image (occluded or out of frame)

                landmark_depths (batch_size x n_landmarks x 1 tensor):
                    the depths of the selected (sampled) landmarks in each image
        """
        batch_size = cameras.batch_size
        landmark_pos = repeat(
            landmarks["landmark_pos"],
            "n_landmarks three -> batch_size n_landmarks three",
            batch_size=batch_size,
        )
        n_landmarks = landmark_pos.shape[1]
        img_height = gt_depths.shape[2]
        img_width = gt_depths.shape[3]

        # need to also get the landmark depths, so can't use the PinholeCamera.project method
        landmark_cam_coords = transform_points(cameras.extrinsics, landmark_pos)
        landmark_depths = landmark_cam_coords[..., 2:]
        landmark_pixel_coords = convert_points_from_homogeneous(
            transform_points(cameras.intrinsics, landmark_cam_coords)
        )

        # occlusion check
        landmark_pixel_idxs = torch.round(landmark_pixel_coords).long()
        landmark_pixel_idxs = torch.clamp(
            landmark_pixel_idxs,
            torch.tensor([0, 0], device=self.device),
            torch.tensor([img_width - 1, img_height - 1], device=self.device),
        )
        gt_landmark_depths = batched_2d_index_select(
            gt_depths, torch.flip(landmark_pixel_idxs, dims=(-1,))
        )

        # add some small tolerance
        landmark_visibility_mask = (
            (landmark_depths < gt_landmark_depths + 0.05)
            & (landmark_pixel_coords[..., 0:1] > 0 - 1e-3)
            & (landmark_pixel_coords[..., 1:2] > 0 - 1e-3)
            & (landmark_pixel_coords[..., 0:1] < img_width + 1e-3)
            & (landmark_pixel_coords[..., 1:2] < img_height + 1e-3)
        )

        landmarks.update(
            {
                "landmark_pixel_coords": landmark_pixel_coords,
                "landmark_visibility_mask": landmark_visibility_mask,
                "landmark_depths": landmark_depths,
            }
        )

    def _filter_landmarks(self, landmarks):
        """Filter out unsuitable landmarks.

        Parameters
        ----------
        landmarks (dict):
            Information about the landmarks.

        Updates
        -------
        landmarks (dict):
            filters out some unsuitable landmarks
        """
        # filter out landmarks that are only visible in one image
        mask = landmarks["landmark_visibility_mask"].sum(dim=0) > 1
        mask = mask.squeeze(-1)

        for key in landmarks.keys():
            if len(landmarks[key].shape) == 2:
                landmarks[key] = landmarks[key][mask]
            elif len(landmarks[key].shape) == 3:
                landmarks[key] = landmarks[key][:, mask]
            else:
                raise ValueError(f"Unexpected shape for {key}: {landmarks[key].shape}")

    def _compute_landmark_embeddings(self, landmarks, features):
        """Modifies the landmarks dictionary in-place, adding the computed landmark embeddings.

        Parameters
        ----------
        landmarks (dict):
            the dictionary containing information about the landmarks.

        features (batch_size x feature_dim x feature_height x feature_width tensor):
            the extracted features

        Updates
        -------
        landmarks (dict):
            landmark_embeddings (n_landmarks x feature_dim tensor):
                the embeddings for the landmarks; should be normalised.
        """
        if self.landmark_embedding_method == "sampled_patch":
            return
        else:
            raise NotImplementedError

    def _calculate_contrastive_masks(self, landmarks, cameras, depths):
        """Calculate the masks for which patches are positive pairs, and which ones are used in as the entire set of patches (all pairs).

        Parameters:
        -----------
        landmarks (dict):
            the dictionary containing information about the landmarks.

        cameras (PinholeCamera):
            A Kornia PinholeCamera object with the intrinsics and extrinsics for the batch.

        depths (batch_size x 1 x feature_height x feature_width tensor):
            ground truth depth in meters for each pixel patch

        Returns:
        --------
        pos_patches: n_landmarks x (batch_size * height * width) bool tensor
            A boolean tensor indicating which patches are positive patches for a given landmark

        pos_neg_patches: n_landmarks x (batch_size * height * width) bool tensor
            A boolean tensor indicating which patches are positive or negative patches for a given landmark
        """
        # first, for every image for every landmark, calculate the rectangular patch within the negative patch radius (the larger of the two)
        # then index all these patches and calculate the 3D locations of only these patches -- this means a lot fewer points to calculate
        # then filter based on distance to the landmarks

        batch_size = cameras.batch_size
        n_landmarks = landmarks["landmark_pos"].shape[0]
        height, width = depths.shape[2:4]

        # calculate rectangular patches that should contain all the patches we might use
        # if neg_radius < pos_radius, use all patches as negative patches --> don't need to calculate 3D locations for them
        if self.neg_radius >= self.pos_radius:
            coarse_patch_window_size = (
                self.neg_radius
                * torch.stack((cameras.fx, cameras.fy), dim=-1).view(batch_size, 1, 2)
                / (landmarks["landmark_depths"] + 1e-8)
            )  # shape batch_size x n_landmarks x 2
        else:
            coarse_patch_window_size = (
                self.pos_radius
                * torch.stack((cameras.fx, cameras.fy), dim=-1).view(batch_size, 1, 2)
                / (landmarks["landmark_depths"] + 1e-8)
            )  # shape batch_size x n_landmarks x 2

        patch_pixel_coords = (
            create_meshgrid(
                height, width, normalized_coordinates=False, device=self.device
            )
            + 0.5
        ) * self.feature_subsample
        patch_pixel_coords = repeat(
            patch_pixel_coords,
            "() h w c -> b n (h w) c",
            b=batch_size,
            n=n_landmarks,
            h=height,
            w=width,
            c=2,
        )

        landmark_pixel_coords = landmarks["landmark_pixel_coords"]

        coarse_patch_mask = (
            (
                patch_pixel_coords[..., 0]
                <= landmark_pixel_coords[..., None, 0]
                + coarse_patch_window_size[..., None, 0]
            )
            & (
                patch_pixel_coords[..., 0]
                >= landmark_pixel_coords[..., None, 0]
                - coarse_patch_window_size[..., None, 0]
            )
            & (
                patch_pixel_coords[..., 1]
                <= landmark_pixel_coords[..., None, 1]
                + coarse_patch_window_size[..., None, 1]
            )
            & (
                patch_pixel_coords[..., 1]
                >= landmark_pixel_coords[..., None, 1]
                - coarse_patch_window_size[..., None, 1]
            )
        )
        coarse_patch_mask = reduce(coarse_patch_mask, "b n hw -> b hw", "sum")

        patch_idxs = torch.nonzero(coarse_patch_mask, as_tuple=True)

        patch_cameras = PinholeCamera(
            cameras.intrinsics[patch_idxs[0]],
            cameras.extrinsics[patch_idxs[0]],
            cameras.height[patch_idxs[0]],
            cameras.width[patch_idxs[0]],
        )
        patch_pixel_coords = patch_pixel_coords[:, 0][patch_idxs].unsqueeze(1)
        patch_depths = rearrange(depths, "b () h w -> b (h w) () ()")[patch_idxs]

        patch_world_pos = patch_cameras.unproject(
            patch_pixel_coords, patch_depths
        ).squeeze(1)

        patch_landmark_dists = torch.linalg.norm(
            patch_world_pos[None] - landmarks["landmark_pos"][:, None], ord=2, dim=-1
        )  # shape n_landmarks x n_patches

        pos_patch_mask = patch_landmark_dists < self.pos_radius
        if self.landmark_embedding_method == "sampled_patch":
            # remove self-pairs
            pos_patch_mask[:, patch_landmark_dists.argmin(dim=1)] = False

        pos_patches = torch.zeros(
            (n_landmarks, batch_size, height * width),
            dtype=torch.bool,
            device=self.device,
        )
        pos_patches[:, patch_idxs[0], patch_idxs[1]] = pos_patch_mask
        pos_patches = rearrange(pos_patches, "n b hw -> n (b hw)")

        if self.neg_radius >= self.pos_radius:
            pos_neg_patch_mask = patch_landmark_dists < self.neg_radius
            if self.landmark_embedding_method == "sampled_patch":
                # remove self-pairs
                pos_neg_patch_mask[:, patch_landmark_dists.argmin(dim=1)] = False

            pos_neg_patches = torch.zeros(
                (n_landmarks, batch_size, height * width),
                dtype=torch.bool,
                device=self.device,
            )
            pos_neg_patches[:, patch_idxs[0], patch_idxs[1]] = pos_neg_patch_mask
            pos_neg_patches = rearrange(pos_neg_patches, "n b hw -> n (b hw)")
        else:
            pos_neg_patches = torch.ones(
                (n_landmarks, batch_size * height * width),
                dtype=torch.bool,
                device=self.device,
            )

        return pos_patches, pos_neg_patches

    def _subsample_patches(self, pos_patches, pos_neg_patches, features):
        """Subsample patches used for computational efficiency. Also collects the features for the patches.

        Parameters
        ----------
        pos_patches: n_landmarks x (batch_size * height * width) bool tensor
            A boolean tensor indicating which patches are positive patches for a given landmark

        pos_neg_patches: n_landmarks x (batch_size * height * width) bool tensor
            A boolean tensor indicating which patches are positive or negative patches for a given landmark

        features (batch_size x n_channels x feature_height x feature_width tensor):
            the features extracted.

        Returns
        -------
        selected_patch_features (n_patches x n_channels tensor):
            the features for the patches that are used

        pos_patches (n_landmarks x n_patches bool tensor):
            A mask for the patches that are used as positive patches

        pos_neg_patches (n_landmarks x n_patches bool tensor):
            A mask for the patches that are used as either positive or negative patches.
            There should be at least one True entry in each row.
        """
        if self.keep_all_positive_patches:
            # make sure all positive patches are always used
            pos_patch_mask = torch.any(pos_patches, dim=0)
        else:
            pos_patch_mask = torch.zeros_like(pos_patches[0])

        n_total_patches = pos_patches.shape[1]
        n_selected_patches = floor(self.frac_patch_subsample * n_total_patches)

        patch_mask = torch.zeros_like(pos_patches[0])
        patch_mask[
            torch.randperm(n_total_patches, device=self.device)[:n_selected_patches]
        ] = True
        patch_mask = patch_mask | pos_patch_mask

        selected_patch_features = rearrange(features, "b c h w -> (b h w) c")[
            patch_mask
        ]
        pos_patches = pos_patches[:, patch_mask]
        pos_neg_patches = pos_neg_patches[:, patch_mask]

        return selected_patch_features, pos_patches, pos_neg_patches

src/lightning/module.py
dakloepfer/locus
import torch
from einops import rearrange


class VecSmoothAP(torch.nn.Module):
    """Computes the Vectorised Smooth Average Precision loss."""

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

        self.config = config
        self.sigmoid_temperature = config.SIGMOID_TEMPERATURE

    def forward(self, landmarks, pos_patches, pos_neg_patches, patch_features):
        """Computes the Vectorised Smooth Average Precision loss given landmarks, masks for positive patches and all patches, and the features.

        Parameters
        ----------
        landmarks (dict):
            minimal keys:
                landmark_embeddings (n_landmarks x feature_dim tensor):
                    the embeddings for the landmarks; should be normalised

        pos_patches (n_landmarks x n_all_patches bool tensor):
            A mask for the patches that are used as positive patches

        pos_neg_patches (n_landmarks x n_all_patches bool tensor):
            A mask for the patches that are used as either positive or negative patches.

        patch_features (n_all_patches x feature_dim tensor):
            The features for all the patches; should be normalised

        Returns
        -------
        loss (scalar tensor):
            -Vectorised Smooth AP as the loss

        log_dict (dict):
            a dictionary with information to log
        """
        landmark_embeddings = landmarks["landmark_embeddings"]

        n_all_patches = patch_features.shape[0]
        n_landmarks = landmark_embeddings.shape[0]

        similarities = landmark_embeddings @ patch_features.t()

        similarities = rearrange(
            similarities,
            "n_landmarks n_all_patches -> (n_landmarks n_all_patches)",
            n_landmarks=n_landmarks,
            n_all_patches=n_all_patches,
        )
        pos_patches = rearrange(
            pos_patches,
            "n_landmarks n_all_patches -> (n_landmarks n_all_patches)",
            n_landmarks=n_landmarks,
            n_all_patches=n_all_patches,
        )
        pos_neg_patches = rearrange(
            pos_neg_patches,
            "n_landmarks n_all_patches -> (n_landmarks n_all_patches)",
            n_landmarks=n_landmarks,
            n_all_patches=n_all_patches,
        )

        # similarity_differences[i, j] = similarity of the j-th positive/negative pair - similarity of the i-th positive pair
        # shape = (n_positive_pairs, n_pos_neg_pairs)
        similarity_differences = similarities[pos_neg_patches].unsqueeze(
            dim=0
        ) - similarities[pos_patches].unsqueeze(dim=1)

        # pass through sigmoid
        similarity_differences = torch.sigmoid(
            similarity_differences / self.sigmoid_temperature
        )

        # Compute the rankings of the positive pairs with respect to the entire set
        ranking_pos_neg = 1 + torch.sum(similarity_differences, dim=1)

        # Compute the rankings of the positive pairs with respect to the positive pairs
        ranking_pos = 1 + torch.sum(
            similarity_differences * pos_patches[pos_neg_patches].unsqueeze(dim=0),
            dim=1,
        )

        vec_smooth_ap = torch.mean(ranking_pos / ranking_pos_neg)

        loss = -vec_smooth_ap

        log_dict = {
            "vec_smooth_ap": vec_smooth_ap.detach(),
            "loss": loss.detach(),
            "avg_pos_similarity": similarities[pos_patches].mean().detach(),
            "avg_neg_similarity": similarities[pos_neg_patches & (~pos_patches)]
            .mean()
            .detach(),
            "avg_dontcare_similarity": similarities[~pos_neg_patches].mean().detach(),
        }

        return loss, log_dict

src/losses/vec_smooth_ap.py
dakloepfer/locus
import os
import pickle
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import ConcatDataset, DataLoader
from torchvision.models import resnet50, ResNet50_Weights
from sklearn.decomposition import PCA
from einops import rearrange

from src.data.matterport_dataset import MatterportDataset


class PCA_Baseline(nn.Module):
    """Base class for baseline models that handels all the PCA stuff to reduce the feature dimensionality"""

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.data_conf = config.DATASET
        self.feature_dim = config.MODEL.FEATURE_DIM
        self.network = None

        self.pca_path = (
            f"{config.MODEL.TYPE}_{config.MODEL.FEATURE_DIM}features_pca.pkl"
        )
        self.training_features_for_pca = []
        self.pca = None

    def _setup_dataset(
        self,
    ):
        """To make it a bit easier to set up different datasets"""
        data_source = self.data_conf.TRAIN_DATA_SOURCE
        data_root = self.data_conf.TRAIN_DATA_ROOT
        scene_list_path = self.data_conf.TRAIN_SCENE_LIST
        intrinsics_path = self.data_conf.TRAIN_INTRINSICS_PATH
        matterport_config = {
            "horizontal_only": self.data_conf.DATASET.MATTERPORT_HORIZONTAL_IMGS_ONLY,
            "normalize": self.data_conf.DATASET.MATTERPORT_NORMALIZE,
            "img_height": self.data_conf.DATASET.IMG_HEIGHT,
            "img_width": self.data_conf.DATASET.IMG_WIDTH,
        }
        mode = "train"

        with open(scene_list_path, "r") as f:
            scene_names = [name.split()[0] for name in f.readlines()]

        datasets = []
        for scene_name in tqdm(
            scene_names,
            desc=f"Loading {mode} datasets",
        ):
            if data_source == "Matterport":
                datasets.append(
                    MatterportDataset(
                        data_root,
                        scene_name,
                        intrinsics_path,
                        mode=mode,
                        pose_dir=None,
                        augment_fn=None,
                        **matterport_config,
                    )
                )
            else:
                raise NotImplementedError(f"Data source {data_source} not implemented")

        return ConcatDataset(datasets)

    def make_pca(self, imgs):  # TODO also load samples here
        if os.path.exists(self.pca_path):
            with open(self.pca_path, "rb") as f:
                self.load_pca(pickle.load(f))
        else:
            dataset = self._setup_dataset()
            loader = DataLoader(dataset, batch_size=32, num_workers=8, pin_memory=True)
            with torch.no_grad():
                n_batches = 0
                for batch in tqdm(
                    loader, f"Getting the PCA for the {self.config.MODEL.TYPE} baseline"
                ):
                    imgs = (
                        batch["img"].to("cuda:0")
                        if torch.cuda.is_available()
                        else batch["img"]
                    )

                    features = self.network(imgs)
                    self.training_features_for_pca.append(
                        rearrange(features.cpu(), "b c h w -> (b h w) c")
                    )
                    n_batches += 1
                    if n_batches >= 400:
                        break
                self.calculate_pca()

            with open(
                self.pca_path,
                "wb+",
            ) as pca_file:
                pickle.dump(self.pca, pca_file)

    def calculate_pca(self):
        """Use the stored training features to calculate principal components"""
        if len(self.training_features_for_pca) == 0:
            raise ValueError("No training features stored")

        self.training_features_for_pca = torch.cat(
            self.training_features_for_pca, dim=0
        ).numpy()
        self.pca = PCA(self.feature_dim)
        self.pca.fit(self.training_features_for_pca)
        self.training_features_for_pca = None
        self.pca_mean = (
            torch.tensor(
                self.pca.mean_, dtype=torch.float, device=next(self.parameters()).device
            )
            if self.pca.mean_ is not None
            else torch.zeros(1, self.feature_dim, device=next(self.parameters()).device)
        )
        self.pca_components = torch.tensor(
            self.pca.components_,
            dtype=torch.float,
            device=next(self.parameters()).device,
        )

    def load_pca(self, pca):
        self.pca = pca
        self.training_features_for_pca = None
        self.pca_mean = (
            torch.tensor(
                self.pca.mean_, dtype=torch.float, device=next(self.parameters()).device
            )
            if self.pca.mean_ is not None
            else torch.zeros(1, self.feature_dim, device=next(self.parameters()).device)
        )
        self.pca_components = torch.tensor(
            self.pca.components_,
            dtype=torch.float,
            device=next(self.parameters()).device,
        )

    def forward(self, imgs):
        features = self.network(imgs)
        B, _, H, W = features.shape

        if self.feature_dim != self.out_dim:
            if self.pca is not None:
                features = rearrange(features, "b c h w -> (b h w) c")
                features = features - self.pca_mean
                features = torch.matmul(features, self.pca_components.T)
                features = rearrange(
                    features, "(b h w) c -> b c h w", b=B, h=H, w=W, c=self.feature_dim
                )
            else:
                raise ValueError("PCA not calculated yet")

        return features


class DINO_Network(nn.Module):
    def __init__(self, n_blocks=11, model="vitb8"):
        super().__init__()

        assert n_blocks <= 11
        assert model in ["vitb8", "vits8"]
        self.output_subsample = 8

        self.n_blocks = n_blocks

        self.dino_model = torch.hub.load(
            "facebookresearch/dino:main", "dino_{}".format(model)
        ).eval()

        self.out_dim = 384 if model == "vits8" else 768

        self.dino_model.blocks = self.dino_model.blocks[:n_blocks]
        for param in self.dino_model.parameters():
            param.requires_grad = False

    def forward(self, x):
        B, C, H, W = x.shape

        with torch.no_grad():
            x = self.dino_model.prepare_tokens(x)

            # if I stop before the last block, the elements should have all the relevant information about the image patches
            # the attention masks that they use to segment images attend to these features
            for blk in self.dino_model.blocks[: self.n_blocks]:
                x = blk(x)

        return rearrange(
            x[:, 1:], "b (h w) c -> b c h w", h=int(H // 8), w=int(W // 8)
        )  # first token does not derive from an image patch


class DINO_Baseline(PCA_Baseline):
    # Just a wrapper to keep the code consistent with the actual models
    def __init__(self, config):
        super().__init__(config)

        n_blocks = config.MODEL.N_PRETRAINED_BLOCKS
        model = config.MODEL.VIT_BACKBONE

        assert n_blocks <= 11
        self.network = DINO_Network(n_blocks=n_blocks, model=model)
        self.out_dim = self.network.out_dim
        self.output_subsample = self.network.output_subsample

        if self.feature_dim != self.out_dim:
            self.make_pca()


class ResNet50Baseline(PCA_Baseline):
    def __init__(self, config):
        super().__init__(config)

        self.weights = ResNet50_Weights.DEFAULT
        resnet = resnet50(weights=self.weights)
        resnet.eval()
        self.network = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
            resnet.layer1,
            resnet.layer2,
            resnet.layer3,
            resnet.layer4,
        )
        self.network.eval()

        self.output_subsample = 32
        self.out_dim = 512

        if self.feature_dim != self.out_dim:
            self.make_pca()

src/models/baselines.py
dakloepfer/locus
import torch
import torch.nn as nn
from sklearn.decomposition import PCA
from einops import rearrange


class DINO_Network(nn.Module):
    def __init__(self, n_blocks=11, model="vitb8"):
        super().__init__()

        assert n_blocks <= 11
        assert model in ["vitb8", "vits8"]
        self.output_subsample = 8

        self.n_blocks = n_blocks

        self.dino_model = torch.hub.load(
            "facebookresearch/dino:main", "dino_{}".format(model)
        ).eval()

        self.out_dim = 384 if model == "vits8" else 768

        self.dino_model.blocks = self.dino_model.blocks[:n_blocks]
        for param in self.dino_model.parameters():
            param.requires_grad = False

    def forward(self, x):
        B, C, H, W = x.shape

        with torch.no_grad():
            x = self.dino_model.prepare_tokens(x)

            # if I stop before the last block, the elements should have all the relevant information about the image patches
            # the attention masks that they use to segment images attend to these features
            for blk in self.dino_model.blocks[: self.n_blocks]:
                x = blk(x)

        return rearrange(
            x[:, 1:], "b (h w) c -> b c h w", h=int(H // 8), w=int(W // 8)
        )  # first token does not derive from an image patch


class DINO_Baseline(nn.Module):
    # Just a wrapper to keep the code consistent with the actual models
    def __init__(self, config):
        super().__init__()

        n_blocks = config.N_PRETRAINED_BLOCKS
        feature_dim = config.FEATURE_DIM
        model = config.VIT_BACKBONE

        assert n_blocks <= 11
        self.network = DINO_Network(n_blocks=n_blocks, model=model)
        self.out_dim = self.network.out_dim
        self.feature_dim = feature_dim
        self.training_features_for_pca = []
        self.pca = None

    def calc_and_store_features_for_pca(self, imgs):
        features = self.network(imgs)
        self.training_features_for_pca.append(
            rearrange(features.cpu(), "b c h w -> (b h w) c")
        )

    def calculate_pca(self):
        """Use the stored training features to calculate principal components"""
        if len(self.training_features_for_pca) == 0:
            raise ValueError("No training features stored")

        self.training_features_for_pca = torch.cat(
            self.training_features_for_pca, dim=0
        ).numpy()
        self.pca = PCA(self.feature_dim)
        self.pca.fit(self.training_features_for_pca)
        self.training_features_for_pca = None
        self.pca_mean = (
            torch.tensor(
                self.pca.mean_, dtype=torch.float, device=next(self.parameters()).device
            )
            if self.pca.mean_ is not None
            else torch.zeros(1, self.feature_dim, device=next(self.parameters()).device)
        )
        self.pca_components = torch.tensor(
            self.pca.components_,
            dtype=torch.float,
            device=next(self.parameters()).device,
        )

    def load_pca(self, pca):
        self.pca = pca
        self.training_features_for_pca = None
        self.pca_mean = (
            torch.tensor(
                self.pca.mean_, dtype=torch.float, device=next(self.parameters()).device
            )
            if self.pca.mean_ is not None
            else torch.zeros(1, self.feature_dim, device=next(self.parameters()).device)
        )
        self.pca_components = torch.tensor(
            self.pca.components_,
            dtype=torch.float,
            device=next(self.parameters()).device,
        )

    def forward(self, imgs):
        features = self.network(imgs)
        B, _, H, W = features.shape

        if self.feature_dim != self.out_dim:
            if self.pca is not None:
                features = rearrange(features, "b c h w -> (b h w) c")
                features = features - self.pca_mean
                features = torch.matmul(features, self.pca_components.T)
                features = rearrange(
                    features, "(b h w) c -> b c h w", b=B, h=H, w=W, c=self.feature_dim
                )
            else:
                raise ValueError("PCA not calculated yet")

        return features

src/models/dino_baseline.py
dakloepfer/locus
"""Feature extractor that uses frozen, pre-trained DINO model and several linear layers on top."""

import torch
import torch.nn as nn
from einops import rearrange
from timm.models import vision_transformer


######################## Copied from facebookresearch_dino_main/vision_transformer.py ########################
def drop_path(x, drop_prob: float = 0.0, training: bool = False):
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""

    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


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

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


class Attention(nn.Module):
    def __init__(
        self,
        dim,
        num_heads=8,
        qkv_bias=False,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
    ):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape

        qkv = (
            self.qkv(x)
            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
            .permute(2, 0, 3, 1, 4)
        )
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x, attn


class Block(nn.Module):
    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, x, return_attention=False):
        y, attn = self.attn(self.norm1(x))
        if return_attention:
            return attn
        x = x + self.drop_path(y)
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


############################################################################################################


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

        assert config.N_PRETRAINED_BLOCKS <= 11

        self.output_subsample = 8
        self.n_pretrained_blocks = config.N_PRETRAINED_BLOCKS

        self.dino_model = torch.hub.load(
            "facebookresearch/dino:main", "dino_{}".format(config.VIT_BACKBONE)
        ).eval()

        backbone_dim = 384 if config.VIT_BACKBONE == "vits8" else 768

        self.trainable_head = nn.Sequential(
            *(
                [nn.Linear(backbone_dim, 128)]
                + [
                    Block(128, num_heads=4, qkv_bias=True, qk_scale=0.125)
                    for _ in range(2)
                ]
                + [nn.Linear(128, config.FEATURE_DIM)]
            )
        )

        self.dino_model.blocks = self.dino_model.blocks[: self.n_pretrained_blocks]

        for param in self.dino_model.parameters():
            param.requires_grad = False

    def forward(self, x):
        B, C, H, W = x.shape

        with torch.no_grad():
            x = self.dino_model.prepare_tokens(x)
            for blk in self.dino_model.blocks[: self.n_pretrained_blocks]:
                x = blk(x)

        # turn into extracted features
        x = self.trainable_head(x)
        x = x[:, 1:]

        x = rearrange(x, "b (h w) c -> b c h w", h=int(H // 8), w=int(W // 8))

        return x

    def get_dino_features(self, x):
        B, C, H, W = x.shape

        with torch.no_grad():
            x = self.dino_model.prepare_tokens(x)
            for blk in self.dino_model.blocks[:-1]:
                x = blk(x)

        return rearrange(
            x[:, 1:], "b (h w) c -> b c h w", h=int(H // 8), w=int(W // 8)
        )  # first token does not derive from an image patch

src/models/dino_feature_extractor.py
dakloepfer/locus
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet50, ResNet50_Weights
from sklearn.decomposition import PCA
from einops import rearrange


class ResNet50Baseline(nn.Module):
    def __init__(self, config):
        super().__init__()
        feature_dim = config.FEATURE_DIM
        assert feature_dim == 64

        self.weights = ResNet50_Weights.DEFAULT
        resnet = resnet50(weights=self.weights)
        resnet.eval()
        self.network = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
            resnet.layer1,
            resnet.layer2,
            resnet.layer3,
            resnet.layer4,
        )
        self.network.eval()

        self.output_subsample = 32

        self.out_dim = 512
        self.feature_dim = feature_dim
        self.training_features_for_pca = []
        self.pca = None

    def calc_and_store_features_for_pca(self, imgs):
        features = self.network(imgs)
        self.training_features_for_pca.append(
            rearrange(features.cpu(), "b c h w -> (b h w) c")
        )

    def calculate_pca(self):
        """Use the stored training features to calculate principal components"""
        if len(self.training_features_for_pca) == 0:
            raise ValueError("No training features stored")

        self.training_features_for_pca = torch.cat(
            self.training_features_for_pca, dim=0
        ).numpy()
        self.pca = PCA(self.feature_dim)
        self.pca.fit(self.training_features_for_pca)
        self.training_features_for_pca = None
        self.pca_mean = (
            torch.tensor(
                self.pca.mean_, dtype=torch.float, device=next(self.parameters()).device
            )
            if self.pca.mean_ is not None
            else torch.zeros(1, self.feature_dim, device=next(self.parameters()).device)
        )
        self.pca_components = torch.tensor(
            self.pca.components_,
            dtype=torch.float,
            device=next(self.parameters()).device,
        )

    def load_pca(self, pca):
        self.pca = pca
        self.training_features_for_pca = None
        self.pca_mean = (
            torch.tensor(
                self.pca.mean_, dtype=torch.float, device=next(self.parameters()).device
            )
            if self.pca.mean_ is not None
            else torch.zeros(1, self.feature_dim, device=next(self.parameters()).device)
        )
        self.pca_components = torch.tensor(
            self.pca.components_,
            dtype=torch.float,
            device=next(self.parameters()).device,
        )

    def forward(self, imgs):
        features = self.network(imgs)
        B, _, H, W = features.shape
        assert H == 32, "{} not 32".format(H)
        assert W == 40, "{} not 40".format(W)

        if self.feature_dim != self.out_dim:
            if self.pca is not None:
                features = rearrange(features, "b c h w -> (b h w) c")
                features = features - self.pca_mean
                features = torch.matmul(features, self.pca_components.T)
                features = rearrange(
                    features, "(b h w) c -> b c h w", b=B, h=H, w=W, c=self.feature_dim
                )
            else:
                raise ValueError("PCA not calculated yet")
        return features

src/models/resnet50_baseline.py
dakloepfer/locus
import albumentations as A


class DarkAug(object):
    """
    Extreme dark augmentation aiming at Aachen Day-Night
    """

    def __init__(self) -> None:
        self.augmentor = A.Compose(
            [
                A.RandomBrightnessContrast(
                    p=0.75, brightness_limit=(-0.6, 0.0), contrast_limit=(-0.5, 0.3)
                ),
                A.Blur(p=0.1, blur_limit=(3, 9)),
                A.MotionBlur(p=0.2, blur_limit=(3, 25)),
                A.RandomGamma(p=0.1, gamma_limit=(15, 65)),
                A.HueSaturationValue(p=0.1, val_shift_limit=(-100, -40)),
            ],
            p=0.75,
        )

    def __call__(self, x):
        return self.augmentor(image=x)["image"]


class MobileAug(object):
    """
    Random augmentations aiming at images of mobile/handhold devices.
    """

    def __init__(self, outdoor=True):
        if outdoor:
            self.augmentor = A.Compose(
                [
                    A.MotionBlur(p=0.15),
                    A.ColorJitter(p=0.5),
                    A.RandomRain(p=0.05),  # random occlusion
                    A.RandomSunFlare(p=0.1),
                    A.JpegCompression(p=0.25),
                    A.ISONoise(p=0.25),
                ],
                p=1.0,
            )
        else:
            self.augmentor = A.Compose(
                [
                    A.MotionBlur(p=0.15),
                    A.ColorJitter(p=0.5),
                    A.JpegCompression(p=0.25),
                    A.ISONoise(p=0.25),
                ],
                p=1.0,
            )

    def __call__(self, x):
        return self.augmentor(image=x)["image"]


class DroneAug(object):
    """
    Random augmentations aiming at images taken by drones.
    """

    def __init__(self, outdoor=True):
        if outdoor:
            self.augmentor = A.Compose(
                [
                    A.MotionBlur(p=0.25),
                    A.ColorJitter(p=0.5),
                    A.RandomRain(p=0.1),  # random occlusion
                    A.RandomSunFlare(p=0.1),
                    A.JpegCompression(p=0.25),
                    A.ISONoise(p=0.25),
                ],
                p=1.0,
            )
        else:
            self.augmentor = A.Compose(
                [
                    A.MotionBlur(p=0.25),
                    A.ColorJitter(p=0.5),
                    A.JpegCompression(p=0.25),
                    A.ISONoise(p=0.25),
                ],
                p=1.0,
            )

    def __call__(self, x):
        return self.augmentor(image=x)["image"]


def build_augmentor(method=None, data_source=None, **kwargs):
    if method == "dark":
        return DarkAug()
    elif method == "mobile":
        if data_source.lower() in ["eurocmav", "scannet", "inloc"]:
            outdoor = False
        else:
            outdoor = True

        return MobileAug(outdoor=outdoor)
    elif method == "drone":
        if data_source.lower() in ["eurocmav", "scannet", "inloc"]:
            outdoor = False
        else:
            outdoor = True

        return DroneAug(outdoor=outdoor)

    elif method is None:
        return None
    else:
        raise ValueError(f"Invalid augmentation method: {method}")


if __name__ == "__main__":
    augmentor = build_augmentor("FDA")

src/utils/augment.py
dakloepfer/locus
import os
import inspect
import contextlib
import joblib
from typing import Union, List, Tuple, Dict, Any
from loguru import _Logger, logger
from itertools import chain
from argparse import ArgumentParser, _ArgumentGroup, Namespace

import torch
from yacs.config import CfgNode as CN
from pytorch_lightning.utilities import rank_zero_only
from lightning import Trainer
import cv2
import numpy as np
from einops import rearrange, repeat


def batched_2d_index_select(input_tensor, indices):
    """
    Arguments:
        input_tensor: b x c x h x w
        indices: b x n x 2 (y, x in the last dimension)
    Returns:
        output: b x n x c

    output[b,n,c] = input_tensor[b, c, indices[b,n,0], indices[b,n,1]]
    """
    b, c, h, w = input_tensor.shape
    input_tensor = rearrange(input_tensor, "b c h w -> b (h w) c")
    indices = indices[:, :, 0] * w + indices[:, :, 1]
    indices = repeat(indices, "b n -> b n c", c=c)
    output = torch.gather(input_tensor, dim=1, index=indices)
    return output


def invert_se3(T):
    """Invert an SE(3) transformation matrix."""
    assert T.shape[-2:] == (4, 4), "T must be of shape (..., 4, 4)"

    rot = T[..., :3, :3]
    trans = T[..., :3, 3]

    if type(T) == torch.Tensor:
        inv_T = torch.zeros_like(T)
        inv_rot = rot.transpose(-1, -2)
        inv_trans = torch.einsum("...ij,...j->...i", -inv_rot, trans)

    else:  # numpy
        inv_T = np.zeros_like(T)
        inv_rot = np.swapaxes(rot, -1, -2)
        inv_trans = np.einsum("...ij,...j->...i", -inv_rot, trans)

    inv_T[..., :3, :3] = inv_rot
    inv_T[..., :3, 3] = inv_trans
    inv_T[..., 3, 3] = 1.0

    return inv_T


def lower_config(yacs_cfg):
    if not isinstance(yacs_cfg, CN):
        return yacs_cfg
    return {k.lower(): lower_config(v) for k, v in yacs_cfg.items()}


def upper_config(dict_cfg):
    if not isinstance(dict_cfg, dict):
        return dict_cfg
    return {k.upper(): upper_config(v) for k, v in dict_cfg.items()}


def apply_to_subconfig(func, cfg: CN, subcfg_name: str, **kwargs):
    subcfg_name_list = subcfg_name.strip().split(".")
    if len(subcfg_name_list) == 1:
        setattr(cfg, subcfg_name, func(getattr(cfg, subcfg_name), **kwargs))
    else:
        apply_to_subconfig(
            func,
            getattr(cfg, subcfg_name_list[0]),
            ".".join(subcfg_name_list[1:]),
            **kwargs,
        )


def make_list(x, n=-1):
    if isinstance(x, list):
        return x
    else:
        return [x for _ in range(n)]


def average_dict_list(dict_list: List[Dict[Any, Any]]) -> Dict[Any, Any]:
    """Average a list of dictionaries, recursively if needed.

    Parameters
    ----------
    dict_list (List[Dict[Any, Any]]):
        the list of dictionaries to average. All dictionaries must have the same keys.

    Returns
    -------
    Dict[Any, Any]:
        a dictionary with the same keys, but with values averaged over the list.
    """
    if len(dict_list) == 0:
        return {}

    if len(dict_list) == 1:
        return dict_list[0]

    return_dict = {}
    for k in dict_list.keys():
        if isinstance(dict_list[0][k], dict):
            return_dict[k] = average_dict_list([d[k] for d in dict_list])
        else:
            return_dict[k] = sum([d[k] for d in dict_list]) / len(dict_list)

    return return_dict


def log_on(condition, message, level):
    if condition:
        assert level in ["INFO", "DEBUG", "WARNING", "ERROR", "CRITICAL"]
        logger.log(level, message)


def get_rank_zero_only_logger(logger: _Logger):
    if rank_zero_only.rank == 0:
        return logger
    else:
        for _level in logger._core.levels.keys():
            level = _level.lower()
            setattr(logger, level, lambda x: None)
        logger._log = lambda x: None
    return logger


def flattenList(x):
    return list(chain(*x))


@contextlib.contextmanager
def tqdm_joblib(tqdm_object):
    """Context manager to patch joblib to report into tqdm progress bar given as argument

    Usage:
        with tqdm_joblib(tqdm(desc="My calculation", total=10)) as progress_bar:
            Parallel(n_jobs=16)(delayed(sqrt)(i**2) for i in range(10))

    When iterating over a generator, directly use of tqdm is also a solutin (but monitor the task queuing, instead of finishing)
        ret_vals = Parallel(n_jobs=args.world_size)(
                    delayed(lambda x: _compute_cov_score(pid, *x))(param)
                        for param in tqdm(combinations(image_ids, 2),
                                          desc=f'Computing cov_score of [{pid}]',
                                          total=len(image_ids)*(len(image_ids)-1)/2))
    Src: https://stackoverflow.com/a/58936697
    """

    class TqdmBatchCompletionCallback(joblib.parallel.BatchCompletionCallBack):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)

        def __call__(self, *args, **kwargs):
            tqdm_object.update(n=self.batch_size)
            return super().__call__(*args, **kwargs)

    old_batch_callback = joblib.parallel.BatchCompletionCallBack
    joblib.parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
    try:
        yield tqdm_object
    finally:
        joblib.parallel.BatchCompletionCallBack = old_batch_callback
        tqdm_object.close()


def str_to_bool_or_str(val: str) -> Union[str, bool]:
    """Possibly convert a string representation of truth to bool. Returns the input otherwise. Based on the python
    implementation distutils.utils.strtobool.

    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values are 'n', 'no', 'f', 'false', 'off', and '0'.
    """
    lower = val.lower()
    if lower in ("y", "yes", "t", "true", "on", "1"):
        return True
    if lower in ("n", "no", "f", "false", "off", "0"):
        return False
    return val


def str_to_bool(val: str) -> bool:
    """Convert a string representation of truth to bool.

    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values
    are 'n', 'no', 'f', 'false', 'off', and '0'.

    Raises:
        ValueError:
            If ``val`` isn't in one of the aforementioned true or false values.

    >>> str_to_bool('YES')
    True
    >>> str_to_bool('FALSE')
    False
    """
    val_converted = str_to_bool_or_str(val)
    if isinstance(val_converted, bool):
        return val_converted
    raise ValueError(f"invalid truth value {val_converted}")


def str_to_bool_or_int(val: str) -> Union[bool, int, str]:
    """Convert a string representation to truth of bool if possible, or otherwise try to convert it to an int.

    >>> str_to_bool_or_int("FALSE")
    False
    >>> str_to_bool_or_int("1")
    True
    >>> str_to_bool_or_int("2")
    2
    >>> str_to_bool_or_int("abc")
    'abc'
    """
    val_converted = str_to_bool_or_str(val)
    if isinstance(val_converted, bool):
        return val_converted
    try:
        return int(val_converted)
    except ValueError:
        return val_converted


def _gpus_allowed_type(x) -> Union[int, str]:
    if "," in x:
        return str(x)
    else:
        return int(x)


def _int_or_float_type(x) -> Union[int, float]:
    if "." in str(x):
        return float(x)
    else:
        return int(x)


def _get_abbrev_qualified_cls_name(cls):
    assert isinstance(cls, type), repr(cls)
    if cls.__module__.startswith("pytorch_lightning."):
        # Abbreviate.
        return f"pl.{cls.__name__}"
    else:
        # Fully qualified.
        return f"{cls.__module__}.{cls.__qualname__}"


def get_init_arguments_and_types(cls) -> List[Tuple[str, Tuple, Any]]:
    r"""Scans the class signature and returns argument names, types and default values.

    Returns:
        List with tuples of 3 values:
        (argument name, set with argument types, argument default value).

    Examples:

        >>> from pytorch_lightning import Trainer
        >>> args = get_init_arguments_and_types(Trainer)

    """
    cls_default_params = inspect.signature(cls).parameters
    name_type_default = []
    for arg in cls_default_params:
        arg_type = cls_default_params[arg].annotation
        arg_default = cls_default_params[arg].default
        try:
            arg_types = tuple(arg_type.__args__)
        except AttributeError:
            arg_types = (arg_type,)

        name_type_default.append((arg, arg_types, arg_default))

    return name_type_default


def _parse_args_from_docstring(docstring: str) -> Dict[str, str]:
    arg_block_indent = None
    current_arg = None
    parsed = {}
    for line in docstring.split("\n"):
        stripped = line.lstrip()
        if not stripped:
            continue
        line_indent = len(line) - len(stripped)
        if stripped.startswith(("Args:", "Arguments:", "Parameters:")):
            arg_block_indent = line_indent + 4
        elif arg_block_indent is None:
            continue
        elif line_indent < arg_block_indent:
            break
        elif line_indent == arg_block_indent:
            current_arg, arg_description = stripped.split(":", maxsplit=1)
            parsed[current_arg] = arg_description.lstrip()
        elif line_indent > arg_block_indent:
            parsed[current_arg] += f" {stripped}"
    return parsed


def parse_argparser(cls, arg_parser: Union[ArgumentParser, Namespace]) -> Namespace:
    """Parse CLI arguments, required for custom bool types. Copied from https://pytorch-lightning.readthedocs.io/en/1.3.8/_modules/pytorch_lightning/utilities/argparse.html#parse_argparser"""
    args = (
        arg_parser.parse_args()
        if isinstance(arg_parser, ArgumentParser)
        else arg_parser
    )

    types_default = {
        arg: (arg_types, arg_default)
        for arg, arg_types, arg_default in get_init_arguments_and_types(cls)
    }

    modified_args = {}
    for k, v in vars(args).items():
        if k in types_default and v is None:
            # We need to figure out if the None is due to using nargs="?" or if it comes from the default value
            arg_types, arg_default = types_default[k]
            if bool in arg_types and isinstance(arg_default, bool):
                # Value has been passed as a flag => It is currently None, so we need to set it to True
                # We always set to True, regardless of the default value.
                # Users must pass False directly, but when passing nothing True is assumed.
                # i.e. the only way to disable something that defaults to True is to use the long form:
                # "--a_default_true_arg False" becomes False, while "--a_default_false_arg" becomes None,
                # which then becomes True here.

                v = True

        modified_args[k] = v
    return Namespace(**modified_args)


def add_pl_argparse_args(cls, parent_parser: ArgumentParser):
    """Extends existing argparse by default attributes for a Pytorch Lightning Trainer.
    Replaces the (annoyingly deprecated) PytorchLightning function pl.Trainer.add_argparse_args().

    Code largely copied from https://pytorch-lightning.readthedocs.io/en/1.3.8/_modules/pytorch_lightning/utilities/argparse.html#add_argparse_args , with use_argument_group=True.

    Args:
        cls: a Lightning class (in practice a Trainer)
        parent_parser:
            The custom CLI arguments parser, which will be extended by the default Trainer default arguments.

    Returns:
        the parent parser with added default Trainer arguments.
    """
    parent_parser_args = [action.dest for action in parent_parser._actions]
    group_name = _get_abbrev_qualified_cls_name(cls)
    parser = parent_parser.add_argument_group(group_name)

    ignore_arg_names = ["self", "args", "kwargs"]
    if hasattr(cls, "get_deprecated_arg_names"):
        ignore_arg_names += cls.get_deprecated_arg_names()

    allowed_types = (str, int, float, bool)

    # Get symbols from cls or init function.
    for symbol in (cls, cls.__init__):
        args_and_types = get_init_arguments_and_types(symbol)
        args_and_types = [x for x in args_and_types if x[0] not in ignore_arg_names]
        if len(args_and_types) > 0:
            break

    args_help = _parse_args_from_docstring(cls.__init__.__doc__ or cls.__doc__ or "")

    for arg, arg_types, arg_default in args_and_types:
        if arg in parent_parser_args:
            continue
        arg_types = [at for at in allowed_types if at in arg_types]
        if not arg_types:
            # skip argument with not supported type
            continue
        arg_kwargs = {}
        if bool in arg_types:
            arg_kwargs.update(nargs="?", const=True)
            # if the only arg type is bool
            if len(arg_types) == 1:
                use_type = str_to_bool
            elif int in arg_types:
                use_type = str_to_bool_or_int
            elif str in arg_types:
                use_type = str_to_bool_or_str
            else:
                # filter out the bool as we need to use more general
                use_type = [at for at in arg_types if at is not bool][0]
        else:
            use_type = arg_types[0]

        if arg == "gpus" or arg == "tpu_cores":
            use_type = _gpus_allowed_type

        # hack for types in (int, float)
        if len(arg_types) == 2 and int in set(arg_types) and float in set(arg_types):
            use_type = _int_or_float_type

        # hack for track_grad_norm
        if arg == "track_grad_norm":
            use_type = float

        parser.add_argument(
            f"--{arg}",
            dest=arg,
            default=arg_default,
            type=use_type,
            help=args_help.get(arg),
            **arg_kwargs,
        )

    return parent_parser


def pl_from_argparse_args(cls, args: Union[Namespace, ArgumentParser], **kwargs):
    """Manually re-implemented function to create a Pytorch Lightning Trainer from CLI arguments, as originally used in Trainer.from_argparse_args(), which has (annoyingly) been deprecated.

    Code largely taken copied from https://pytorch-lightning.readthedocs.io/en/1.3.8/_modules/pytorch_lightning/utilities/argparse.html#from_argparse_args , with the cls argument fixed as cls=Trainer.

    Args:
        cls: a Lightning class (in practice, a Trainer)
        args: The parser or namespace to take arguments from. Only known arguments will be
            parsed and passed to the :class:`Trainer`.
        **kwargs: Additional keyword arguments that may override ones in the parser or namespace.
            These must be valid Trainer arguments.

    Returns:
        :class:`Trainer`
            Initialised Trainer object.
    """

    if isinstance(args, ArgumentParser):
        args = parse_argparser(cls, args)

    params = vars(args)

    # we only want to pass in valid Trainer args from the parser, the rest may be user specific
    valid_kwargs = inspect.signature(cls.__init__).parameters
    trainer_kwargs = dict(
        (name, params[name]) for name in valid_kwargs if name in params
    )
    trainer_kwargs.update(**kwargs)

    return cls(**trainer_kwargs)

src/utils/misc.py
dakloepfer/locus
import os
import torch
from lightning.pytorch.profilers import SimpleProfiler, PassThroughProfiler
from contextlib import contextmanager
from pytorch_lightning.utilities import rank_zero_only


class InferenceProfiler(SimpleProfiler):
    """
    This profiler records duration of actions with cuda.synchronize()
    Use this in test time.
    """

    def __init__(self):
        super().__init__()
        self.start = rank_zero_only(self.start)
        self.stop = rank_zero_only(self.stop)
        self.summary = rank_zero_only(self.summary)

    @contextmanager
    def profile(self, action_name: str) -> None:
        try:
            torch.cuda.synchronize()
            self.start(action_name)
            yield action_name
        finally:
            torch.cuda.synchronize()
            self.stop(action_name)


def build_profiler(args):
    name = args.profiler_name
    if name == "inference":
        return InferenceProfiler()
    elif name == "pytorch":
        dirpath = os.path.join(args.save_dir, args.exp_name)
        filename = f"{name}_profiler.txt"
        from pytorch_lightning.profilers import PyTorchProfiler

        return PyTorchProfiler(
            dirpath=dirpath,
            filename=filename,
            use_cuda=True,
            profile_memory=True,
            row_limit=100,
        )
    elif name is None:
        return PassThroughProfiler()
    else:
        raise ValueError(f"Invalid profiler: {name}")

src/utils/profiler.py
dakloepfer/locus
import os
import argparse
import pprint
from tqdm import tqdm
import torch
from distutils.util import strtobool
from loguru import logger as loguru_logger

from lightning import pytorch as pl
from pytorch_lightning.utilities import rank_zero_only
from pytorch_lightning.loggers import TensorBoardLogger
from lightning.pytorch.callbacks import (
    ModelSummary,
    DeviceStatsMonitor,
    TQDMProgressBar,
)
from lightning import Trainer


from src.default_config.default import get_cfg_defaults
from src.utils.misc import (
    get_rank_zero_only_logger,
    add_pl_argparse_args,
    pl_from_argparse_args,
)
from src.utils.profiler import build_profiler
from src.lightning.data import MultiSceneDataModule
from src.lightning.module import PLModule
from src.lightning.linear_probe import LinearProbeModule
from src.data.matterport_segmentation_dataset import (
    MatterportSegmentationDataset,
    N_OBJECTS_PER_SCENE,
)

loguru_logger = get_rank_zero_only_logger(loguru_logger)


def parse_args():
    # init a costum parser which will be added into pl.Trainer parser
    # check documentation: https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html#trainer-flags
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("main_cfg_path", type=str, help="main config path")
    parser.add_argument("data_cfg_path", type=str, help="data config path")
    parser.add_argument(
        "--save_dir",
        type=str,
        default="locus_logs/test",
    )
    parser.add_argument("--exp_name", type=str, default="default_exp_name")
    parser.add_argument("--task", type=str, default="patch_retrieval")
    parser.add_argument("--device_list", type=int, nargs="+", default=None)
    parser.add_argument(
        "--batch_size", type=int, default=32, help="total batch size (across all GPUs)"
    )
    parser.add_argument("--num_workers", type=int, default=16)
    parser.add_argument(
        "--pin_memory",
        type=lambda x: bool(strtobool(x)),
        nargs="?",
        default=True,
        help="whether loading data to pinned memory or not",
    )
    parser.add_argument(
        "--ckpt_path",
        type=str,
        default=None,
        help="pretrained checkpoint path",
    )
    parser.add_argument(
        "--profiler_name",
        type=str,
        default=None,
        help="options: [inference, pytorch], or leave it unset",
    )
    parser.add_argument(
        "--tqdm_refresh_rate",
        type=int,
        default=1,
        help="refresh rate for tqdm progress bar",
    )
    parser.add_argument("--linear_probe_max_epochs", type=int, default=1000)

    parser = add_pl_argparse_args(Trainer, parser)
    parser = add_pl_argparse_args(PLModule, parser)
    parser = add_pl_argparse_args(TensorBoardLogger, parser)

    args = parser.parse_args()

    return args


def precalculate_features(args, feature_extractor, dataset):
    # calculate the features, return a TensorDataset with the features and ground truth segmentations as labels
    # Don't downsample the gt_segmentations to the feature resolution
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=args.pin_memory,
    )
    feature_extractor.eval()

    features = []
    gt_segmentations = []

    for batch in tqdm(dataloader, "Pre-calculating Segmentation Features"):
        gt_segmentations.append(batch["segmentation"])
        with torch.no_grad():
            features.append(feature_extractor(batch["img"].to(args.device_list[0])))
    features = torch.cat(features, dim=0)
    gt_segmentations = torch.cat(gt_segmentations, dim=0)
    feature_dataset = torch.utils.data.TensorDataset(features, gt_segmentations)

    return feature_dataset


def main():
    loguru_logger.info(f"Number of available GPUs: {torch.cuda.device_count()}")
    # parse arguments
    args = parse_args()
    rank_zero_only(pprint.pprint)(vars(args))

    # init default-cfg and merge it with the main- and data-cfg
    config = get_cfg_defaults()
    config.merge_from_file(args.main_cfg_path)
    config.merge_from_file(args.data_cfg_path)
    pl.seed_everything(config.TRAINER.SEED)  # reproducibility

    config.freeze()

    profiler = build_profiler(args)

    if args.task in ["patch_retrieval"]:
        # lightning data
        data_module = MultiSceneDataModule(args, config)
        loguru_logger.info(f"DataModule initialized!")

        # lightning module
        assert args.ckpt_path is not None, "ckpt_path must be specified for testing!"
        model = PLModule(
            config,
            ckpt_path=args.ckpt_path,
            devices=args.device_list,
            profiler=profiler,
            test_task=args.task,
        )
        loguru_logger.info(f"LightningModule initialized!")

        # TensorBoard Logger
        logger = TensorBoardLogger(
            save_dir=os.path.join(args.save_dir, args.exp_name),
            name=args.task,
            default_hp_metric=False,
        )

        # Callbacks
        model_summary = ModelSummary(max_depth=10)
        tqdm_progress_bar = TQDMProgressBar(refresh_rate=args.tqdm_refresh_rate)
        callbacks = [
            model_summary,
            tqdm_progress_bar,
        ]
        if args.profiler_name is not None:
            device_stats_monitor = DeviceStatsMonitor(cpu_stats=True)
            callbacks.append(device_stats_monitor)

        # Lightning Trainer
        trainer = pl_from_argparse_args(
            Trainer,
            args,
            callbacks=callbacks,
            logger=logger,
            reload_dataloaders_every_n_epochs=False,  # avoid repeated samples!
            profiler=profiler,
            detect_anomaly=args.detect_anomaly,
            devices=1,  # only use one GPU for the trainer
        )
        loguru_logger.info(f"Trainer initialized!")
        loguru_logger.info(f"Start testing!")

        trainer.test(model, datamodule=data_module)

    elif args.task == "segmentation":
        feature_extractor = PLModule(
            config,
            ckpt_path=args.ckpt_path,
            devices=args.device_list,
            profiler=profiler,
            test_task=args.task,
        )

        with open(config.DATASET.TEST_SCENE_LIST, "r") as f:
            scene_names = [name.split()[0] for name in f.readlines()]

        data_root = config.DATASET.TEST_DATA_ROOT

        overall_results = {
            "mAP": [],
            "mIoU": [],
            "Jaccard": [],
            "Accuracy": [],
            "object_mAP": [],
            "object_mIoU": [],
            "object_Jaccard": [],
            "object_Accuracy": [],
            "stuff_mAP": [],
            "stuff_mIoU": [],
            "stuff_Jaccard": [],
            "stuff_Accuracy": [],
            "n_images": [],
        }
        output_file = os.path.join(
            args.save_dir, args.exp_name, "segmentation_results.txt"
        )
        if not os.path.exists(os.path.dirname(output_file)):
            os.makedirs(os.path.dirname(output_file))

        # write a header to the output_file
        with open(output_file, "w+") as f:
            f.write(
                "|\tScene Name\t|\tmAP\t\t|\tmIoU\t\t|\tJaccard\t\t|\tAccuracy\t|\tobject_mAP\t|\tobject_mIoU\t|\tobject_Jaccard\t|\tobject_Accuracy\t|\tstuff_mAP\t|\tstuff_mIoU\t|\tstuff_Jaccard\t|\tstuff_Accuracy\t|\tn_images\t|\n"
            )
            f.write(
                "_________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________\n"
            )

        for scene_name in scene_names:
            loguru_logger.info('Testing scene "{}"'.format(scene_name))

            segmentation_dataset = MatterportSegmentationDataset(
                data_root,
                scene_name,
                horizontal_only=config.DATASET.MATTERPORT_HORIZONTAL_IMGS_ONLY,
                normalize=config.DATASET.MATTERPORT_NORMALIZE,
                img_height=config.DATASET.IMG_HEIGHT,
                img_width=config.DATASET.IMG_WIDTH,
            )
            feature_dataset = precalculate_features(
                args, feature_extractor, segmentation_dataset
            )

            linear_probe_train_loader = torch.utils.data.DataLoader(
                feature_dataset, batch_size=args.batch_size, shuffle=True
            )
            linear_probe = LinearProbeModule(config, scene_name)

            logger = TensorBoardLogger(
                save_dir=os.path.join(args.save_dir, args.exp_name),
                name=f"linear_probe_training_scene{scene_name}",
                default_hp_metric=False,
            )

            trainer = pl.Trainer(
                max_epochs=args.linear_probe_max_epochs,
                logger=logger,
                enable_progress_bar=True,
                devices=1,
            )
            loguru_logger.info(
                'Start training linear probe for "{}"'.format(scene_name)
            )
            trainer.fit(linear_probe, linear_probe_train_loader)
            loguru_logger.info(f"Finished training linear probe for {scene_name}!")

            linear_probe_test_loader = torch.utils.data.DataLoader(
                feature_dataset, batch_size=args.batch_size, shuffle=False
            )
            loguru_logger.info(
                f"Start calculating segmentation performance on scene {scene_name}!"
            )
            trainer.test(linear_probe, linear_probe_test_loader)

            scene_results = linear_probe.get_results()
            scene_results["n_images"] = len(feature_dataset)

            with open(output_file, "a") as f:
                f.write(
                    f"|\t{scene_name}\t|\t{scene_results['mAP']:.5f}\t\t|\t{scene_results['mIoU']:.5f}\t\t|\t{scene_results['Jaccard']:.5f}\t\t|\t{scene_results['Accuracy']:.5f}\t\t|\t{scene_results['object_mAP']:.5f}\t\t|\t{scene_results['object_mIoU']:.5f}\t\t|\t{scene_results['object_Jaccard']:.5f}\t\t|\t{scene_results['object_Accuracy']:.5f}\t\t|\t{scene_results['stuff_mAP']:.5f}\t\t|\t{scene_results['stuff_mIoU']:.5f}\t\t|\t{scene_results['stuff_Jaccard']:.5f}\t\t|\t{scene_results['stuff_Accuracy']:.5f}\t\t|\t{scene_results['n_images']}\t\t|\n"
                )

            for key in overall_results.keys():
                if key == "n_images":
                    overall_results[key].append(len(feature_dataset))
                    continue
                overall_results[key].append(scene_results[key])
            loguru_logger.info(
                'Finished segmentation testing on scene "{}"'.format(scene_name)
            )
        mean = lambda x: sum(x) / len(x)
        with open(output_file, "a") as f:
            f.write(
                "_________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________________\n"
            )
            f.write(
                f"|\tAverage\t\t|\t{mean(overall_results['mAP']):.5f}\t\t|\t{mean(overall_results['mIoU']):.5f}\t\t|\t{mean(overall_results['Jaccard']):.5f}\t\t|\t{mean(overall_results['Accuracy']):.5f}\t\t|\t{mean(overall_results['object_mAP']):.5f}\t\t|\t{mean(overall_results['object_mIoU']):.5f}\t\t|\t{mean(overall_results['object_Jaccard']):.5f}\t\t|\t{mean(overall_results['object_Accuracy']):.5f}\t\t|\t{mean(overall_results['stuff_mAP']):.5f}\t\t|\t{mean(overall_results['stuff_mIoU']):.5f}\t\t|\t{mean(overall_results['stuff_Jaccard']):.5f}\t\t|\t{mean(overall_results['stuff_Accuracy']):.5f}\t\t|\t{sum(overall_results['n_images'])}\t\t|\n"
            )


if __name__ == "__main__":
    main()

test.py
dakloepfer/locus
import argparse
import pprint
import torch
from distutils.util import strtobool
from pathlib import Path
from loguru import logger as loguru_logger

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_only
from pytorch_lightning.loggers import TensorBoardLogger
from lightning.pytorch.callbacks import (
    ModelCheckpoint,
    LearningRateMonitor,
    ModelSummary,
    DeviceStatsMonitor,
    TQDMProgressBar,
)
from lightning import Trainer


from src.default_config.default import get_cfg_defaults
from src.utils.misc import (
    get_rank_zero_only_logger,
    add_pl_argparse_args,
    pl_from_argparse_args,
)
from src.utils.profiler import build_profiler
from src.lightning.data import MultiSceneDataModule
from src.lightning.module import PLModule

loguru_logger = get_rank_zero_only_logger(loguru_logger)


def parse_args():
    # init a costum parser which will be added into pl.Trainer parser
    # check documentation: https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html#trainer-flags
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("main_cfg_path", type=str, help="main config path")
    parser.add_argument("data_cfg_path", type=str, help="data config path")
    parser.add_argument(
        "--save_dir",
        type=str,
        default="locus_logs/train",
    )
    parser.add_argument("--exp_name", type=str, default="default_exp_name")
    parser.add_argument("--device_list", type=int, nargs="+", default=None)
    parser.add_argument(
        "--batch_size", type=int, default=16, help="total batch size (across all GPUs)"
    )
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument(
        "--pin_memory",
        type=lambda x: bool(strtobool(x)),
        nargs="?",
        default=True,
        help="whether loading data to pinned memory or not",
    )
    parser.add_argument(
        "--ckpt_path",
        type=str,
        default=None,
        help="pretrained checkpoint path",
    )
    parser.add_argument(
        "--disable_ckpt",
        action="store_true",
        help="disable checkpoint saving (useful for debugging).",
    )
    parser.add_argument(
        "--profiler_name",
        type=str,
        default=None,
        help="options: [inference, pytorch], or leave it unset",
    )
    parser.add_argument(
        "--tqdm_refresh_rate",
        type=int,
        default=1,
        help="refresh rate for tqdm progress bar",
    )

    parser = add_pl_argparse_args(Trainer, parser)
    parser = add_pl_argparse_args(PLModule, parser)
    parser = add_pl_argparse_args(TensorBoardLogger, parser)

    args = parser.parse_args()

    return args


def main():
    loguru_logger.info(f"Number of available GPUs: {torch.cuda.device_count()}")
    # parse arguments
    args = parse_args()
    rank_zero_only(pprint.pprint)(vars(args))

    # init default-cfg and merge it with the main- and data-cfg
    config = get_cfg_defaults()
    config.merge_from_file(args.main_cfg_path)
    config.merge_from_file(args.data_cfg_path)
    pl.seed_everything(config.TRAINER.SEED)  # reproducibility

    config.freeze()

    # lightning module
    profiler = build_profiler(args)
    model = PLModule(
        config, ckpt_path=args.ckpt_path, devices=args.device_list, profiler=profiler
    )
    loguru_logger.info(f"LightningModule initialized!")

    # lightning data
    data_module = MultiSceneDataModule(args, config)
    loguru_logger.info(f"DataModule initialized!")

    # TensorBoard Logger
    logger = TensorBoardLogger(
        save_dir=args.save_dir,
        name=args.exp_name,
        default_hp_metric=False,
    )
    ckpt_dir = Path(logger.log_dir) / "checkpoints"

    # Callbacks
    ckpt_callback = ModelCheckpoint(
        monitor="loss",
        verbose=True,
        save_top_k=10,
        mode="min",
        save_last=True,
        dirpath=str(ckpt_dir),
        filename="{epoch}-{avg_loss:.3f}",
    )

    lr_monitor = LearningRateMonitor(logging_interval="step")
    model_summary = ModelSummary(max_depth=10)
    tqdm_progress_bar = TQDMProgressBar(refresh_rate=args.tqdm_refresh_rate)
    callbacks = [
        lr_monitor,
        model_summary,
        tqdm_progress_bar,
    ]
    if not args.disable_ckpt:
        callbacks.append(ckpt_callback)
    if args.profiler_name is not None:
        device_stats_monitor = DeviceStatsMonitor(cpu_stats=True)
        callbacks.append(device_stats_monitor)

    # Lightning Trainer
    trainer = pl_from_argparse_args(
        Trainer,
        args,
        gradient_clip_val=config.TRAINER.GRADIENT_CLIPPING,
        accumulate_grad_batches=config.TRAINER.ACCUMULATE_GRAD_BATCHES,
        callbacks=callbacks,
        logger=logger,
        reload_dataloaders_every_n_epochs=False,  # avoid repeated samples!
        profiler=profiler,
        detect_anomaly=args.detect_anomaly,
        devices=1,  # only use one GPU for the trainer
    )
    loguru_logger.info(f"Trainer initialized!")
    loguru_logger.info(f"Start training!")

    trainer.fit(model, datamodule=data_module)


if __name__ == "__main__":
    main()

train.py
dakloepfer/locus