affordance_transfer_pipeline.py

"""
The pipeline is as follows.

1. Load labels to database
2. Render images to the results folder
3. Calculate affordance map and save
4. Render and vote
5. Evaluate the results
"""

from collections import defaultdict
import pickle as pkl
import numpy as np
import tyro
import imageio
import pycolmap_scene_manager as pycolmap
import torch
import os
import faiss
import base64
import cv2
import json
import PIL.Image as Image
from torchvision import transforms as T
from gsplat import rasterization

if not torch.cuda.is_available():
    raise RuntimeError("CUDA is required for this script")
torch.set_default_device("cuda")
device = "cuda"

SIZE = 224 * 4
FEATURE_MAP_SIZE = 64

LABEL_TO_IDX = {
    "background": 0,
    "grasp": 1,
    "cut": 2,
    "scoop": 3,
    "contain": 4,
    "pound": 5,
    "support": 6,
    "wrap grasp": 7,
}

IDX_TO_LABEL = ["background", "grasp", "cut", "scoop", "contain", "pound", "support", "wrap grasp"]

transform = T.Compose([T.Resize((SIZE, SIZE)), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
PATCH_SIZE = 14

DIM = 1024

LABEL_TO_IDX = {
    "background": 0,
    "grasp": 1,
    "cut": 2,
    "scoop": 3,
    "contain": 4,
    "pound": 5,
    "support": 6,
    "wrap grasp": 7,
}

IDX_TO_LABEL = ["background", "grasp", "cut", "scoop", "contain", "pound", "support", "wrap grasp"]



feature_extractor = (
    torch.hub.load("facebookresearch/dinov2:main", "dinov2_vitl14_reg")
    .to(device)
    .eval()
)

def torch_to_cv(tensor):
    img_cv = tensor.detach().cpu().numpy()[..., ::-1]
    img_cv = np.clip(img_cv * 255, 0, 255).astype(np.uint8)
    return img_cv


def _detach_tensors_from_dict(d, inplace=True):
    if not inplace:
        d = d.copy()
    for key in d:
        if isinstance(d[key], torch.Tensor):
            d[key] = d[key].detach()
    return d


def load_checkpoint(
    checkpoint: str, data_dir: str, rasterizer: str = "original", data_factor: int = 1
):

    colmap_project = pycolmap.SceneManager(f"{data_dir}/sparse/0")
    colmap_project.load_cameras()
    colmap_project.load_images()
    colmap_project.load_points3D()
    model = torch.load(checkpoint)  # Make sure it is generated by 3DGS original repo
    if rasterizer == "original":
        model_params, _ = model
        splats = {
            "active_sh_degree": model_params[0],
            "means": model_params[1],
            "features_dc": model_params[2],
            "features_rest": model_params[3],
            "scaling": model_params[4],
            "rotation": model_params[5],
            "opacity": model_params[6].squeeze(1),
        }
    elif rasterizer == "gsplat":

        model_params = model["splats"]
        splats = {
            "active_sh_degree": 3,
            "means": model_params["means"],
            "features_dc": model_params["sh0"],
            "features_rest": model_params["shN"],
            "scaling": model_params["scales"],
            "rotation": model_params["quats"],
            "opacity": model_params["opacities"],
        }
    else:
        raise ValueError("Invalid rasterizer")

    _detach_tensors_from_dict(splats)

    # Assuming only one camera
    for camera in colmap_project.cameras.values():
        camera_matrix = torch.tensor(
            [
                [camera.fx, 0, camera.cx],
                [0, camera.fy, camera.cy],
                [0, 0, 1],
            ]
        )
        break

    camera_matrix[:2, :3] /= data_factor

    splats["camera_matrix"] = camera_matrix
    splats["colmap_project"] = colmap_project
    splats["colmap_dir"] = data_dir

    return splats


def get_viewmat_from_colmap_image(image):
    viewmat = torch.eye(4).float()  # .to(device)
    viewmat[:3, :3] = torch.tensor(image.R()).float()  # .to(device)
    viewmat[:3, 3] = torch.tensor(image.t).float()  # .to(device)
    return viewmat


def create_checkerboard(width, height, size=64):
    checkerboard = np.zeros((height, width, 3), dtype=np.uint8)
    for y in range(0, height, size):
        for x in range(0, width, size):
            if (x // size + y // size) % 2 == 0:
                checkerboard[y : y + size, x : x + size] = 255
            else:
                checkerboard[y : y + size, x : x + size] = 128
    return checkerboard


def render_to_dir(output_dir: str, splats, feedback: bool = False):
    if feedback:
        cv2.destroyAllWindows()
        cv2.namedWindow("Initial Rendering", cv2.WINDOW_NORMAL)
    os.makedirs(output_dir, exist_ok=True)
    colmap_project = splats["colmap_project"]
    frame_idx = 0
    for image in sorted(colmap_project.images.values(), key=lambda x: x.name):
        image_name = image.name  # .split(".")[0] + ".jpg"
        image_path = f"{output_dir}/{image_name}"
        means = splats["means"]
        colors_dc = splats["features_dc"]
        colors_rest = splats["features_rest"]
        colors = torch.cat([colors_dc, colors_rest], dim=1)
        opacities = torch.sigmoid(splats["opacity"])
        scales = torch.exp(splats["scaling"])
        quats = splats["rotation"]
        viewmat = get_viewmat_from_colmap_image(image)
        K = splats["camera_matrix"]
        output, _, info = rasterization(
            means,
            quats,
            scales,
            opacities,
            colors,
            viewmats=viewmat[None],
            Ks=K[None],
            sh_degree=3,
            width=K[0, 2] * 2,
            height=K[1, 2] * 2,
        )
        output_cv = torch_to_cv(output[0])
        imageio.imsave(image_path, output_cv[:, :, ::-1])
        if feedback:
            cv2.imshow("Initial Rendering", output_cv)
            cv2.waitKey(1)
        frame_idx += 1






def load_labels(labels_dir, results_dir):
    example_images = os.listdir(labels_dir)
    example_images = [img for img in example_images if img.endswith(".webp")]
    features = []
    labels = []
    for example in example_images:
        example_path = os.path.join(labels_dir, example)
        annotations = json.load(open(example_path.replace(".webp", ".json")))

        img_example = cv2.imread(example_path)[..., ::-1]
        img_example = Image.fromarray(img_example)
        img_example_th = transform(img_example).to(device)
        feats = feature_extractor.forward_features(img_example_th[None])[
            "x_norm_patchtokens"
        ][0]
        feats = feats.reshape((FEATURE_MAP_SIZE, FEATURE_MAP_SIZE, -1))
        feats = feats / torch.norm(feats, dim=-1, keepdim=True)
        feats_flatten = feats.detach().cpu().numpy().reshape((-1, DIM))

        background_mask_inverted = np.zeros((FEATURE_MAP_SIZE, FEATURE_MAP_SIZE))

        for shape in annotations["shapes"]:
            shape_label = shape["label"]
            assert shape_label in LABEL_TO_IDX
            label_idx = LABEL_TO_IDX[shape_label]
            mask64 = shape["mask"]
            mask_bytes = base64.b64decode(mask64)
            mask = (
                cv2.imdecode(np.frombuffer(mask_bytes, np.uint8), cv2.IMREAD_UNCHANGED)
                * 255
            )
            boundary_points = np.array(shape["points"]).astype(np.int32)
            blank_image = np.zeros_like(img_example)
            blank_image[
                boundary_points[0][1] : boundary_points[1][1] + 1,
                boundary_points[0][0] : boundary_points[1][0] + 1,
            ] = mask[..., None]
            blank_image = blank_image[..., 0]
            mask_np = cv2.resize(
                blank_image,
                (FEATURE_MAP_SIZE, FEATURE_MAP_SIZE),
                interpolation=cv2.INTER_NEAREST,
            )
            mask_flatten = mask_np.flatten()
            feats_current_mask = feats_flatten[mask_flatten > 0]
            labels_current_mask = np.ones((feats_current_mask.shape[0], 1)) * label_idx

            features = (
                np.concatenate([features, feats_current_mask], axis=0)
                if len(features) > 0
                else feats_current_mask
            )
            labels = (
                np.concatenate([labels, labels_current_mask], axis=0)
                if len(labels) > 0
                else labels_current_mask
            )
            background_mask_inverted = np.logical_or(background_mask_inverted, mask_np)
            cv2.imshow("mask", mask_np)
            cv2.waitKey(1)
            cv2.destroyAllWindows()

        background_mask = ~background_mask_inverted
        features_background = feats_flatten[background_mask.flatten()]
        labels_background = np.zeros((features_background.shape[0], 1))
        features = (
            np.concatenate([features, features_background], axis=0)
            if len(features) > 0
            else features_background
        )
        labels = (
            np.concatenate([labels, labels_background], axis=0)
            if len(labels) > 0
            else labels_background
        )
        cv2.imshow("background_mask", background_mask.astype(np.uint8) * 255)
        cv2.waitKey(1)
        cv2.destroyAllWindows()
        data = {"features": features, "labels": labels}
        os.makedirs(results_dir, exist_ok=True)
        pkl.dump(data, open(os.path.join(results_dir, "features_and_labels.pkl"), "wb"))


def render_images(
    data_dir: str,
    checkpoint: str,
    results_dir: str,
):
    output_dir = os.path.join(results_dir, "images")
    splats = load_checkpoint(checkpoint, data_dir, rasterizer="gsplat")
    render_to_dir(output_dir, splats, feedback=True)

def render_and_vote(data_dir, checkpoint, results_dir):
    splats = load_checkpoint(checkpoint, data_dir, rasterizer="gsplat")
    means = splats["means"]
    colors_dc = splats["features_dc"]
    colors_rest = splats["features_rest"]
    colors = torch.cat([colors_dc, colors_rest], dim=1)
    opacities = torch.sigmoid(splats["opacity"])
    scales = torch.exp(splats["scaling"])
    quats = splats["rotation"]
    K = splats["camera_matrix"]
    width = int(K[0, 2] * 2)
    height = int(K[1, 2] * 2)
    colmap_project = splats["colmap_project"]

    affordance_dir = os.path.join(results_dir, "affordance_maps")
    affordance_map_images_dir = os.path.join(results_dir, "affordance_map_images")
    affordance_maps_3d_dir = os.path.join(results_dir, "affordance_map_images_3dgs")
    os.makedirs(affordance_maps_3d_dir, exist_ok=True)




    votes = torch.zeros((8, colors.shape[0])).to(device) # Hardcoding for now

    colors.requires_grad = True

    for image in sorted(colmap_project.images.values(), key=lambda x: x.name):
        viewmat = get_viewmat_from_colmap_image(image)

        output, _, meta = rasterization(
            means,
            quats,
            scales,
            opacities,
            colors,
            viewmat[None],
            K[None],
            width=width,
            height=height,
            sh_degree=3,
            # render_mode="RGB+ED",
            backgrounds=torch.tensor([[0.0, 0.0, 0.0]]).to(device),
        )

        label_map_name = image.name.split(".")[0] + ".npy"
        label_map = np.load(os.path.join(affordance_dir, label_map_name))
        

        mask = np.zeros((height, width, 3), dtype=np.uint8)
        label_map = cv2.resize(
            label_map, (width, height), interpolation=cv2.INTER_NEAREST
        )
        

        # Voting
        for i in range(8):
            mask[label_map == i, :] = i
        

        mask = torch.from_numpy(mask).to(device).float()


        for i in range(8):
            loss = ((mask == i) * output[0]).mean()
            loss.backward(retain_graph=True)
            votes[i] += colors.grad[..., 0, :3].norm(dim=[1])
            colors.grad.zero_()
        
    votes_path = os.path.join(results_dir, "votes.npy")
    votes_np = votes.cpu().numpy()
    np.save(votes_path, votes_np)

    colors_new = colors.clone()
    color_palette = np.array([[125, 125, 125], [255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 255, 0], [255, 0, 255], [0, 255, 255], [128, 0, 0]])

    for i in range(8):
        mask = torch.argmax(votes, dim=0) == i
        colors_new[mask, 0, :3] = colors_new[mask, 0, :3] * 0.5 + 0.5*((torch.from_numpy(color_palette[i])/255.).to(device)-0.5)/ (1/np.sqrt(4*np.pi))
        colors_new[:, 1:, :3] = 0.1 * colors_new[:, 1:, :3]

    for image in sorted(colmap_project.images.values(), key=lambda x: x.name):
        viewmat = get_viewmat_from_colmap_image(image)

        output, _, meta = rasterization(
            means,
            quats,
            scales,
            opacities,
            colors_new,
            viewmat[None],
            K[None],
            width=width,
            height=height,
            sh_degree=3,
            backgrounds=torch.tensor([[0.0, 0.0, 0.0]]).to(device),
        )

        output_cv = torch_to_cv(output[0, ..., :3].detach())
        cv2.imshow("Mapped affordance regions", output_cv)
        cv2.imwrite(f"{affordance_maps_3d_dir}/{image.name}", output_cv)
        cv2.waitKey(1)


def calculate_affordance_map(data_dir, labels_dir, results_dir, k=5):
    def most_frequent(array):
        return np.bincount(array).argmax()
    images_dir = os.path.join(results_dir, "images")
    affordance_map_dir = os.path.join(results_dir, "affordance_maps")
    affordance_map_images_dir = os.path.join(results_dir, "affordance_map_images")
    os.makedirs(affordance_map_images_dir, exist_ok=True)
    os.makedirs(affordance_map_dir, exist_ok=True)
    features_and_labels = pkl.load(open(os.path.join(results_dir, "features_and_labels.pkl"), "rb"))

    features = features_and_labels["features"]
    labels = features_and_labels["labels"]
    feature_index = faiss.IndexFlatIP(DIM)
    feature_index.add(features)
    for image_name in sorted(os.listdir(images_dir)):
        image_path = os.path.join(images_dir, image_name)
        image = cv2.imread(image_path)[..., ::-1]
        image = Image.fromarray(image)
        image_th = transform(image).to(device)
        feats = feature_extractor.forward_features(image_th[None])["x_norm_patchtokens"][0]
        feats = feats.reshape((FEATURE_MAP_SIZE, FEATURE_MAP_SIZE, -1))
        feats = feats / torch.norm(feats, dim=-1, keepdim=True)
        feats_flatten = feats.detach().cpu().numpy().reshape((-1, DIM))
        D, I = feature_index.search(feats_flatten, k)
        label_set = labels[I].astype(np.uint8)
        label_set = np.apply_along_axis(most_frequent, axis=1, arr=label_set)
        affordance_map = label_set.reshape((FEATURE_MAP_SIZE, FEATURE_MAP_SIZE)).astype(np.uint8)
        # affordance_map = labels[I[:, 0]].reshape((FEATURE_MAP_SIZE, FEATURE_MAP_SIZE)).astype(np.uint8)

        # Showing feedback
        color_palette = np.array([[0, 0, 0], [0, 0, 255], [0, 255, 0], [255, 0, 0], [0, 255, 255], [255, 0, 255], [255, 255, 0], [0, 0, 128]])

        labels_mask = color_palette[affordance_map.flatten()]
        labels_mask = labels_mask.reshape((FEATURE_MAP_SIZE, FEATURE_MAP_SIZE, 3))
        np.save(os.path.join(affordance_map_dir, image_name.replace(".jpg", ".npy")), affordance_map)
        labels_mask_big = cv2.resize(labels_mask, image.size, interpolation=cv2.INTER_NEAREST)
        img = cv2.imread(image_path)

        img_out = img * 0.5 + 0.5 * labels_mask_big
        img_out = img_out.astype(np.uint8)

        for i in range(8):
            color = color_palette[i].tolist()
            cv2.putText(img_out, IDX_TO_LABEL[i], (10, 10 + 20*i), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)

        cv2.imshow("2D-2D Affordance Transfer", img_out)
        cv2.imwrite(os.path.join(affordance_map_images_dir, image_name), img_out)
        cv2.waitKey(1)

    cv2.destroyAllWindows()




from scipy.io import loadmat

def evaluate_results(data_dir, checkpoint, results_dir):
    gt_dir = os.path.join(data_dir, "gt")
    affordance_maps_dir = os.path.join(results_dir, "affordance_maps")
    
    # Not rendering now
    splats = load_checkpoint(checkpoint, data_dir, rasterizer="gsplat")
    colmap_project = splats["colmap_project"]
    label_files = os.listdir(gt_dir)
    label_files = [label_file for label_file in label_files if label_file.endswith("label.mat")]
    label_files.sort()
    mIoU = defaultdict(list)
    recall = defaultdict(list)
    assert len(label_files) == len(colmap_project.images)

    # For affordance map

    print("\n\nEvaluating 2D-2D affordance transfer.")
    for image, label_file in zip(sorted(colmap_project.images.values(), key=lambda x: x.name), label_files):
        # gt_name = image.name.replace("_rgb.jpg", "_label.mat")
        # image_path = os.path.join(data_dir, "images", image.name)
        # gt_image_path = os.path.join(gt_dir, label_file.replace("_label.mat", "_rgb.jpg"))
        gt_path = os.path.join(gt_dir, label_file)
        gt = loadmat(gt_path)
        if gt["gt_type"] == "automatic":
            continue
        gt_label = gt["gt_label"]

        affordance_map_name = image.name.replace(".jpg", ".npy")
        affordance_map_path = os.path.join(affordance_maps_dir, affordance_map_name)
        affordance_map = np.load(affordance_map_path)
        affordance_map = cv2.resize(affordance_map, (gt_label.shape[1], gt_label.shape[0]), interpolation=cv2.INTER_NEAREST)

        for i in range(1, 8):
            gt_mask = gt_label == i
            affordance_mask = affordance_map == i
            intersection = np.logical_and(gt_mask, affordance_mask).sum()
            union = np.logical_or(gt_mask, affordance_mask).sum()
            if union == 0:
                continue
            if intersection == 0:
                iou = 0
            else:
                iou = intersection / union
            mIoU[i].append(iou)

            if gt_mask.sum() == 0:
                continue
            if intersection == 0:
                rec = 0
            else:
                rec = intersection / gt_mask.sum() # Actually recall
            recall[i].append(rec)
            # print(f"Class {i}: {iou}")
        # Calculate mIoU
    res = 0
    for i in range(1, 8):
        if len(mIoU[i]) == 0:
            continue
        res += np.mean(mIoU[i])
        # print(f"mIoU {i}: {np.mean(mIoU[i])}")
    res /= (np.unique(gt_label).size - 1)
    print(f"mIoU: {res}")

    res_recall = 0
    for i in range(1, 8):
        if len(recall[i]) == 0:
            continue
        res_recall += np.mean(recall[i])
        # print(f"Recall {i}: {np.mean(recall[i])}")
    res_recall /= (np.unique(gt_label).size - 1)
    print(f"Recall: {res_recall}")


    print("\nEvaluating 2D-3D affordance transfer.")
    mIoU = defaultdict(list)
    recall = defaultdict(list)
    means = splats["means"]
    colors_dc = splats["features_dc"]
    colors_rest = splats["features_rest"]
    colors = torch.cat([colors_dc, colors_rest], dim=1)
    opacities = torch.sigmoid(splats["opacity"])
    scales = torch.exp(splats["scaling"])
    quats = splats["rotation"]
    K = splats["camera_matrix"]
    width = int(K[0, 2] * 2)
    height = int(K[1, 2] * 2)

    votes_path = os.path.join(results_dir, "votes.npy")
    votes = np.load(votes_path)
    votes = torch.from_numpy(votes).to(device)
    for image, label_file in zip(sorted(colmap_project.images.values(), key=lambda x: x.name), label_files):
        gt_name = image.name.replace("_rgb.jpg", "_label.mat")
        image_path = os.path.join(data_dir, "images", image.name)
        gt_image_path = os.path.join(gt_dir, label_file.replace("_label.mat", "_rgb.jpg"))

        gt_path = os.path.join(gt_dir, label_file)
        gt = loadmat(gt_path)
        if gt["gt_type"] == "automatic":
            continue
        gt_label = gt["gt_label"]

        affordance_map_name = image.name.replace(".jpg", ".npy")
        affordance_map_path = os.path.join(affordance_maps_dir, affordance_map_name)
        affordance_map = np.load(affordance_map_path)
        affordance_map = cv2.resize(affordance_map, (gt_label.shape[1], gt_label.shape[0]), interpolation=cv2.INTER_NEAREST)

        viewmat = get_viewmat_from_colmap_image(image)


        for i in range(1, 8):
            colors_new = colors.clone()[:,0]
            mask = torch.argmax(votes, dim=0) == i
            colors_new[mask] = 1.0
            colors_new[~mask] = 0.0

            output, _, meta = rasterization(
                means,
                quats,
                scales,
                opacities,
                colors_new,
                viewmat[None],
                K[None],
                width=width,
                height=height,
                backgrounds=torch.tensor([[0.0, 0.0, 0.0]]).to(device),
            )
            output_cv = torch_to_cv(output[0, ..., :3].detach())
            cv2.imshow("Mapped affordance regions", output_cv)
            cv2.waitKey(1)
            gt_mask = gt_label == i
            affordance_mask = output_cv > 64
            affordance_mask = affordance_mask[...,0]
            intersection = np.logical_and(gt_mask, affordance_mask).sum()
            union = np.logical_or(gt_mask, affordance_mask).sum()
            if union == 0:
                continue
            if intersection == 0:
                iou = 0
            else:
                iou = intersection / union
            mIoU[i].append(iou)
            if gt_mask.sum() == 0:
                continue
            if intersection == 0:
                rec = 0
            else:
                rec = intersection / gt_mask.sum()
            recall[i].append(rec)
    res = 0
    for i in range(1, 8):
        if len(mIoU[i]) == 0:
            continue
        res += np.mean(mIoU[i])
        # print(f"mIoU {i}: {np.mean(mIoU[i])}")
    res /= (np.unique(gt_label).size - 1)
    print(f"mIoU: {res}")

    res_recall = 0
    for i in range(1, 8):
        if len(recall[i]) == 0:
            continue
        res_recall += np.mean(recall[i])
        # print(f"Recall {i}: {np.mean(recall[i])}")
    res_recall /= (np.unique(gt_label).size - 1)
    print(f"Recall: {res_recall}")


def main(
    data_dir: str = "./data/processed_scene_01",
    checkpoint: str = "./data/processed_scene_01/ckpts/ckpt_29999_rank0.pt",
    labels_dir: str = "./data/affordance_labels",
    results_dir: str = "./results/scene_01",
):
    # Load labels to database
    load_labels(labels_dir, results_dir)

    # Render images to the results folder
    render_images(data_dir, checkpoint, results_dir)

    # Calculate affordance map and save
    calculate_affordance_map(data_dir, labels_dir, results_dir)

    # Render and vote
    render_and_vote(data_dir, checkpoint, results_dir)

    # Evaluate the results. Comment out if you don't want to evaluate
    evaluate_results(data_dir, checkpoint, results_dir)


if __name__ == "__main__":
    tyro.cli(main)