Predicted Image id and box from SSD

Question

How to find predicted image id and Box from SSD, I am using this GitHub link here is the test function which I want to save the image id and box

def test(loader, net, criterion, device):
net.eval()
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
num = 0
for _, data in enumerate(loader):
images, boxes, labels = data
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
num += 1

    with torch.no_grad():
        confidence, locations = net(images)
        regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
        loss = regression_loss + classification_loss

    running_loss += loss.item()
    running_regression_loss += regression_loss.item()
    running_classification_loss += classification_loss.item()
return running_loss / num, running_regression_loss / num, run

Answer 1

Assume

y = net(x)
detections = y.data

you can print the detection info with the following

# this will loop over predictions class by class
for i in range(detections.size(1)): 
    # this will loop over each detection in the class
    for j in range(detections.size(2)):
        score = detection[0,i,j,0]
        coords = detections[0,i,j,1:]
        print(f"Class id: {i} \t Score: {score} \t Coords: {coords}")

See the demo for more details.