Getting Visualizations as Tensor (RT-DETR)

Discussion
,
1

I want to detect using RT-DETR and also get it’s heatmap. However, it seems like the only way get the heatmap is turn visualization on and it will save it in a directory. Is there no way for me to get the array of the heatmap directly as an output?

2

It would take some serious modifications to the source code for this, as there is no attribute handling for this in the Results object. Additionally, given that there would be an activation map for each layer in the model for each image, this would rapidly consume RAM usage, which is why it’s preferable to save these to disk.

There should be a .npy file saved that has the array data for the feature maps. You should be able to load these using numpy to retrieve the array data from disk.

import numpy as np
from ultralytics import YOLO, ASSETS

im = ASSETS / "bus.jpg"
model = YOLO("yolo11n.pt")
results = model.predict(im, visualize=True)

>>> runs/detect/predict/bus/stage0_Conv_features.png... (16/16)
>>> ...

array = np.load(
    "runs/detect/predict/bus/stage0_Conv_features.npy"
)
array.shape
>>> (16, 320, 240)  # iterate first dimension for each feature map
3

Thanks for the response BurhanQ.

I also couldn’t find any easy way, so currently I have just resorted to forking the code and commenting out everything I dont need from the visualization function so that it still runs fast.

1 Like