wandb_add_detections.py

from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
from PIL import Image
import requests
from tqdm import tqdm

def add_detections(model, processor, samples, field_name, device=None, gt_field="ground_truth"):
    # The model predicts object classes with integer IDs, but FiftyOne expected class strings
    id_label_map = {i:c for i,c in enumerate(dataset.distinct(gt_field+".detections.label"))}
    
    if device is None:
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    values = {}
    model.to(device)
    for fp in tqdm(samples.values("filepath")):
        image = Image.open(fp)
        imw, imh = image.size

        inputs = processor(images=image, return_tensors="pt").to(device)
        outputs = model(**inputs)

        # convert outputs (bounding boxes and class logits) to FiftyOne format
        # let's only keep detections with score > 0.2
        target_sizes = torch.tensor([image.size[::-1]])
        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.2)[0]
        dets = []
        for score, label, box in zip(results["scores"].tolist(), results["labels"].tolist(), results["boxes"].tolist()):
            tlx, tly, brx, bry = box
            class_name = id_label_map[int(label)]
            box_w = brx-tlx
            box_h = bry-tly
            fo_box = [max(0,tlx/imw), max(tly/imh,0), min(1,box_w/imw), min(1,box_h/imh)]
            det = fo.Detection(label=class_name, confidence=score, bounding_box=fo_box)
            dets.append(det)
        values[fp] = fo.Detections(detections=dets)
        
    samples.set_values(field_name, values, key_field="filepath")