Ugly, dirty example of interactive feature-map visualization in Plotly Dash

vis.py

from dash import Dash, Patch, dcc, html, Input, Output, callback
import plotly.express as px
import numpy as np
import torch
import torchvision

app = Dash(__name__)


def load_image(image_path):
    import cv2

    image = cv2.imread(str(image_path))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    image = cv2.resize(image, (768, 1024))  # Scale down for demo purposes
    return image


img = load_image("res/dog.jpg")
net = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)


def forward(net, x):
    x = net.conv1(x)
    x = net.bn1(x)
    x = net.relu(x)
    x = net.maxpool(x)

    x = net.layer1(x)
    x = net.layer2(x)
    x = net.layer3(x)
    # x = net.layer4(x)

    return x


net.eval()

_fmap = forward(net, torch.tensor(img)[None, ...].permute(0, 3, 1, 2).float()) / 255.0
fmap = _fmap.detach()[0].permute(1, 2, 0).numpy()

hmap = np.mean(fmap, axis=-1)
fig = px.imshow(hmap)

app.layout = html.Div(
    style={"display": "flex"},
    children=[
        html.Div(
            style={"flex": "1"},
            children=[dcc.Graph(id="image", figure=px.imshow(img))],
        ),
        html.Div(
            style={"flex": "1"},
            children=[dcc.Graph(id="feature-map", figure=fig)],
        ),
    ],
)


@callback(
    Output("feature-map", "figure"),
    Input("feature-map", "clickData"),
    prevent_initial_call=True,
)
def update_on_click(click_data):
    if not click_data:
        return Patch()  # Do nothing

    x = click_data["points"][0]["x"]
    y = click_data["points"][0]["y"]

    patched_fig = Patch()

    H, W, C = fmap.shape
    hmap = np.dot(fmap[None, y, x, :], fmap.reshape(-1, C).T).reshape(
        H, W
    ) / np.linalg.norm(fmap, axis=-1)

    patched_fig["data"][0]["z"] = hmap

    return patched_fig


if __name__ == "__main__":
    app.run(debug=True)