inference_cam.py

#!/usr/bin/env python
""" Inference script

Hacked together by Ross Wightman (https://github.com/rwightman)
"""
import argparse
import json
import os
from contextlib import suppress
from pathlib import Path

import numpy as np
import open3d as o3d
import timm.utils as utils
import torch
import torch.nn.parallel
import tqdm
from pytorch_grad_cam import AblationCAM, EigenCAM, GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
from sklearn.neighbors import NearestNeighbors
from timm.utils import setup_default_logging

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = ""  # USE CPU only

import cv2  # noqa: E402

import dataset.transformations as transformations
from effdet import create_dataset, create_loader, create_model, rotation6d
from effdet.data import resolve_input_config

try:
    from timm.layers import set_layer_config
except ImportError:
    from timm.models.layers import set_layer_config

has_apex = False
try:
    from apex import amp

    has_apex = True
except ImportError:
    pass

has_native_amp = False
try:
    if getattr(torch.cuda.amp, "autocast") is not None:
        has_native_amp = True
except AttributeError:
    pass

torch.backends.cudnn.benchmark = True


def add_bool_arg(parser, name, default=False, help=""):  # FIXME move to utils
    dest_name = name.replace("-", "_")
    group = parser.add_mutually_exclusive_group(required=False)
    group.add_argument("--" + name, dest=dest_name, action="store_true", help=help)
    group.add_argument("--no-" + name, dest=dest_name, action="store_false", help=help)
    parser.set_defaults(**{dest_name: default})


parser = argparse.ArgumentParser(description="PyTorch ImageNet Validation")
parser.add_argument("root", metavar="DIR", help="path to dataset root")
parser.add_argument(
    "--dataset",
    default="coco",
    type=str,
    metavar="DATASET",
    help='Name of dataset (default: "coco"',
)
parser.add_argument("--split", default="val", help="validation split")
parser.add_argument(
    "--model",
    "-m",
    metavar="MODEL",
    default="tf_efficientdet_d1",
    help="model architecture (default: tf_efficientdet_d1)",
)
add_bool_arg(
    parser,
    "redundant-bias",
    default=None,
    help="override model config for redundant bias layers",
)
add_bool_arg(
    parser, "soft-nms", default=None, help="override model config for soft-nms"
)
add_bool_arg(
    parser, "visualize", default=False, help="Visualize detections and ground truth"
)
add_bool_arg(parser, "refine", default=True, help="ICP Refinement of predictions")
add_bool_arg(parser, "export", default=False, help="Export predictions to JSON file")
parser.add_argument(
    "--num-classes",
    type=int,
    default=None,
    metavar="N",
    help="Override num_classes in model config if set. For fine-tuning from pretrained.",
)
parser.add_argument(
    "-j",
    "--workers",
    default=4,
    type=int,
    metavar="N",
    help="number of data loading workers (default: 4)",
)
add_bool_arg(parser, "persistent-workers", default=False, help="Keep workers alive.")
parser.add_argument(
    "-b",
    "--batch-size",
    default=1,
    type=int,
    metavar="N",
    help="mini-batch size (default: 128)",
)
parser.add_argument(
    "--img-size",
    default=None,
    type=int,
    metavar="N",
    help="Input image dimension, uses model default if empty",
)
parser.add_argument(
    "--mean",
    type=float,
    nargs="+",
    default=None,
    metavar="MEAN",
    help="Override mean pixel value of dataset",
)
parser.add_argument(
    "--std",
    type=float,
    nargs="+",
    default=None,
    metavar="STD",
    help="Override std deviation of of dataset",
)
parser.add_argument(
    "--interpolation",
    default="bilinear",
    type=str,
    metavar="NAME",
    help="Image resize interpolation type (overrides model)",
)
parser.add_argument(
    "--fill-color",
    default=None,
    type=str,
    metavar="NAME",
    help='Image augmentation fill (background) color ("mean" or int)',
)
parser.add_argument(
    "--log-freq",
    default=10,
    type=int,
    metavar="N",
    help="batch logging frequency (default: 10)",
)
parser.add_argument(
    "--checkpoint",
    default="",
    type=str,
    metavar="PATH",
    help="path to latest checkpoint (default: none)",
)
parser.add_argument(
    "--pretrained", dest="pretrained", action="store_true", help="use pre-trained model"
)
parser.add_argument("--num-gpu", type=int, default=1, help="Number of GPUS to use")
parser.add_argument(
    "--no-prefetcher",
    action="store_true",
    default=False,
    help="disable fast prefetcher",
)
parser.add_argument(
    "--pin-mem",
    action="store_true",
    default=False,
    help="Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.",
)
parser.add_argument(
    "--use-ema",
    dest="use_ema",
    action="store_true",
    help="use ema version of weights if present",
)
parser.add_argument(
    "--amp",
    action="store_true",
    default=False,
    help="Use AMP mixed precision. Defaults to Apex, fallback to native Torch AMP.",
)
parser.add_argument(
    "--apex-amp",
    action="store_true",
    default=False,
    help="Use NVIDIA Apex AMP mixed precision",
)
parser.add_argument(
    "--native-amp",
    action="store_true",
    default=False,
    help="Use Native Torch AMP mixed precision",
)
parser.add_argument(
    "--torchscript",
    dest="torchscript",
    action="store_true",
    help="convert model torchscript for inference",
)
parser.add_argument(
    "--torchcompile",
    nargs="?",
    type=str,
    default=None,
    const="inductor",
    help="Enable compilation w/ specified backend (default: inductor).",
)
parser.add_argument(
    "--results",
    default="",
    type=str,
    metavar="FILENAME",
    help="JSON filename for evaluation results",
)
parser.add_argument(
    "--z-offset",
    default=0.0,
    type=float,
    metavar="Z_OFFSET",
    help="Offset the z-predictions in mm (ImagePose only!).",
)
parser.add_argument(
    "--min-score",
    default=0.9,
    type=float,
    help="Minimum score for predictions to be considered valid.",
)
parser.add_argument(
    "--topk", default=0, type=int, help="Only use the top k predictions per image."
)
parser.add_argument(
    "--z-aug", default=0, type=float, help="Sets the z-augmentation factor."
)

def inference(args):
    setup_default_logging()

    if args.amp:
        if has_native_amp:
            args.native_amp = True
        elif has_apex:
            args.apex_amp = True
    assert not args.apex_amp or not args.native_amp, "Only one AMP mode should be set."
    args.pretrained = (
        args.pretrained or not args.checkpoint
    )  # might as well try to validate something
    args.prefetcher = not args.no_prefetcher

    # create model
    with set_layer_config(scriptable=args.torchscript):
        extra_args = {}
        if args.img_size is not None:
            extra_args = dict(image_size=(args.img_size, args.img_size))
        bench = create_model(
            args.model,
            bench_task="predict",
            num_classes=args.num_classes,
            pretrained=args.pretrained,
            redundant_bias=args.redundant_bias,
            soft_nms=args.soft_nms,
            checkpoint_path=args.checkpoint,
            checkpoint_ema=args.use_ema,
            **extra_args,
        )
    model_config = bench.config

    param_count = sum([m.numel() for m in bench.parameters()])
    print("Model %s created, param count: %d" % (args.model, param_count))

    bench = bench.cuda()

    if args.torchscript:
        assert (
            not args.apex_amp
        ), "Cannot use APEX AMP with torchscripted model, force native amp with `--native-amp` flag"
        bench = torch.jit.script(bench)
    elif args.torchcompile:
        bench = torch.compile(bench, backend=args.torchcompile)

    amp_autocast = suppress
    if args.apex_amp:
        bench = amp.initialize(bench, opt_level="O1")
        print("Using NVIDIA APEX AMP. Validating in mixed precision.")
    elif args.native_amp:
        amp_autocast = torch.cuda.amp.autocast
        print("Using native Torch AMP. Validating in mixed precision.")
    else:
        print("AMP not enabled. Validating in float32.")

    if args.num_gpu > 1:
        bench = torch.nn.DataParallel(bench, device_ids=list(range(args.num_gpu)))

    dataset = create_dataset(args.dataset, args.root, args.split)
    input_config = resolve_input_config(args, model_config)

    loader = create_loader(
        dataset,
        input_size=input_config["input_size"],
        batch_size=args.batch_size,
        use_prefetcher=args.prefetcher,
        interpolation=input_config["interpolation"],
        fill_color=input_config["fill_color"],
        mean=input_config["mean"],
        std=input_config["std"],
        num_workers=args.workers,
        persistent_workers=args.persistent_workers,
        pin_mem=args.pin_mem,
    )

    has_groundtruth = dataset.parser.has_labels
    score_threshold = args.min_score

    if dataset.NAME == "imagepose":
        z_offset = args.z_offset
    elif dataset.NAME == "depthpose":
        z_offset = dataset.MEAN[0] - args.mean[0] if args.mean else 0
        z_offset *= 1000  # convert to mm.

    bench.eval()
    cam = EigenCAM(bench, target_layers=[bench.model.backbone.blocks[-1][-2]], use_cuda=False)
    # cam = EigenCAM(bench, target_layers=[bench.model.fpn.cell[-1].fnode[-1]], use_cuda=False)
    # cam = EigenCAM(bench, target_layers=[
        # list(bench.model.translation_z_net.children())[-1],
        # list(bench.model.translation_xy_net.children())[-1],
        # list(bench.model.rotation_net.children())[-1],
    # ])
    # cam = EigenCAM(bench, target_layers=[
    #     list(bench.model.fpn.modules())[-3]
    # ])
    with torch.inference_mode():
        tqdm_loader = tqdm.tqdm(
            loader,
            unit_scale=args.batch_size
        )

        for input, target in tqdm_loader:
            with amp_autocast():
                output = bench(input)[0]

            batch_target = {key: value[0] for key, value in target.items()}
            targets = postprocess_targets(batch_target, has_groundtruth)
            ortho_scale = targets["ortho_scale"]
            K = targets["K"]
            img_size = targets["img_size"]
            img_id = targets["img_idx"]

            T_blender_world = transformations.get_T_blender_world(ortho_scale)
            T_world_cam = transformations.get_T_world_cam(
                targets["cam_R"], targets["cam_t"]
            )

            predictions = postprocess_predictions(
                output,
                score_threshold,
                dataset,
                ortho_scale,
                img_size,
                T_blender_world,
                T_world_cam,
                K,
                z_offset,
                top_k=args.topk,
            )

            normalized_img = input.squeeze().permute(1, 2, 0).cpu().numpy()
            mean = loader.dataset.MEAN
            std = loader.dataset.STD
            img = np.clip(normalized_img * std + mean, 0, 1)
            img = cv2.cvtColor(img.astype(np.float32), cv2.COLOR_RGB2GRAY)
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)  # 3 channel gray image

            cam_out = cam(input_tensor=input)[0]
            cam_img = show_cam_on_image(img, cam_out)
            cv2.imshow("cam", cam_img)
            # cv2.waitKey(0)

            renormalized_cam_image = renormalize_cam_in_bounding_boxes(
                img, cam_out, predictions["bbox"])
            cv2.imshow("cam_normalized", renormalized_cam_image)
            cv2.waitKey(0)

        print("Done!")


def draw_detections(img, boxes, colors=None, names=None):
    if colors is None:
        colors = [(0, 255, 0) for _ in boxes]
    if names is None:
        names = ["fleckenzwerg" for _ in boxes]

    for box, color, name in zip(boxes.astype(np.int16), colors, names):
        xmin, ymin, xmax, ymax = box
        cv2.rectangle(
            img,
            (xmin, ymin),
            (xmax, ymax),
            color,
            2)

        cv2.putText(img, name, (xmin, ymin - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2,
                    lineType=cv2.LINE_AA)
    return img


def renormalize_cam_in_bounding_boxes(image_float_np, grayscale_cam, boxes, colors=None, names=None, use_rgb=False):
    """Normalize the CAM to be in the range [0, 1]
    inside every bounding boxes, and zero outside of the bounding boxes. """
    renormalized_cam = np.zeros(grayscale_cam.shape, dtype=np.float32)
    for x1, y1, x2, y2 in np.clip(boxes.astype(np.int16), 0, 639):
        renormalized_cam[y1:y2, x1:x2] = scale_cam_image(grayscale_cam[y1:y2, x1:x2].copy())
    renormalized_cam = scale_cam_image(renormalized_cam)
    eigencam_image_renormalized = show_cam_on_image(image_float_np, renormalized_cam, use_rgb=use_rgb)
    image_with_bounding_boxes = draw_detections(eigencam_image_renormalized, boxes, colors, names)
    return image_with_bounding_boxes


def postprocess_predictions(
    output,
    threshold,
    dataset,
    ortho_scale,
    img_size,
    T_blender_world,
    T_world_cam,
    K,
    z_offset,
    top_k: int = 0,
):
    # output format: [y, x, y, x, score, class, r1, r2, r3, r4, r5, r6, tx, ty, tz]

    # Filter by score.
    scores = output[:, 4].cpu().numpy()
    if top_k > 0:
        # Scores were already sorted, so just use the top k.
        solid_detections = np.arange(min(top_k, len(scores)))
    else:
        solid_detections = scores > threshold

    scores = scores[solid_detections]

    # Extract predictions.
    bbox_preds = output[solid_detections, :4].cpu().numpy()  # 2D bounding boxes
    bbox_preds = bbox_preds[:, [1, 0, 3, 2]]  # yxyx -> xyxy
    class_preds = output[solid_detections, 5].cpu().numpy()
    rotation6d_preds = output[solid_detections, 6:12]
    rotation_preds = (
        rotation6d.compute_rotation_matrix_from_ortho6d(rotation6d_preds).cpu().numpy()
    )
    centerpoint_preds = output[solid_detections, 12:14].cpu().numpy()
    centerpoint_preds = centerpoint_preds[:, [1, 0]]  # yx -> xy
    translation_z_preds = output[solid_detections, 14].cpu().numpy()

    if dataset.NAME == "depthpose":
        translation_z_preds -= z_offset
        translation_preds = transformations.orthographic_unproject(
            centerpoint_preds,
            translation_z_preds,
            img_size,
            ortho_scale,
        )
    elif dataset.NAME == "imagepose":
        translation_z_preds += z_offset
        translation_preds = transformations.perspective_unproject(
            centerpoint_preds,
            translation_z_preds,
            T_blender_world,
            T_world_cam,
            K,
        )

    return {
        "scores": scores,
        "bbox": bbox_preds,
        "rotation": rotation_preds,  # Rotation Matrix
        "centerpoint": centerpoint_preds,
        "translation": translation_preds,
    }


def postprocess_targets(target, has_groundtruth=True):
    targets = {
        "img_idx": target["img_idx"].cpu().numpy(),
        "img_size": target["img_size"][0].cpu().numpy(),
        "ortho_scale": target["ortho_scale"].cpu().numpy(),
        "cam_R": target["cam_R"].cpu().numpy(),
        "cam_t": target["cam_t"].cpu().numpy(),
        "K": target["K"].cpu().numpy().reshape(3, 3),
    }

    if has_groundtruth:
        centerpoint_targets = target["centerpoint"].cpu().numpy()
        solid_targets = centerpoint_targets[:, 0] != -1
        centerpoint_targets = centerpoint_targets[solid_targets]
        bbox_targets = target["bbox"][solid_targets].cpu().numpy()
        rotation_targets = target["rotation"][solid_targets].cpu().numpy()
        rotation_targets = rotation_targets.reshape(-1, 3, 3)
        rotation_targets_from6d = rotation6d.compute_rotation_matrix_from_ortho6d(
            target["rotation6d"][solid_targets]
        )
        rotation_targets_from6d = rotation_targets_from6d.cpu().numpy()

        # Quick sanity check to see whether the conversion is working correctly.
        assert np.allclose(
            rotation_targets_from6d, rotation_targets
        ), "6D Rotation conversion failed."

        translation_targets = target["translation"][solid_targets].cpu().numpy()
        targets["bbox"] = bbox_targets
        targets["rotation"] = rotation_targets_from6d
        targets["centerpoint"] = centerpoint_targets
        targets["translation"] = translation_targets

    return targets


def main():
    args = parser.parse_args()
    inference(args)


if __name__ == "__main__":
    main()