Hook RoIHeads layer

Hi, I’ve been trying to hook in the Faster R-CNN ResNet50 pre-trained model in the RoIHeads layer.

First, I load the model and feed the image to the model. After that, I hook in the layer to access its tensors. My objective is to obtain the boxes with a confidence threshold of more than 50%.

My question is, did I write the correct code? If not, what should I do instead?

import torch
import torchvision
import pickle
from PIL import Image

# Define the model
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()

!wget -nv https://www.goodfreephotos.com/cache/other-photos/car-and-traffic-on-the-road-coming-towards-me.jpg -O traffic.jpg
image = Image.open('traffic.jpg')
image_tensor = torchvision.transforms.ToTensor()(image)
image_tensor = image_tensor.unsqueeze(0)

outputs = []
hook = model.backbone.register_forward_hook(lambda self, image_tensor, output: outputs.append(output))
res = model(image_tensor)
hook.remove()
selected_rois = model.roi_heads.box_roi_pool(outputs[0], [r['boxes'] for r in res], [i.shape[-2:] for i in image_tensor])
print(selected_rois.shape)

# Define the threshold value
threshold = 0.5

# Get the box features and scores
box_features = res[0]
scores = box_features.get('scores') # Assuming the last column contains the scores

# find indices of grather score
above_threshold_indices = torch.where(scores > threshold)[0] #tensor([0, 1, 2, 3])

#get value of scores
above_threshold_scores = scores[above_threshold_indices] #tensor([0.9840, 0.9784, 0.9386, 0.8663]

#finde the features
above_threshold_box_features = selected_rois[above_threshold_indices]
above_threshold_box_features.shape