DETR实现目标检测(二)-利用自己训练的模型进行预测

1、图片预测(CPU)

关于DETR模型训练自己的数据集参考上篇文章:

DETR实现目标检测(一)-训练自己的数据集-CSDN博客

训练完成后的模型文件保存位置如下:

准备好要预测的图片:

然后直接调用模型进行预测,并设置置信度阈值来输出检测框:

最后用plot函数来画出图片及预测框,效果如下:

最后附上完整代码:

python 复制代码
from PIL import Image
import matplotlib.pyplot as plt
import torchvision.transforms as T
from hubconf import *
from util.misc import nested_tensor_from_tensor_list

torch.set_grad_enabled(False)

# COCO classes
CLASSES = [
    '1'
]

# colors for visualization
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098]]

# standard PyTorch mean-std input image normalization
transform = T.Compose([
    T.Resize(800),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])


# for output bounding box post-processing
def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=1)


def rescale_bboxes(out_bbox, size):
    img_w, img_h = size
    b = box_cxcywh_to_xyxy(out_bbox)
    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
    return b


def predict(im, model, transform):
    # mean-std normalize the input image (batch-size: 1)
    anImg = transform(im)
    data = nested_tensor_from_tensor_list([anImg])

    # propagate through the model
    outputs = model(data)

    # keep only predictions with 0.7+ confidence
    probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]

    keep = probas.max(-1).values > 5*1e-8  # 置信度阈值

    # convert boxes from [0; 1] to image scales
    bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)
    return probas[keep], bboxes_scaled


def plot_results(pil_img, prob, boxes):
    plt.figure(figsize=(16, 10))
    plt.imshow(pil_img)
    ax = plt.gca()
    colors = COLORS * 100
    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
                                   fill=False, color=c, linewidth=3))
        cl = p.argmax()
        text = f'{CLASSES[cl]}: {p[cl]:0.2f}'
        ax.text(xmin, ymin, text, fontsize=15,
                bbox=dict(facecolor='yellow', alpha=0.5))
    plt.axis('off')
    plt.show()


if __name__ == "__main__":
    model = detr_resnet50(False, 1) # 这里与前面的num_classes数值相同,就是最大的category id值 + 1
    state_dict = torch.load(r"C:\Users\90539\Downloads\detr-main\detr-main\data\output\checkpoint.pth", map_location='cpu')
    model.load_state_dict(state_dict["model"])
    model.eval()

    # im = Image.open('data/coco_frame_count/train2017/001554.jpg')
    im = Image.open(r'C:\Users\90539\Downloads\detr-main\detr-main\data/coco_frame_count/val2017/09-12-52-0.png')

    scores, boxes = predict(im, model, transform)
    plot_results(im, scores, boxes)
相关推荐
豆浩宇4 小时前
Conda环境隔离和PyCharm配置,完美同时运行PaddlePaddle和PyTorch
人工智能·pytorch·算法·计算机视觉·pycharm·conda·paddlepaddle
㱘郳10 小时前
cifar10分类对比:使用PyTorch卷积神经网络和SVM
pytorch·分类·cnn
Tiger Z1 天前
《动手学深度学习v2》学习笔记 | 2.4 微积分 & 2.5 自动微分
pytorch·深度学习·ai
先做个垃圾出来………1 天前
PyTorch 模型文件介绍
人工智能·pytorch·python
我不是小upper1 天前
一文详解深度学习中神经网络的各层结构与功能!
人工智能·pytorch·深度学习
钱彬 (Qian Bin)1 天前
一文掌握工业缺陷检测项目实战(Pytorch算法训练、部署、C++ DLL制作、Qt集成)
c++·pytorch·python·qt·实战·工业缺陷检测·faster rcnn
vvilkim1 天前
PyTorch 中的循环神经网络 (RNN/LSTM):时序数据处理实战指南
pytorch·rnn·lstm
Hello Mr.Z1 天前
使用pytorch创建/训练/推理OCR模型
人工智能·pytorch·python
点云SLAM2 天前
PyTorch 中.backward() 详解使用
人工智能·pytorch·python·深度学习·算法·机器学习·机器人
山烛2 天前
深度学习:CNN 模型训练中的学习率调整(基于 PyTorch)
人工智能·pytorch·python·深度学习·cnn·调整学习率