当前位置： 首页 > article >正文

Faster-RCNN代码解读2：快速上手使用

article 2024/11/15 6:02:40

Faster-RCNN代码解读2：快速上手使用

前言

​ 因为最近打算尝试一下Faster-RCNN的复现，不要多想，我还没有厉害到可以一个人复现所有代码。所以，是参考别人的代码，进行自己的解读。

​ 代码来自于B站的UP主（大佬666），其把代码都放到了GitHub上了，我把链接都放到下面了（应该不算侵权吧，毕竟代码都开源了_）：

b站链接：https://www.bilibili.com/video/BV1of4y1m7nj/?vd_source=afeab8b555e5eb1bfa1e7f267262cbf2

GitHub链接：https://github.com/WZMIAOMIAO/deep-learning-for-image-processing

目的

​ 其实UP主已经做了很好的视频讲解了他的代码，只是有时候我还是喜欢阅读博客来学习，另外视频很长，6个小时，我看的时候容易睡着_，所以才打算写博客记录一下学习笔记。

目前完成的内容

​ 第一篇：VOC数据集详细介绍

​ 第二篇：Faster-RCNN代码解读2：快速上手使用（本文）

目录结构

文章目录

    • Faster-RCNN代码解读2：快速上手使用
      • 1. 前言：
      • 2. 下载项目代码：
      • 3. 下载数据和权重文件：
      • 4. predict.py文件解读：
      • 5. pascal_voc_classes.json文件介绍：
      • 6. 快速上手：
      • 7. 总结：

1. 前言：

​ 本篇文章的作用是准备好一些必备的数据或权重文件，以实现直接快速使用代码的目的。

2. 下载项目代码：

​ 打开大佬的GitHub链接，然后，进入pytorch_object_detection文件内：

在这里插入图片描述

​ 然后，把Faster-RCNN文件夹下载下来即可。不过，GitHub本身不支持单个文件夹的下载，这时候推荐一下浏览器的插件GitZip for github ，把这个插件安装后，即可下载单独的文件夹，如下图所示：

在这里插入图片描述

​ 下载完成后的目录结构如下：

在这里插入图片描述

3. 下载数据和权重文件：

​ 打开README.md文件，里面说明了预训练权重文件和数据集的下载地址：

  • ResNet50+FPN权重文件下载：
官方的权重文件：https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth

up主自己训练后的权重地址：
https://pan.baidu.com/s/1ifilndFRtAV5RDZINSHj5w 提取码:dsz8
  • 数据集下载地址：
http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar

​ 完成上述下载后，可以得到下图的文件：

在这里插入图片描述

4. predict.py文件解读：

​ 打开predict.py文件，这个文件的作用就是加载已经训练过的模型，对一张图片进行目标检测。

main函数：

​ 看main函数，主要分为四个部分：

  • 设置权重文件路径（需要我们改的参数），并用模型加载：
# 选定GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device))

# 创建模型:21=20个类别+1个背景
model = create_model(num_classes=21)

# 加载权重参数
# weights_path = "./save_weights/model.pth"   # 权重保存路径，作者自己定义的
weights_path = "./fasterrcnn_voc2012.pth"   # 权重保存路径，我们下载后自己的路径
assert os.path.exists(weights_path), "{} file dose not exist.".format(weights_path)
# 开始加载权重文件
weights_dict = torch.load(weights_path, map_location='cpu')  # 加载之前训练保存的字典
weights_dict = weights_dict["model"] if "model" in weights_dict else weights_dict # 选定model参数
model.load_state_dict(weights_dict) # 加载
model.to(device) # 放入GPU
  • 读取“类别----数字值”的json文件，并生成一个字典，以方便后期将预测的类别（比如：1、2这样的数字）转为字符串（比如：person、bicycle等）
# 读取json文件
label_json_path = './pascal_voc_classes.json'
assert os.path.exists(label_json_path), "json file {} dose not exist.".format(label_json_path)
with open(label_json_path, 'r') as f:
class_dict = json.load(f)
# 将值转为字典
category_index = {str(v): str(k) for k, v in class_dict.items()}
  • 加载一张测试图片（需要改为我们自己的），并改为符合[batch,channel,w,h]的格式
# 加载一张测试图片
original_img = Image.open("./test.jpg") # 需要改为自己的路径

# 将PIL图像格式转为tensor格式
data_transform = transforms.Compose([transforms.ToTensor()])
img = data_transform(original_img)
# 增加一个batch维度，符合训练图片格式
img = torch.unsqueeze(img, dim=0)
  • 开始预测图片，并计算运行的时间（一般不计算第一次的时间，因为GPU调用需要时间）和画出对应图像
model.eval()  # 进入验证模式
with torch.no_grad():
    # init
    # 初始化，原始图像的宽、高
    img_height, img_width = img.shape[-2:]
    # 将图像放入GPU中，并变为model可以识别的格式[batch_size,channel,w,h]
    init_img = torch.zeros((1, 3, img_height, img_width), device=device)
    # 验证
    model(init_img)

    # 计算预测时间，不过不能直接计算第一次，因为需要启动gpu等
    t_start = time_synchronized()
    predictions = model(img.to(device))[0]
    t_end = time_synchronized()
    print("inference+NMS time: {}".format(t_end - t_start))

    # 得到预测的相关参数
    predict_boxes = predictions["boxes"].to("cpu").numpy()
    predict_classes = predictions["labels"].to("cpu").numpy()
    predict_scores = predictions["scores"].to("cpu").numpy()

    if len(predict_boxes) == 0:
    	print("没有检测到任何目标!")

    # 绘制图像
    plot_img = draw_objs(original_img,
                             predict_boxes,
                             predict_classes,
                             predict_scores,
                             category_index=category_index,
                             box_thresh=0.5,
                             line_thickness=3,
                             font='arial.ttf',
                             font_size=20)
    plt.imshow(plot_img)
    plt.show()
    # 保存预测的图片结果
    plot_img.save("test_result.jpg")

create_model函数

​ 了解了main函数后，我们再看看create_model函数，这个函数的作用就是创建模型。作者在该项目中采取了很多模型，比如VGG16、mobilenetv2、resnet等等，而这里我们用的是刚刚下载的权重文件对应的模型，即resNet50+fpn+faster-rcnn，因此需要把其它的模型代码注释掉：

def create_model(num_classes):
    # mobileNetv2+faster_RCNN
    # backbone = MobileNetV2().features
    # backbone.out_channels = 1280
    #
    # anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512),),
    #                                     aspect_ratios=((0.5, 1.0, 2.0),))
    #
    # roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
    #                                                 output_size=[7, 7],
    #                                                 sampling_ratio=2)
    #
    # model = FasterRCNN(backbone=backbone,
    #                    num_classes=num_classes,
    #                    rpn_anchor_generator=anchor_generator,
    #                    box_roi_pool=roi_pooler)

    # resNet50+fpn+faster_RCNN
    # 注意，这里的norm_layer要和训练脚本中保持一致
    backbone = resnet50_fpn_backbone(norm_layer=torch.nn.BatchNorm2d)
    model = FasterRCNN(backbone=backbone, num_classes=num_classes, rpn_score_thresh=0.5)

    return model

5. pascal_voc_classes.json文件介绍：

​ 我们再看看上面涉及json文件，这个文件就是voc数据集的类别和数字值的对应关系，比如：

{
    "aeroplane": 1,
    "bicycle": 2,
    "bird": 3,
    "boat": 4,
    "bottle": 5,
    "bus": 6,
    "car": 7,
    "cat": 8,
    "chair": 9,
    "cow": 10,
    "diningtable": 11,
    "dog": 12,
    "horse": 13,
    "motorbike": 14,
    "person": 15,
    "pottedplant": 16,
    "sheep": 17,
    "sofa": 18,
    "train": 19,
    "tvmonitor": 20
}

​ 需要注意的是，这里的值是从1开始的，是因为0一般是留给背景的。

6. 快速上手：

​ 有了上面的解读后，我们可以快速上手看看效果。

​ 这里再次声明一下predict.py文件需要修改**权重文件路径和自己搞一张测试图片并修改路径。**完成修改后，直接运行该文件即可，我测试了几张图片，结果如下图：

在这里插入图片描述
在这里插入图片描述

7. 总结：

​ 上面主要简单介绍了如何快速上手，看到结果，给自己一种这个很简单的错觉。后面，主要就是对一些主要的文件进行解读。


http://www.kler.cn/a/10375.html

相关文章：

  • Jetpack 之 Ink API初探
  • INQUIRE：一个包含五百万张自然世界图像，涵盖10,000个不同物种的专为专家级文本到图像检索任务设计的新型基准数据集。
  • Spring MVC 与 JSP 数据传输
  • 基于混合配准策略的多模态医学图像配准方法研究
  • 94个属于一区且接受医工交叉领域投稿的期刊汇总｜个人观点·24-11-13
  • leetcode hot100【LeetCode 114.二叉树展开为链表】java实现
  • 【数据结构与算法】栈和队列(StackQueue)
  • 【BAT】查询局域内所有的IP
  • 2023 年十大 API 管理趋势
  • 基于matlab进行雷达信号模拟
  • 朝花夕拾 - 卷王的自白（光头祭天，法力无边》
  • 独立看门狗（IWDG）实验
  • 第十三届蓝桥杯大赛软件赛省赛 C/C++ 大学 B 组思考+总结
  • ChatGPT宝藏插件丨装上之后，上网、语音聊天、一键分享对话……简直让你爽到起飞！
  • Java题目训练——年终奖和迷宫问题
  • 自己再造一个大规模预训练语言模型？可以的
  • 【C++】Vector
  • 软件测试培训
  • JUC源码系列-AQS的Condition的接口实现
  • 程序员在职场中如何让自己的技能快速提高，WEB前端开发工程师如何让自己快速成为团队的核心开发人员？
  • Follow My Heart Of Apirl. 2023
  • LAZADA平台开放接口的接入和参数说明（目前支持以下基本接口：item_get 获得淘宝商品详情item_search 获得淘宝商品详情）
  • 为何ChatGPT如此擅长编造故事？
  • 操作系统 - 学习笔记
  • 不一样的websocket封装简洁版
  • 设计模式之迭代器模式(C++)