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

news2025/1/24 5:36:48

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

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

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.coloradmin.cn/o/1977962.html

如若内容造成侵权/违法违规/事实不符，请联系多彩编程网进行投诉反馈，一经查实，立即删除！

相关文章

（五）activiti-modeler 编辑器初步优化

（五）activiti-modeler 编辑器初步优化

最终效果： 1..首先去掉顶部的logo，没什么用，还占用空间。 修改modeler.html文件，添加样式： <style type="text/css"> #main-header{display: none; } #main{padding: 0px; } </style> 2.左边组件选择区域太宽了，一般用不到那么宽。 修改editor…
阅读更多...
Linux驱动入门实验班day03-GPIO子系统概述

Linux驱动入门实验班day03-GPIO子系统概述

3.通用框架1——最简单方式1&#xff1a;执行命令cat /sys/kernel/debug/gpio查看串口信息 gpio4对应的下列 方式2&#xff1a; 对于按键GPIO4_14:对应第四组第14个引脚 gpiochip3 ,从96开始&#xff0c; 9614110&#xff1b;
阅读更多...
如何获取公网IP

如何获取公网IP

前言 有些时候我们需要获取电脑或者服务器的公网 IP&#xff0c;例如我们访问的目标地址需要限制 IP 白名单或者限制访问来源&#xff0c;又或者我们使用了代理&#xff0c;想试试有没有生效。要获取公网 IP 无法从电脑或服务器的本地配置中获取&#xff0c;如 ipconfig&#…
阅读更多...
利用 Angular 发挥环境的力量

利用 Angular 发挥环境的力量

一.介绍 您是否曾想过如何在不同的环境中为同一应用设置不同的颜色、标题或 API 调用&#xff1f;可以肯定的是&#xff0c;生产 API 和测试 API 是不同的&#xff0c;应谨慎使用。部署时&#xff0c;我们不会在项目的所有地方手动更改所有 API 调用。不应这样做&#xff0c;因…
阅读更多...
基于Yolov8面部七种表情检测与识别C++模型部署

基于Yolov8面部七种表情检测与识别C++模型部署

表情识别 七种表情识别是一个多学科交叉的研究领域&#xff0c;它结合了心理学、认知科学、计算机视觉和机器学习等学科的知识和技术。 基本概念 表情的定义&#xff1a;表情是人们在情绪体验时面部肌肉活动的结果&#xff0c;是人类情感交流的基本方式之一。基本表情理论&a…
阅读更多...
matplotlib显示opencv读取的图片颜色异常，BGR转RGB的两种方式:cv2.cvtColor与img[:,:,::-1]

matplotlib显示opencv读取的图片颜色异常，BGR转RGB的两种方式:cv2.cvtColor与img[:,:,::-1]

《博主简介》 小伙伴们好&#xff0c;我是阿旭。专注于人工智能、AIGC、python、计算机视觉相关分享研究。 ✌更多学习资源&#xff0c;可关注公-仲-hao:【阿旭算法与机器学习】&#xff0c;共同学习交流~ &#x1f44d;感谢小伙伴们点赞、关注&#xff01; 《------往期经典推…
阅读更多...
github添加ssh密钥，通过ssh方式推送代码

github添加ssh密钥，通过ssh方式推送代码

左手编程&#xff0c;右手年华。大家好&#xff0c;我是一点&#xff0c;关注我&#xff0c;带你走入编程的世界。 公众号&#xff1a;一点sir&#xff0c;关注领取python编程资料 很多人在使用github的时候&#xff0c;如果还是使用https的方式推送代码的话&#xff0c;可能会…
阅读更多...
《LeetCode热题100》---＜5.①普通数组篇五道＞

《LeetCode热题100》---＜5.①普通数组篇五道＞

本篇博客讲解LeetCode热题100道普通数组篇中的五道题 第一道&#xff1a;最大子数组和&#xff08;中等&#xff09; 第二道&#xff1a;合并区间&#xff08;中等&#xff09; 第一道&#xff1a;最大子数组和&#xff08;中等&#xff09; 法一&#xff1a;贪心算法 class So…
阅读更多...
「文件加密软件精选」13个惊艳无比的软件推荐请查收!

「文件加密软件精选」13个惊艳无比的软件推荐请查收!

信息安全是企业生存与发展的基石。 一次不慎的数据泄露&#xff0c;就可能给企业带来难以估量的损失。因此&#xff0c;文件加密成为了保护企业核心资产的重要手段。 某日&#xff0c;两位员工小李和小张在茶水间闲聊&#xff1a; “你知道吗&#xff1f;最近公司开始推广文…
阅读更多...
飞浆OCR模型训练详细教程

飞浆OCR模型训练详细教程

目录 飞浆使用文档一 环境查看1.1 Python 环境1.2 版本选择1.3 飞浆测试 二 数据集的标准备2.1 PPOCRLabelv22.1.1 安装 PaddlePaddle2.1.2 安装与运行 PPOCRLabel2.1.3 数据集划分 三 PaddleOCR 的环境搭建3.1 环境搭建3.2 模型下载 四 模型训练4.1 检测模型训练4.2 识别模型训…
阅读更多...
萱仔环境记录——git的安装流程

萱仔环境记录——git的安装流程

最近由于我有一个大模型的offer&#xff0c;由于我只在实验室的电脑上装了git&#xff0c;我准备在自己的笔记本上本地安装一个git&#xff0c;也给我的一个师弟讲解一下git安装和使用的过程&#xff0c;给我的环境安装章节添砖加瓦。 github是基于git的一个仓库托管平台。 g…
阅读更多...
用Python实现亚马逊Amazon高性能爬虫采集销量信息

用Python实现亚马逊Amazon高性能爬虫采集销量信息

用Python实现亚马逊Amazon高性能爬虫采集销量信息 引言 亚马逊作为全球最大的电商平台&#xff0c;拥有丰富的商品种类和庞大的用户基数。因此&#xff0c;采集亚马逊的销量信息对于市场分析、竞争对手研究以及运营优化有着重要的作用。本文将详细介绍如何用Python实现高性能的…
阅读更多...
机械学习—零基础学习日志（高数20——洛必达法则）

机械学习—零基础学习日志（高数20——洛必达法则）

零基础为了学人工智能&#xff0c;真的开始复习高数 这里讲解一个历史&#xff0c;洛必达法则其实并不是洛必达想出来的&#xff0c;洛必达整理了第一本微积分的书籍&#xff0c;是真正的知识传播者。洛必达法则是洛必达从伯努利哪里买过来的&#xff0c;并结合了莱布尼兹的论…
阅读更多...
本地部署 Llama-3-EvoVLM-JP-v2

本地部署 Llama-3-EvoVLM-JP-v2

本地部署 Llama-3-EvoVLM-JP-v2 0. 引言1. 关于 Llama-3-EvoVLM-JP-v22. 本地部署2-0. 克隆代码2-1. 安装依赖模块2-2. 创建 Web UI2-3.启动 Web UI2-4. 访问 Web UI 0. 引言 Sakana AI 提出了一种称为进化模型合并的方法&#xff0c;并使用该方法创建大规模语言模型&#xff…
阅读更多...
数论——线性同余方程、扩欧求解线性同余方程、线性组合、原根求解

数论——线性同余方程、扩欧求解线性同余方程、线性组合、原根求解

线性同余方程 线性同余方程是形如 的方程&#xff0c;其中a 、b、m 为给定的整数&#xff0c;x 是未知整数。 扩欧求解线性同余方程 void mod_slover(int a, int b, int n) {int d, x, y, x0;d extend_gcd(a, n, x, y);if (b % d ! 0)cout << "no answer";…
阅读更多...
联邦学习研究综述【联邦学习】

联邦学习研究综述【联邦学习】

文章目录 0 前言机器学习两大挑战&#xff1a; 1 什么是联邦学习&#xff1f;联邦学习的一次迭代过程如下&#xff1a;联邦学习技术具有以下几个特点&#xff1a; 2 联邦学习的算法原理目标函数本地目标函数联邦学习的迭代过程 3 联邦学习分类横向联邦学习纵向联邦学习联邦迁移…
阅读更多...
科普文：微服务之Spring Cloud OpenFeign服务调用调用过程分析

科普文：微服务之Spring Cloud OpenFeign服务调用调用过程分析

概叙 Feign和OpenFeign的关系 其实到现在&#xff0c;至少是2018年之后&#xff0c;我们说Feign&#xff0c;或者说OpenFeign其实是一个东西&#xff0c;就是OpenFeign&#xff0c;是2018.12 Netflix停止维护后&#xff0c;Spring cloud整合Netflix生态系统的延续。后面其实都…
阅读更多...
ComfyUI-BrushNet(局部重绘)节点安装及效果、模型下载及详细使用方法✨

ComfyUI-BrushNet(局部重绘)节点安装及效果、模型下载及详细使用方法✨

&#x1f35c;背景介绍 ComfyUI 中BrushNet的节点已经发布了三个月左右的时间了&#xff0c;后来陆续更新了更多的功能和模型接入&#xff0c;整体效果看起来还是不错的。这个节点随着能力的更新&#xff0c;接入了更多的模型&#xff0c;而每个模型默认的名字又比较相似&…
阅读更多...
RoboDK的插件

RoboDK的插件

目录 collision-free-planner&#xff1a; opc-ua&#xff1a; collision-free-planner&#xff1a; RoboDK 的无碰撞规划器插件使用概率路线图 (PRM) 自动在机器人工作空间内创建无碰撞路径。 有关无碰撞规划器的更多信息&#xff0c;请访问我们的 文档。 生成参数无碰撞…
阅读更多...
揭秘：查询大数据信用报告的三大作用

揭秘：查询大数据信用报告的三大作用

相信很多朋友都听说过大数据信用是什么&#xff0c;其实早在几年前&#xff0c;不少的网贷平台都是用人行征信加上大数据信用作为平台风控审核的依据&#xff0c;直到大数据技术的发展&#xff0c;现在不少的银行等机构都将大数据信用作为银行放贷风控审核的重要指标&#xff0…
阅读更多...
推荐文章
最新文章

Copyright @ 2022~2023