一键体验Detectron2框架中的所有预训练模型

已于 2024-08-01 19:01:12 修改
阅读量399 收藏 11
点赞数 10
分类专栏： Detectron2 文章标签： 机器学习
于 2024-08-01 11:44:36 首次发布
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/fydw_715/article/details/140825847
版权

Detectron2是由Facebook AI Research (FAIR)推出的基于PyTorch的模块化物体检测库，发布于2019年10月10日。该平台原是2018年推出的Detectron的第二代版本，它完全重写于maskrcnn-benchmark，并采用了PyTorch语言实现。与原版相比，Detectron2具有更灵活和可扩展的设计，可以在单个或多个GPU服务器上快速训练。

主要特点：

  1. 实现新颖：

    • Detectron2是从头重新设计的，它采用了PyTorch而不是原来的Caffe2。这种转变使得编程模型更加直观，提高了模型设计和实验的迭代速度。
    • PyTorch的大型社区和不断改进也为Detectron2用户提供了巨大利益。

  2. 模块化和可扩展性：

    • Detectron2采用模块化设计，允许用户在几乎任何部分插入自定义模块，使新研究项目可以在数百行代码内实现，并且清晰地与核心库分离。
    • 这种设计实现了模型和功能的快速集成与扩展，例如高级姿态检测模型的快速设计和训练。

  3. 新功能和模型：

    • Detectron2整合了DensePose、Panoptic FPN等最新算法，并支持新的数据集如LVIS。
    • 它支持各种任务，包括对象检测、实例分割、语义分割和全景分割。

  4. 实现质量：

    • 通过重新设计，解决了原版Detectron中的一些实现问题，大大提高了整体性能。
    • 整个训练管道已迁移至GPU，提升了训练速度且易于拓展至多GPU服务器，便于处理大规模数据集。

  5. Detectron2go：

    • Detectron2go是一个附加的软件层，旨在简化模型的生产部署，包括标准训练工作流程、网络量化和模型转换等。

预训练模型

下面的表格列出了detectron2/configs源码目录下的模型及其简要的中文描述。

分类Model Config File Path(–model-name 参数)简要描述
Base ModelsBase-RCNN-C4.yaml基于C4的基础RCNN模型
Base-RCNN-DilatedC5.yaml基于扩张卷积C5的基础RCNN模型
Base-RCNN-FPN.yaml基于FPN的基础RCNN模型
Base-RetinaNet.yaml基础的RetinaNet模型
COCO目标检测COCO-Detection/fast_rcnn_R_50_FPN_1x.yaml采用FPN的ResNet-50的Fast R-CNN模型，1倍学习速率
COCO-Detection/faster_rcnn_R_101_C4_3x.yaml采用C4的ResNet-101的Faster R-CNN模型，3倍学习速率
COCO-Detection/faster_rcnn_R_101_DC5_3x.yaml采用DC5的ResNet-101的Faster R-CNN模型，3倍学习速率
COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml采用FPN的ResNet-101的Faster R-CNN模型，3倍学习速率
COCO-Detection/faster_rcnn_R_50_C4_1x.yaml采用C4的ResNet-50的Faster R-CNN模型，1倍学习速率
COCO-Detection/faster_rcnn_R_50_C4_3x.yaml采用C4的ResNet-50的Faster R-CNN模型，3倍学习速率
COCO-Detection/faster_rcnn_R_50_DC5_1x.yaml采用DC5的ResNet-50的Faster R-CNN模型，1倍学习速率
COCO-Detection/faster_rcnn_R_50_DC5_3x.yaml采用DC5的ResNet-50的Faster R-CNN模型，3倍学习速率
COCO-Detection/faster_rcnn_R_50_FPN_1x.yaml采用FPN的ResNet-50的Faster R-CNN模型，1倍学习速率
COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml采用FPN的ResNet-50的Faster R-CNN模型，3倍学习速率
COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml采用FPN的ResNeXt-101的Faster R-CNN模型，3倍学习速率
COCO-Detection/fcos_R_50_FPN_1x.py采用FPN的ResNet-50的FCOS模型，1倍学习速率
COCO-Detection/retinanet_R_101_FPN_3x.yaml采用FPN的ResNet-101的RetinaNet模型，3倍学习速率
COCO-Detection/retinanet_R_50_FPN_1x.py采用FPN的ResNet-50的RetinaNet模型，1倍学习速率
COCO-Detection/retinanet_R_50_FPN_1x.yaml采用FPN的ResNet-50的RetinaNet模型，1倍学习速率
COCO-Detection/retinanet_R_50_FPN_3x.yaml采用FPN的ResNet-50的RetinaNet模型，3倍学习速率
COCO-Detection/rpn_R_50_C4_1x.yaml采用C4的ResNet-50的RPN模型，1倍学习速率
COCO-Detection/rpn_R_50_FPN_1x.yaml采用FPN的ResNet-50的RPN模型，1倍学习速率
COCO实例分割COCO-InstanceSegmentation/mask_rcnn_R_101_C4_3x.yaml采用C4的ResNet-101的Mask R-CNN模型，3倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_101_DC5_3x.yaml采用DC5的ResNet-101的Mask R-CNN模型，3倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml采用FPN的ResNet-101的Mask R-CNN模型，3倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_50_C4_1x.py采用C4的ResNet-50的Mask R-CNN模型，1倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_50_C4_1x.yaml采用C4的ResNet-50的Mask R-CNN模型，1倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_50_C4_3x.yaml采用C4的ResNet-50的Mask R-CNN模型，3倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_50_DC5_1x.yaml采用DC5的ResNet-50的Mask R-CNN模型，1倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_50_DC5_3x.yaml采用DC5的ResNet-50的Mask R-CNN模型，3倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.py采用FPN的ResNet-50的Mask R-CNN模型，1倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml采用FPN的ResNet-50的Mask R-CNN模型，1倍学习速率
COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x_giou.yaml采用FPN的ResNet-50的Mask R-CNN模型，1倍学习速率，使用GIoU
COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml采用FPN的ResNet-50的Mask R-CNN模型，3倍学习速率
COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml采用FPN的ResNeXt-101的Mask R-CNN模型，3倍学习速率
COCO-InstanceSegmentation/mask_rcnn_regnetx_4gf_dds_fpn_1x.py采用RegNetX-4GF的Mask R-CNN模型，1倍学习速率
COCO-InstanceSegmentation/mask_rcnn_regnety_4gf_dds_fpn_1x.py采用RegNetY-4GF的Mask R-CNN模型，1倍学习速率
COCO关键点检测COCO-Keypoints/keypoint_rcnn_R_101_FPN_3x.yaml采用FPN的ResNet-101的关键点检测RCNN模型，3倍学习速率
COCO-Keypoints/keypoint_rcnn_R_50_FPN_1x.py采用FPN的ResNet-50的关键点检测RCNN模型，1倍学习速率
COCO-Keypoints/keypoint_rcnn_R_50_FPN_1x.yaml采用FPN的ResNet-50的关键点检测RCNN模型，1倍学习速率
COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml采用FPN的ResNet-50的关键点检测RCNN模型，3倍学习速率
COCO-Keypoints/keypoint_rcnn_X_101_32x8d_FPN_3x.yaml采用FPN的ResNeXt-101的关键点检测RCNN模型，3倍学习速率
COCO全景分割COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml采用FPN的ResNet-101的全景分割模型，3倍学习速率
COCO-PanopticSegmentation/panoptic_fpn_R_50_1x.py采用FPN的ResNet-50的全景分割模型，1倍学习速率
COCO-PanopticSegmentation/panoptic_fpn_R_50_1x.yaml采用FPN的ResNet-50的全景分割模型，1倍学习速率
COCO-PanopticSegmentation/panoptic_fpn_R_50_3x.yaml采用FPN的ResNet-50的全景分割模型，3倍学习速率
CityscapesCityscapes/mask_rcnn_R_50_FPN.yaml采用FPN的ResNet-50的Cityscapes实例分割模型
Detectron1比较Detectron1-Comparisons/faster_rcnn_R_50_FPN_noaug_1x.yaml采用FPN的ResNet-50的Faster R-CNN模型，无数据增强，1倍学习速率
Detectron1-Comparisons/keypoint_rcnn_R_50_FPN_1x.yaml采用FPN的ResNet-50的关键点检测RCNN模型，1倍学习速率
Detectron1-Comparisons/mask_rcnn_R_50_FPN_noaug_1x.yaml采用FPN的ResNet-50的Mask R-CNN模型，无数据增强，1倍学习速率
LVIS v0.5实例分割LVISv0.5-InstanceSegmentation/mask_rcnn_R_101_FPN_1x.yaml采用FPN的ResNet-101的Mask R-CNN模型，1倍学习速率，适用于LVIS v0.5
LVISv0.5-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml采用FPN的ResNet-50的Mask R-CNN模型，1倍学习速率，适用于LVIS v0.5
LVISv0.5-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_1x.yaml采用FPN的ResNeXt-101的Mask R-CNN模型，1倍学习速率，适用于LVIS v0.5
LVIS v1实例分割LVISv1-InstanceSegmentation/mask_rcnn_R_101_FPN_1x.yaml采用FPN的ResNet-101的Mask R-CNN模型，1倍学习速率，适用于LVIS v1
LVISv1-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml采用FPN的ResNet-50的Mask R-CNN模型，1倍学习速率，适用于LVIS v1
LVISv1-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_1x.yaml采用FPN的ResNeXt-101的Mask R-CNN模型，1倍学习速率，适用于LVIS v1
杂项Misc/cascade_mask_rcnn_R_50_FPN_1x.yaml基于FPN的ResNet-50的级联Mask R-CNN模型，1倍学习速率
Misc/cascade_mask_rcnn_R_50_FPN_3x.yaml基于FPN的ResNet-50的级联Mask R-CNN模型，3倍学习速率
Misc/cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml基于FPN的ResNeXt-152的级联Mask R-CNN模型，使用组归一化和扩张卷积
Misc/mask_rcnn_R_50_FPN_1x_cls_agnostic.yaml基于FPN的ResNet-50的类别无关的Mask R-CNN模型，1倍学习速率
Misc/mask_rcnn_R_50_FPN_1x_dconv_c3-c5.yaml基于FPN的ResNet-50的带有扩张卷积的Mask R-CNN模型，1倍学习速率
Misc/mask_rcnn_R_50_FPN_3x_dconv_c3-c5.yaml基于FPN的带有扩张卷积的Mask R-CNN模型，3倍学习速率
Misc/mask_rcnn_R_50_FPN_3x_gn.yaml基于FPN的ResNet-50的带有组归一化的Mask R-CNN模型，3倍学习速率
Misc/mask_rcnn_R_50_FPN_3x_syncbn.yaml基于FPN的ResNet-50的带有同步批归一化的Mask R-CNN模型，3倍学习速率
Misc/panoptic_fpn_R_101_dconv_cascade_gn_3x.yaml基于FPN的ResNet-101的带有扩张卷积和级联组归一化的全景分割模型，3倍学习速率
Misc/scratch_mask_rcnn_R_50_FPN_3x_gn.yaml基于FPN的ResNet-50的从头训练的带组归一化的Mask R-CNN模型，3倍学习速率
Misc/scratch_mask_rcnn_R_50_FPN_9x_gn.yaml基于FPN的ResNet-50的从头训练的带组归一化的Mask R-CNN模型，9倍学习速率
Misc/scratch_mask_rcnn_R_50_FPN_9x_syncbn.yaml基于FPN的ResNet-50的从头训练的带同步批归一化的Mask R-CNN模型，9倍学习速率
Misc/semantic_R_50_FPN_1x.yaml采用FPN的ResNet-50的语义分割模型，1倍学习速率
Pascal VOC检测PascalVOC-Detection/faster_rcnn_R_50_C4.yaml采用C4的ResNet-50的Faster R-CNN模型，PascalVOC检测任务
PascalVOC-Detection/faster_rcnn_R_50_FPN.yaml采用FPN的ResNet-50的Faster R-CNN模型，PascalVOC检测任务

模型体验

上面的表格列出了不同类别的模型和相应的配置文件路径。执行下面的命令下载所有的预训练模型并推理。例如：

python3 infer_allmode.py --input input.jpg --output ./outimgs --cache-dir /liveguard/Detectron2/models/

该命令将 infer_allmode.py 脚本与指定的输入、输出和缓存目录一起执行。具体而言：

  1. 输入图像：input.jpg 是将被加载并用于模型推理的图像文件。
  2. 输出目录：./outimgs 是模型推理的结果图像保存的目录。
  3. 缓存目录：/liveguard/Detectron2/models/ 是模型文件的缓存位置，以避免每次运行脚本时重新下载模型。

执行 infer_allmode.py 脚本后，在不同模型配置文件上的推理结果保存在outimgs目录，目录下的目录结构的层次树，每个子目录对应一种模型类型，文件名的格式为 input_<模型名称>.jpg，已命名的结果图像证明了每个模型的推理任务已经成功执行并保存到了相应的目录中

输出目录结构

顶层目录
/liveguard/Detectron2/outimgs

该目录包含多个子目录，每个子目录对应不同的模型分类。

子目录及其内容

  1. COCO-Detection：

    • 包含对 COCO 数据集进行目标检测的模型推理结果图像。
    • 文件名格式：input_<模型名称>.jpg。
    ├── input_faster_rcnn_R_101_C4_3x.jpg
├── input_faster_rcnn_R_101_DC5_3x.jpg
├── input_faster_rcnn_R_101_FPN_3x.jpg
├── input_faster_rcnn_R_50_C4_1x.jpg
├── input_faster_rcnn_R_50_C4_3x.jpg
├── input_faster_rcnn_R_50_DC5_1x.jpg
├── input_faster_rcnn_R_50_DC5_3x.jpg
├── input_faster_rcnn_R_50_FPN_1x.jpg
├── input_faster_rcnn_R_50_FPN_3x.jpg
├── input_faster_rcnn_X_101_32x8d_FPN_3x.jpg
├── input_retinanet_R_101_FPN_3x.jpg
├── input_retinanet_R_50_FPN_1x.jpg
└── input_retinanet_R_50_FPN_3x.jpg

请添加图片描述

  1. COCO-InstanceSegmentation：

    • 包含对 COCO 数据集进行实例分割的模型推理结果图像。
    • 文件名格式与 COCO-Detection 类似。
    ├── input_mask_rcnn_R_101_C4_3x.jpg
├── input_mask_rcnn_R_101_DC5_3x.jpg
├── input_mask_rcnn_R_101_FPN_3x.jpg
├── input_mask_rcnn_R_50_C4_1x.jpg
├── input_mask_rcnn_R_50_C4_3x.jpg
├── input_mask_rcnn_R_50_DC5_1x.jpg
├── input_mask_rcnn_R_50_DC5_3x.jpg
├── input_mask_rcnn_R_50_FPN_1x.jpg
├── input_mask_rcnn_R_50_FPN_3x.jpg
└── input_mask_rcnn_X_101_32x8d_FPN_3x.jpg

请添加图片描述

  1. COCO-Keypoints：

    • 包含对 COCO 数据集进行关键点检测的模型推理结果图像。
    ├── input_keypoint_rcnn_R_101_FPN_3x.jpg
├── input_keypoint_rcnn_R_50_FPN_1x.jpg
├── input_keypoint_rcnn_R_50_FPN_3x.jpg
└── input_keypoint_rcnn_X_101_32x8d_FPN_3x.jpg

请添加图片描述

  1. COCO-PanopticSegmentation：

    • 包含对 COCO 数据集进行全景分割的模型推理结果图像。
    ├── input_panoptic_fpn_R_101_3x.jpg
├── input_panoptic_fpn_R_50_1x.jpg
└── input_panoptic_fpn_R_50_3x.jpg

请添加图片描述

  1. Cityscapes：

    • 包含对 Cityscapes 数据集进行实例分割的模型推理结果图像。
    └── input_mask_rcnn_R_50_FPN.jpg

请添加图片描述

  1. Detectron1-Comparisons：

    • 包含用于与 Detectron1 进行比较的模型推理结果图像。
    ├── input_faster_rcnn_R_50_FPN_noaug_1x.jpg
├── input_keypoint_rcnn_R_50_FPN_1x.jpg
└── input_mask_rcnn_R_50_FPN_noaug_1x.jpg

请添加图片描述

  1. LVISv0.5-InstanceSegmentation：

    • 包含对 LVIS v0.5 数据集进行实例分割的模型推理结果图像。
    ├── input_mask_rcnn_R_101_FPN_1x.jpg
├── input_mask_rcnn_R_50_FPN_1x.jpg
└── input_mask_rcnn_X_101_32x8d_FPN_1x.jpg

请添加图片描述

  1. Misc：

    • 包含各种类型的模型推理结果图像。
    ├── input_cascade_mask_rcnn_R_50_FPN_1x.jpg
├── input_cascade_mask_rcnn_R_50_FPN_3x.jpg
├── input_cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.jpg
├── input_mask_rcnn_R_50_FPN_1x_dconv_c3-c5.jpg
├── input_mask_rcnn_R_50_FPN_3x_dconv_c3-c5.jpg
├── input_mask_rcnn_R_50_FPN_3x_gn.jpg
├── input_mask_rcnn_R_50_FPN_3x_syncbn.jpg
├── input_panoptic_fpn_R_101_dconv_cascade_gn_3x.jpg
├── input_scratch_mask_rcnn_R_50_FPN_3x_gn.jpg
├── input_scratch_mask_rcnn_R_50_FPN_9x_gn.jpg
└── input_scratch_mask_rcnn_R_50_FPN_9x_syncbn.jpg

请添加图片描述

  1. PascalVOC-Detection：

    • 包含对 Pascal VOC 数据集进行目标检测的模型推理结果图像。
    └── input_faster_rcnn_R_50_C4.jpg

请添加图片描述

总结

该结果展示了执行 infer_allmode.py 脚本后，在不同模型配置文件上的推理结果。每个子目录对应一种模型类型，文件名的格式为 input_<模型名称>.jpg，已命名的结果图像证明了每个模型的推理任务已经成功执行并保存到了相应的目录中。通过这种方式，可以方便地比较不同模型在相同输入图像上的推理结果。

代码及概述

下面是带有详细注释的代码：

import argparse  # 用于处理命令行参数
import os  # 提供与操作系统相关的功能
import cv2  # OpenCV库，用于图像处理
import torch  # PyTorch库，深度学习的基础库
import time  # 提供时间相关的功能
from detectron2 import model_zoo  # Detectron2的模型库
from detectron2.engine import DefaultPredictor  # Detectron2的默认预测器
from detectron2.config import get_cfg  # 获取Detectron2的配置对象
from detectron2.utils.logger import setup_logger  # 设置日志工具
from detectron2.utils.visualizer import Visualizer  # 可视化工具
from detectron2.data import MetadataCatalog  # 用于管理数据集元数据

# 所有模型配置文件的字典
MODEL_CONFIGS = {
    "COCO-Detection": [
        "fast_rcnn_R_50_FPN_1x.yaml",
        "faster_rcnn_R_101_C4_3x.yaml",
        "faster_rcnn_R_101_DC5_3x.yaml",
        "faster_rcnn_R_101_FPN_3x.yaml",
        "faster_rcnn_R_50_C4_1x.yaml",
        "faster_rcnn_R_50_C4_3x.yaml",
        "faster_rcnn_R_50_DC5_1x.yaml",
        "faster_rcnn_R_50_DC5_3x.yaml",
        "faster_rcnn_R_50_FPN_1x.yaml",
        "faster_rcnn_R_50_FPN_3x.yaml",
        "faster_rcnn_X_101_32x8d_FPN_3x.yaml",
        "retinanet_R_101_FPN_3x.yaml",
        "retinanet_R_50_FPN_1x.yaml",
        "retinanet_R_50_FPN_3x.yaml",
        "rpn_R_50_C4_1x.yaml",
        "rpn_R_50_FPN_1x.yaml",
    ],
    "COCO-InstanceSegmentation": [
        "mask_rcnn_R_101_C4_3x.yaml",
        "mask_rcnn_R_101_DC5_3x.yaml",
        "mask_rcnn_R_101_FPN_3x.yaml",
        "mask_rcnn_R_50_C4_1x.yaml",
        "mask_rcnn_R_50_C4_3x.yaml",
        "mask_rcnn_R_50_DC5_1x.yaml",
        "mask_rcnn_R_50_DC5_3x.yaml",
        "mask_rcnn_R_50_FPN_1x.yaml",
        "mask_rcnn_R_50_FPN_3x.yaml",
        "mask_rcnn_X_101_32x8d_FPN_3x.yaml",
    ],
    "COCO-Keypoints": [
        "keypoint_rcnn_R_101_FPN_3x.yaml",
        "keypoint_rcnn_R_50_FPN_1x.yaml",
        "keypoint_rcnn_R_50_FPN_3x.yaml",
        "keypoint_rcnn_X_101_32x8d_FPN_3x.yaml",
    ],
    "COCO-PanopticSegmentation": [
        "panoptic_fpn_R_101_3x.yaml",
        "panoptic_fpn_R_50_1x.yaml",
        "panoptic_fpn_R_50_3x.yaml",
    ],
    "Cityscapes": [
        "mask_rcnn_R_50_FPN.yaml",
    ],
    "Detectron1-Comparisons": [
        "faster_rcnn_R_50_FPN_noaug_1x.yaml",
        "keypoint_rcnn_R_50_FPN_1x.yaml",
        "mask_rcnn_R_50_FPN_noaug_1x.yaml",
    ],
    "LVISv0.5-InstanceSegmentation": [
        "mask_rcnn_R_50_FPN_1x.yaml",
        "mask_rcnn_X_101_32x8d_FPN_1x.yaml",
    ],
    "Misc": [
        "cascade_mask_rcnn_R_50_FPN_1x.yaml",
        "cascade_mask_rcnn_R_50_FPN_3x.yaml",
        "cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml",
        "mask_rcnn_R_50_FPN_1x_dconv_c3-c5.yaml",
        "mask_rcnn_R_50_FPN_3x_dconv_c3-c5.yaml",
        "mask_rcnn_R_50_FPN_3x_gn.yaml",
        "mask_rcnn_R_50_FPN_3x_syncbn.yaml",
        "panoptic_fpn_R_101_dconv_cascade_gn_3x.yaml",
        "scratch_mask_rcnn_R_50_FPN_3x_gn.yaml",
        "scratch_mask_rcnn_R_50_FPN_9x_gn.yaml",
        "scratch_mask_rcnn_R_50_FPN_9x_syncbn.yaml",
        "semantic_R_50_FPN_1x.yaml",
    ],
    "PascalVOC-Detection": [
        "faster_rcnn_R_50_C4.yaml",
        "faster_rcnn_R_50_FPN.yaml",
    ]
}

def draw_model_name(image, text):
    """
    在图像上绘制模型名称的矩形框和文本。

    :param image: 输入图像
    :param text: 要绘制的文本
    :return: 带有绘制文本的图像
    """
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.6
    color = (255, 0, 0)
    thickness = 2
    (text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)
    overlay = image.copy()
    cv2.rectangle(overlay, (5, 5), (10 + text_width, 10 + text_height + 20), (255, 255, 255), -1)
    cv2.putText(overlay, text, (10, 30), font, font_scale, color, thickness)
    return cv2.addWeighted(overlay, 0.7, image, 0.3, 0)

def run_inference(args):
    """
    运行推理，处理各个模型并保存输出图像。

    :param args: 命令行参数
    """
    # 设置日志
    setup_logger()
    # 检查设备（CPU或GPU）
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")
    # 加载输入图像
    im = cv2.imread(args.input)
    input_filename = os.path.splitext(os.path.basename(args.input))[0]
    # 遍历所有模型配置
    for category, models in MODEL_CONFIGS.items():
        for model in models:
            if model.endswith('.py'):
                continue  # 跳过 .py 配置文件
            model_path = f"{category}/{model}"
            print(f"Running inference with model: {model_path}")
            
            # 创建配置
            cfg = get_cfg()
            cfg.merge_from_file(model_zoo.get_config_file(model_path))
            cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
            
            # 设置权重和设备
            try:
                cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_path)
            except RuntimeError as e:
                print(f"Skipping {model_path} due to missing model weights.")
                continue
            
            cfg.MODEL.DEVICE = device
            
            # 创建预测器
            predictor = DefaultPredictor(cfg)
            try:
                start_time = time.time()  # 开始计时
                # 进行推理
                outputs = predictor(im)
                
                # 结束计时
                end_time = time.time()
                elapsed_time = end_time - start_time
                print(f"Inference with model {model_path} took {elapsed_time:.2f} seconds")
                
                # 可视化预测结果
                v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
                out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
                
                # 生成输出路径
                model_filename = os.path.splitext(os.path.basename(model))[0]
                output_filename = f"{input_filename}_{model_filename}.jpg"
                output_dir = os.path.join(args.output, category)
                os.makedirs(output_dir, exist_ok=True)
                output_path = os.path.join(output_dir, output_filename)
                
                # 在输出图像上绘制模型名称
                output_image = out.get_image()[:, :, ::-1]
                output_image = draw_model_name(output_image, model_filename)
                
                # 保存输出文件
                cv2.imwrite(output_path, output_image)
                print(f"Output saved to {output_path}")
            
            except Exception as e:
                print(f"Error inferencing with {model_path}: {e}")
                continue

if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="Detectron2 Inference Script")
    parser.add_argument("--input", required=True, help="Path to the input image")
    parser.add_argument("--output", required=True, help="Path to the output directory")
    parser.add_argument("--cache-dir", required=True, help="Path to the cache directory")
    args = parser.parse_args()
    
    # 确保缓存目录存在
    os.makedirs(args.cache_dir, exist_ok=True)
    os.environ["FVCORE_CACHE"] = args.cache_dir
    
    run_inference(args)

代码概述

  1. 包导入：

    • 处理命令行参数、操作系统功能、图像处理、PyTorch、时间管理、Detectron2相关库的导入。

  2. 模型配置：

    • MODEL_CONFIGS 字典存储了各种模型配置文件路径。

  3. 函数 draw_model_name：

    • 在图像上绘制模型名称，用于标识每个生成图像对应的模型。

  4. 函数 run_inference：

    • 设置日志、检测设备、加载输入图像、遍历所有模型配置、创建配置和预测器、进行推理、可视化预测结果、生成输出路径、绘制模型名称并保存输出文件。

  5. 脚本入口：

    • 解析命令行参数，并调用 run_inference 函数开始处理。

使用方法

在命令行输入以下命令，运行脚本进行推理：

python3 infer_allmode.py --input input.jpg --output ./outimgs --cache-dir /liveguard/Detectron2/models/

运行环境

使用Dockerfile构建一个包含NVIDIA的PyTorch和Detectron2的镜像

nvcr.io/nvidia/pytorch:23.01-py3

运行上有任何问题请在评论区留言

fydw_715
关注
  • 10
    点赞
  • 踩
  • 11
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
专栏目录
大模型训练难于上青天？效率超群、易用的“李白”模型库来了
OneFlow深度学习框架
07-14 4421
这个新生代模型库覆盖了 Hugging Face、Megatron-LM、DeepSpeed、FairSeq 这些所有主流 Transformer 库的优点，分布式训练性能一如既往地力压群雄，更重要的是，全局视角编程最大程度降低了分布式使用门槛，让大模型训练飞入寻常百姓家。...
李白：你的模型权重很不错，可惜被我没收了
OneFlow深度学习框架
11-10 7662
大噶好，年更楼主今天想推的是，主打分布式训练的模型库_李白（LiBai）。https://github.com/Oneflow-Inc/libai对于目前市面上的模型库来说，选择实在是太多了，换了一批又一批，眼睛都挑花了，为什么要用LiBai？（如果你觉得LiBai万一某天能用到，或者这篇文章读下来感觉比较开心，可以去GitHub上点赞，如果能三连就更好了。众所周知，GitH...
拥抱Detectron2
木盏
04-08 342
detectron2是FAIR团队在维护的计算机视觉库，持续收集和实现各种SOTA的目标检测/分割算法。类似的库还有CUHK和商汤维护的MMDetection。对我而言，其最大的优势是可以方便快速下载到各种检测算法的pre-trained模型，同时比torchvision更追新一点。对于demo党而言，能够让你如此轻易实现SOTA算法的demo，那也是很苏糊的。 为什么要拥抱Detectron2而不是Detectron呢？ 答：Detectron2完全由PyTorch框架重写了一遍，为了迎合学术界的趋势。
使用Dockerfile构建一个包含NVIDIA的PyTorch和Detectron2的镜像
最新发布
fydw_715的博客
08-01 510
构建镜像：基于指定的Dockerfile和上下文（包括detectron2目录），构建并标签为。查看镜像：构建成功的镜像大小为20.7GB。保存镜像：将镜像保存为文件，便于分发和备份。这样通过这些步骤，用户可以构建、查看和保存一个包含Detectron2的Docker镜像。这个过程详细展示了每一步的执行和输出，对于理解Docker操作和镜像管理具有很好的参考价值。这一段内容说明了如何使用docker run命令来启动一个包含Detectron2环境的Docker容器，并查看正在运行的容器。
【NLP】大模型训练之难，难于上青天？预训练易用、效率超群的「李白」模型库来了！...
fengdu78的博客
07-08 820
机器之心发布机器之心编辑部LiBai（李白）模型库覆盖了 Hugging Face、Megatron-LM、DeepSpeed、FairSeq 这些所有主流 Transformer 库的优点，让大模型训练飞入寻常百姓家。大模型多了去了，告诉我怎么加速？自 2018 年 BERT 诞生，到 GPT-3、ViT 等拥有数以亿计的参数规模的模型不断涌现，AI 模型参数量的爆发式...
WAIC | 知乎CTO李大海：基于AI的智能社区多模态数据融合研究与实践
喜欢打酱油的老鸟
07-26 458
在 WAIC 2021 AI 开发者论坛上，知乎合伙人兼 CTO 李大海发表主题演讲《基于 AI 的智能社区多模态数据融合研究与实践》。在演讲中，他主要介绍了知乎在智能社区时代多模态数据融合中的研究以及实践进展。 以下为李大海在 WAIC 2021 AI 开发者论坛上的演讲内容，机器之心进行了不改变原意的编辑、整理： 我今天的主题是基于 AI 的智能社区多模态数据融合应用实践。 多模态最近很火，刚刚「好未来」的吴总也讲了这个主题。虽然都是多模态，但大家业务场景不同，使用的侧重面不一样。我们.
2021计算机视觉-包揽所有前沿论文源码 -上半年
热门推荐
江南才尽江南山，年少无知年少狂！
12-19 1万+
大家是否遇到过这种情况，就是在工作或者学习的时候，想去找一些方向的网络，但是呢，尴尬的是，老旧的网络里不想要，前沿的网络又不知道有哪些。为了解决大家的这个困扰，本人决定收集2020年大部分前沿的网络相关链接，之后我会： 每周一更新一次（下面是我多年年收集的链接地址）\color{red}{每周一更新一次（下面是我多年年收集的链接地址） }每周一更新一次（下面是我多年年收集的链接地址） 2019计算机视觉-包揽所有前沿论文源码 2020计算机视觉-包揽所有前沿论文源码 有兴趣的朋友可以加微信：a9442847
[转]2020 年最具潜力 44 个顶级开源项目，涵盖 11 类 AI 学习框架、平台（值得收藏）
江州司马
03-31 716
导语：​Github 开源项目技术图 雷锋网 AI 开发者按：工欲善其事必先利其器，这也是大部分开发者在日常工作中最重要开发原则。选择与开发内容相匹配的工具，常常会使我们事半功倍。但面对人工智能的多个领域，如：机器学习、深度学习、NLP等等，多样的工具有时也让我们也无从选择。 就在最近，一个基于 javascript 的可视化库 D3js（treemap 可视化）对 json 文件生成的技术图...
2020 年最具潜力 44 个顶级开源项目，涵盖 11 类 AI 学习框架、平台（值得收藏）...
Python爬虫和机器学习专栏
02-17 808
文末有干货“Python高校”，马上关注真爱，请置顶或星标作者：杨鲤萍工欲善其事必先利其器，这也是大部分开发者在日常工作中最重要开发原则。选择与开发内容相匹配的工具，常常会使我们事半功...
中文预训练ALBERT模型.zip
07-01
2、运行一键运行脚本(GPU方式): 会自动下载模型和所有任务数据并开始运行。 bash run_classifier_clue.sh 执行该一键运行脚本将会自动下载所有任务数据，并为所有任务找到最优模型，然后测试得到提交结果 模型...
python实例一键抓出PPT中的所有文字
02-07
python实例一键抓出PPT中的所有文字python实例一键抓出PPT中的所有文字python实例一键抓出PPT中的所有文字python实例一键抓出PPT中的所有文字python实例一键抓出PPT中的所有文字python实例一键抓出PPT中的所有文字...
arcgis一键拓扑，模型构建
08-08
"arcgis一键拓扑，模型构建"是ArcGIS中的一个重要应用，涉及到地理空间数据的完整性检查和处理流程的自动化。拓扑是地理数据结构的基础，确保了地图要素之间的逻辑关系正确无误。模型构建器则是ArcGIS中用于创建和...
支持最新版本3dMax：VR一键模型代理/还原插件下载
06-20
本文将详细介绍"VR一键模型代理/还原插件"及其在3ds Max中的应用。 该插件名为"VR一键模型代理"，其主要功能是简化3ds Max中V-RayProxy（VR代理）的管理和操作。V-RayProxy是一种优化大型场景的工具，它可以将复杂...
PaddleHub可以便捷地获取PaddlePaddle生态下的预训练模型，完成模型的管理和一键预测.rar
07-05
PaddleHub可以便捷地获取PaddlePaddle生态下的预训练模型，完成模型的管理和一键预测。配合使用Fine-tune API，可以基于大规模预训练模型快速完成迁移学习，让预训练模型能更好地服务于用户特定场景的应用。 模型...
从零开始学习机器学习，掌握AI未来的关键！
qq_42538588的博客
07-25 1258
欢迎来到"人工智能-机器学习基础"！在本文中，我们将一起探索人工智能的奇妙世界。首先，让我们来介绍一下人工智能的概述。人工智能（Artificial Intelligence）是一门研究如何使机器具备智能能力的科学与技术。简单来说，它是能够像人类一样思考、学习和决策的机器。人工智能可以追溯到上个世纪50年代，当时科学家们就开始尝试模仿和理解人类的思维过程。然而，由于当时计算机技术的限制，人工智能的发展并不顺利。直到最近几年，随着计算能力的迅猛增长和大数据的普及，人工智能取得了巨大突破。
Java面试整理，涵盖基础、JVM、线程并发、框架、MySQL、微服务、Redis、中间件、数据结构与算法等。陆续完善中.zip
07-31
Java面试整理，涵盖基础、JVM、线程并发、框架、MySQL、微服务、Redis、中间件、数据结构与算法等。陆续完善中
chromedriver-mac-x64_121.0.6167.184.zip
07-31
chromedriver-mac-x64_121.0.6167.184.zip
【Java面试+Java技术文章汇总】 涵盖大部分Java程序员所需要掌握的核心知识。.zip
07-31
【Java面试+Java技术文章汇总】 涵盖大部分Java程序员所需要掌握的核心知识。
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值