<gh_stars>0
import argparse
import os

import cv2
import numpy as np
import torch
from torch.nn.functional import cosine_similarity, interpolate
from torchvision import transforms
from torchvision.transforms.functional import crop
from torchvision.utils import make_grid, save_image

from solo.utils.backbones import resnet18, resnet50


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--backbone",
        "-b",
        type=str,
        default="resnet18",
        choices=["resnet18", "resnet50"],
        help="backbone architecture",
    )
    parser.add_argument(
        "--weights", "-w", type=str, required=True, help="path to pretrained weights"
    )
    parser.add_argument(
        "--image", "-i", type=str, required=True, help="path to visualization image"
    )
    parser.add_argument(
        "--coord",
        "-c",
        nargs="+",
        type=int,
        required=True,
        help="x and y coord of point (with [0,0] at top left)",
    )
    parser.add_argument(
        "--output",
        "-o",
        type=str,
        default="visualize_output",
        help="path of output directory",
    )
    parser.add_argument(
        "--suffix", "-s", type=str, default="", help="output file suffix"
    )

    args = parser.parse_args()

    return args


def load_backbone(weights_path: str, arch: str):
    # Load state dict
    weights = torch.load(weights_path)["state_dict"]

    # Fix layer names
    newmodel = {}
    for k, v in weights.items():
        if not k.startswith("backbone"):
            continue

        k = k.replace("backbone.", "")
        newmodel[k] = v

    # Load model
    if arch == "resnet18":
        backbone = resnet18()
    else:
        backbone = resnet50()

    backbone.load_state_dict(newmodel, strict=False)
    return backbone


def draw_dot(img: torch.Tensor, coord: list[int]):
    pt = [round(coord[0] * (224 / 56)), round(coord[1] * (224 / 56))]

    for r in range(-4, 5):
        for c in range(-4, 5):
            img[:, pt[1] + r, pt[0] + c] = torch.tensor([1, 0, 0])

    return img


def main(
    weights_path: str,
    backbone: str,
    image_path: str,
    coord: list[int],
    output_path: str = "visualize_output",
    suffix: str = "",
):
    # Load backbone
    model = load_backbone(weights_path, backbone)

    # Load image
    img = cv2.imread(image_path, 1)[:, :, ::-1]
    h, w = img.shape[:2]
    img = np.float32(img) / 255
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    preprocess = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Resize((896, 896)),
            transforms.Normalize(mean=mean, std=std),
        ]
    )
    input = preprocess(img.copy()).unsqueeze(0)

    # Crop around target point
    # cropped = crop(
    #     input,
    #     round(coord[1] * (896 / h)) - 112,
    #     round(coord[0] * (896 / w)) - 112,
    #     224,
    #     224,
    # )

    # Pass through backbone
    _, out = model(input, dense=True)  # [28,28]
    out = interpolate(out[-1], scale_factor=2)  # [56,56]

    # Get cosine similarity
    pt_r = round(coord[1] * (56 / h))
    pt_c = round(coord[0] * (56 / w))
    sim = torch.zeros((1, 1, 56, 56))
    for r in range(56):
        for c in range(56):
            sim[0, 0, r, c] = cosine_similarity(out[:, :, pt_r, pt_c], out[:, :, r, c])

    # Pass cropped patch through backbone
    # _, out_c = model(cropped, dense=True)  # [7,7]
    # out_c = interpolate(out_c, scale_factor=2)  # [14,14]

    # # Get cosine similarity
    # pt_r_c = 7
    # pt_c_c = 7
    # sim_c = torch.zeros((1, 1, 14, 14))
    # for r in range(14):
    #     for c in range(14):
    #         sim_c[0, 0, r, c] = cosine_similarity(
    #             out_c[:, :, pt_r_c, pt_c_c], out_c[:, :, r, c]
    #         )

    ## Save similarity matrix
    # Process image
    preprocess = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Resize((224, 224)),
        ]
    )
    img = preprocess(img)
    img = draw_dot(img, [pt_c, pt_r])

    # Fix similarity matrix shape
    sim_int = interpolate(sim, size=(224, 224), mode="bilinear", align_corners=True)
    sim_int = sim_int.repeat(1, 3, 1, 1).squeeze(0)

    # Save
    name = os.path.basename(image_path)[: os.path.basename(image_path).find(".")]
    output_path = os.path.join(output_path, "affinity")
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    file_name = f"{ output_path }/{name}_{coord[0]}_{coord[1]}_{suffix}.png"

    grid = make_grid([img, sim_int])
    save_image(grid, file_name)

    ## Save cropped similarity matrices
    # pt_r = round(coord[1] * (56 / h))
    # pt_c = round(coord[0] * (56 / w))

    # sim = sim[:, :, pt_r - 7 : pt_r + 7, pt_c - 7 : pt_c + 7]
    # sim = interpolate(sim, size=(56, 56))
    # sim = sim.repeat(1, 3, 1, 1).squeeze(0)

    # sim_c = interpolate(sim_c, size=(56, 56))
    # sim_c = sim_c.repeat(1, 3, 1, 1).squeeze(0)

    # file_name = f"{ output_path }/{name}_{coord[0]}_{coord[1]}_{suffix}_crop.png"
    # grid = make_grid([sim_c, sim])
    # save_image(grid, file_name)


if __name__ == "__main__":
    args = get_args()
    main(args.weights, args.backbone, args.image, args.coord, args.output, args.suffix)
