英伟达SSD视觉算法分类代码解析

一、官方原代码

#!/usr/bin/env python3
#
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
#

import sys
import argparse

from jetson_inference import imageNet
from jetson_utils import videoSource, videoOutput, cudaFont, Log

# parse the command line
parser = argparse.ArgumentParser(description="Classify a live camera stream using an image recognition DNN.", 
                                 formatter_class=argparse.RawTextHelpFormatter, 
                                 epilog=imageNet.Usage() + videoSource.Usage() + videoOutput.Usage() + Log.Usage())

parser.add_argument("input", type=str, default="", nargs='?', help="URI of the input stream")
parser.add_argument("output", type=str, default="", nargs='?', help="URI of the output stream")
parser.add_argument("--network", type=str, default="googlenet", help="pre-trained model to load (see below for options)")
parser.add_argument("--topK", type=int, default=1, help="show the topK number of class predictions (default: 1)")

try:
	args = parser.parse_known_args()[0]
except:
	print("")
	parser.print_help()
	sys.exit(0)


# load the recognition network
net = imageNet(args.network, sys.argv)

# note: to hard-code the paths to load a model, the following API can be used:
#
# net = imageNet(model="model/resnet18.onnx", labels="model/labels.txt", 
#                 input_blob="input_0", output_blob="output_0")

# create video sources & outputs
input = videoSource(args.input, argv=sys.argv)
output = videoOutput(args.output, argv=sys.argv)
font = cudaFont()

# process frames until EOS or the user exits
while True:
    # capture the next image
    img = input.Capture()

    if img is None: # timeout
        continue  

    # classify the image and get the topK predictions
    # if you only want the top class, you can simply run:
    #   class_id, confidence = net.Classify(img)
    predictions = net.Classify(img, topK=args.topK)

    # draw predicted class labels
    for n, (classID, confidence) in enumerate(predictions):
        classLabel = net.GetClassLabel(classID)
        confidence *= 100.0

        print(f"imagenet:  {confidence:05.2f}% class #{classID} ({classLabel})")

        font.OverlayText(img, text=f"{confidence:05.2f}% {classLabel}", 
                         x=5, y=5 + n * (font.GetSize() + 5),
                         color=font.White, background=font.Gray40)
                         
    # render the image
    output.Render(img)

    # update the title bar
    output.SetStatus("{:s} | Network {:.0f} FPS".format(net.GetNetworkName(), net.GetNetworkFPS()))

    # print out performance info
    net.PrintProfilerTimes()

    # exit on input/output EOS
    if not input.IsStreaming() or not output.IsStreaming():
        break

二、代码解析

代码增加中文注释

#!/usr/bin/env python3
#
# 版权所有 (c) 2020, NVIDIA CORPORATION. 保留所有权利。
#
# 特此免费授予获得此软件和相关文档文件（“软件”）副本的任何人，允许他们在不受限制的情况下处理软件，
# 包括但不限于使用、复制、修改、合并、发布、分发、再许可和/或出售软件副本，并允许提供软件的人
# 这样做，条件如下：
#
# 上述版权声明和本许可声明应包含在软件的所有副本或主要部分中。
#
# 本软件按“原样”提供，不提供任何形式的明示或暗示保证，包括但不限于适销性、
# 适用于特定目的和不侵权的保证。在任何情况下，作者或版权持有人均不对因使用本软件或其他交易，
# 或因使用本软件或其他交易而产生的任何索赔、损害或其他责任负责，无论是在合同诉讼、侵权诉讼还是其他诉讼中。
#

import sys
import argparse

from jetson_inference import imageNet
from jetson_utils import videoSource, videoOutput, cudaFont, Log

# 解析命令行参数
parser = argparse.ArgumentParser(
    description="使用图像识别DNN对实时摄像头流进行分类。",
    formatter_class=argparse.RawTextHelpFormatter,
    epilog=imageNet.Usage() + videoSource.Usage() + videoOutput.Usage() + Log.Usage()
)

parser.add_argument("input", type=str, default="", nargs='?', help="输入流的URI")
parser.add_argument("output", type=str, default="", nargs='?', help="输出流的URI")
parser.add_argument("--network", type=str, default="googlenet", help="要加载的预训练模型（参见下方选项）")
parser.add_argument("--topK", type=int, default=1, help="显示前K个类别预测（默认：1）")

try:
    args = parser.parse_known_args()[0]
except:
    print("")
    parser.print_help()
    sys.exit(0)

# 加载识别网络
net = imageNet(args.network, sys.argv)

# 注意：要硬编码加载模型的路径，可以使用以下API：
# net = imageNet(model="model/resnet18.onnx", labels="model/labels.txt", 
#                input_blob="input_0", output_blob="output_0")

# 创建视频源和输出
input = videoSource(args.input, argv=sys.argv)
output = videoOutput(args.output, argv=sys.argv)
font = cudaFont()

# 处理帧直到输入结束或用户退出
while True:
    # 捕获下一帧图像
    img = input.Capture()

    if img is None:  # 超时
        continue  

    # 对图像进行分类并获取前K个预测
    # 如果只需要最顶层的类别，可以简单地运行：
    #   class_id, confidence = net.Classify(img)
    predictions = net.Classify(img, topK=args.topK)

    # 绘制预测的类别标签
    for n, (classID, confidence) in enumerate(predictions):
        classLabel = net.GetClassLabel(classID)
        confidence *= 100.0

        print(f"imagenet:  {confidence:05.2f}% class #{classID} ({classLabel})")

        font.OverlayText(
            img, 
            text=f"{confidence:05.2f}% {classLabel}", 
            x=5, y=5 + n * (font.GetSize() + 5),
            color=font.White, background=font.Gray40
        )
                         
    # 渲染图像
    output.Render(img)

    # 更新标题栏
    output.SetStatus("{:s} | Network {:.0f} FPS".format(net.GetNetworkName(), net.GetNetworkFPS()))

    # 打印性能信息
    net.PrintProfilerTimes()

    # 输入/输出流结束时退出
    if not input.IsStreaming() or not output.IsStreaming():
        break

这段Python代码是一个使用NVIDIA的Jetson平台进行图像分类的示例程序。代码解析如下：

头部版权声明和许可信息

这部分代码声明了版权信息和软件许可，允许免费使用、复制和分发软件。

导入模块

import sys
import argparse
from jetson_inference import imageNet
from jetson_utils import videoSource, videoOutput, cudaFont, Log

• sys: 处理系统特定的参数和功能。
• argparse: 解析命令行参数。
• jetson_inference和jetson_utils模块用于加载和处理图像分类模型、视频源、视频输出、绘制字体和日志记录。

解析命令行参数

parser = argparse.ArgumentParser(description="Classify a live camera stream using an image recognition DNN.", 
                                 formatter_class=argparse.RawTextHelpFormatter, 
                                 epilog=imageNet.Usage() + videoSource.Usage() + videoOutput.Usage() + Log.Usage())

parser.add_argument("input", type=str, default="", nargs='?', help="URI of the input stream")
parser.add_argument("output", type=str, default="", nargs='?', help="URI of the output stream")
parser.add_argument("--network", type=str, default="googlenet", help="pre-trained model to load (see below for options)")
parser.add_argument("--topK", type=int, default=1, help="show the topK number of class predictions (default: 1)")

try:
	args = parser.parse_known_args()[0]
except:
	print("")
	parser.print_help()
	sys.exit(0)

• 使用argparse模块定义和解析命令行参数，包括输入和输出流的URI、使用的预训练模型和显示前K个预测结果的数量。
• 尝试解析命令行参数，如果解析失败，则显示帮助信息并退出程序。

加载图像分类网络

net = imageNet(args.network, sys.argv)

• 使用imageNet类加载预训练的神经网络模型。

创建视频源和视频输出

input = videoSource(args.input, argv=sys.argv)
output = videoOutput(args.output, argv=sys.argv)
font = cudaFont()

• 使用videoSource类创建视频输入流。
• 使用videoOutput类创建视频输出流。
• 使用cudaFont类创建用于绘制文本的字体。

处理视频帧

while True:
    # capture the next image
    img = input.Capture()

    if img is None: # timeout
        continue  

    # classify the image and get the topK predictions
    predictions = net.Classify(img, topK=args.topK)

    # draw predicted class labels
    for n, (classID, confidence) in enumerate(predictions):
        classLabel = net.GetClassLabel(classID)
        confidence *= 100.0

        print(f"imagenet:  {confidence:05.2f}% class #{classID} ({classLabel})")

        font.OverlayText(img, text=f"{confidence:05.2f}% {classLabel}", 
                         x=5, y=5 + n * (font.GetSize() + 5),
                         color=font.White, background=font.Gray40)
                         
    # render the image
    output.Render(img)

    # update the title bar
    output.SetStatus("{:s} | Network {:.0f} FPS".format(net.GetNetworkName(), net.GetNetworkFPS()))

    # print out performance info
    net.PrintProfilerTimes()

    # exit on input/output EOS
    if not input.IsStreaming() or not output.IsStreaming():
        break

• 使用input.Capture()捕获下一帧图像。
• 如果捕获失败，则继续下一次循环。
• 使用net.Classify()对图像进行分类，并获取前K个预测结果。
• 使用font.OverlayText()在图像上绘制预测的类别标签和置信度。
• 使用output.Render()渲染图像。
• 使用output.SetStatus()更新输出窗口的标题栏，显示网络名称和帧率。
• 使用net.PrintProfilerTimes()打印性能信息。
• 如果输入或输出流停止，则退出循环。

总体来说，这段代码实现了一个实时图像分类系统，使用Jetson平台的预训练模型对视频流中的每一帧进行分类，并在输出视频中叠加预测结果。

# 三、imageNet类有哪些对象和方法
`imageNet`类是Jetson Inference库的一部分，用于处理图像分类任务。以下是`imageNet`类的主要对象和方法：

## 对象属性
- `NetworkType`: 预定义的网络类型枚举。
- `model`: 加载的模型文件路径。
- `labels`: 加载的标签文件路径。
- `input_blob`: 输入层名称。
- `output_blob`: 输出层名称。

## 主要方法

1. **`__init__`**:
   ```python
   def __init__(self, network="googlenet", argv=[], model=None, labels=None, input_blob="input_0", output_blob="output_0"):

• 初始化图像分类网络。
• 参数包括网络类型、命令行参数、模型文件路径、标签文件路径、输入层名称和输出层名称。

2. Classify:

   def Classify(self, img, topK=1):
   

   • 对给定的图像进行分类。
   • 返回预测结果的列表，每个结果包括类别ID和置信度。

3. GetClassLabel:

   def GetClassLabel(self, classID):
   

   • 返回指定类别ID的标签。

4. GetNetworkName:

   def GetNetworkName(self):
   

   • 返回网络的名称。

5. GetNetworkFPS:

   def GetNetworkFPS(self):
   

   • 返回网络处理帧率（FPS）。

6. PrintProfilerTimes:

   def PrintProfilerTimes(self):
   

   • 打印网络的性能分析信息。

7. Usage:

   @staticmethod
   def Usage():
   

   • 返回类的用法说明，通常用于命令行帮助信息。

使用示例

以下是如何使用imageNet类的简单示例：

from jetson_inference import imageNet
from jetson_utils import loadImage

# 初始化图像分类网络
net = imageNet("googlenet")

# 加载图像
img = loadImage("example.jpg")

# 对图像进行分类
class_id, confidence = net.Classify(img)

# 获取类别标签
class_label = net.GetClassLabel(class_id)

print(f"Image is classified as {class_label} with {confidence * 100:.2f}% confidence")

这个示例展示了如何初始化一个imageNet对象，加载一张图像，并对其进行分类，最后打印分类结果和置信度。

三、使用示例

 python3 imagenet.py /dev/video0 display://0 --network=googlenet

四、训练自己的分类模型

以下是包含中文注释的SSD训练代码示例：

1. 安装TAO Toolkit

确保在具有NVIDIA GPU的系统上安装了Docker和NVIDIA Container Toolkit。

2. 拉取TAO Toolkit Docker容器

docker pull nvcr.io/nvidia/tao/tao-toolkit-tf:v3.21.11-tf1.15.5-py3

3. 准备数据

准备训练和验证数据，数据应按照Kitti或Pascal VOC格式组织，包含图像文件和对应的标注文件。

4. 创建SSD配置文件

以下是SSD配置文件的示例，并包含中文注释：

random_seed: 42  # 随机种子，用于确保实验的可重复性
dataset_config {
  data_sources: {
    label_directory_path: "/path/to/labels"  # 训练数据的标签路径
    image_directory_path: "/path/to/images"  # 训练数据的图像路径
  }
  validation_data_sources: {
    label_directory_path: "/path/to/val_labels"  # 验证数据的标签路径
    image_directory_path: "/path/to/val_images"  # 验证数据的图像路径
  }
}
model_config {
  pretrained_model_file: "/path/to/pretrained/model"  # 预训练模型文件路径
  num_layers: 18  # 模型的层数
  all_proposals: 200  # 所有提议框的数量
}
train_config {
  batch_size: 8  # 批次大小
  learning_rate: 0.001  # 学习率
  num_epochs: 80  # 训练轮数
  augmentations: {
    horizontal_flip: true  # 是否进行水平翻转数据增强
    vertical_flip: false  # 是否进行垂直翻转数据增强
  }
}

5. 运行训练

使用以下命令运行训练任务，并包含中文注释：

docker run --gpus all -v /path/to/your/data:/data -v /path/to/your/config:/config -v /path/to/your/output:/output nvcr.io/nvidia/tao/tao-toolkit-tf:v3.21.11-tf1.15.5-py3 ssd train \
  -e /config/ssd_config.yaml \  # 配置文件路径
  -r /output/experiment_dir \  # 实验输出目录
  -k $API_KEY  # TAO Toolkit的API密钥

• --gpus all: 使用所有可用的GPU。
• -v /path/to/your/data:/data: 将本地数据目录挂载到容器内的/data路径。
• -v /path/to/your/config:/config: 将本地配置文件目录挂载到容器内的/config路径。
• -v /path/to/your/output:/output: 将本地输出目录挂载到容器内的/output路径。
• -e /config/ssd_config.yaml: 指定配置文件。
• -r /output/experiment_dir: 指定实验输出目录。
• -k $API_KEY: 指定TAO Toolkit的API密钥。

6. 导出模型

训练完成后，使用以下命令导出模型，并包含中文注释：

docker run --gpus all -v /path/to/your/output:/output nvcr.io/nvidia/tao/tao-toolkit-tf:v3.21.11-tf1.15.5-py3 ssd export \
  -m /output/experiment_dir/model.tlt \  # 输入的TAO模型路径
  -o /output/experiment_dir/model.etlt \  # 输出的ETLT模型路径
  -k $API_KEY  # TAO Toolkit的API密钥

总结

通过TAO Toolkit，你可以方便地对SSD目标检测模型进行训练。准备数据、配置训练参数并运行训练命令，可以帮助你快速训练自定义的目标检测模型并进行部署。详细的指南和更多高级功能可以参考TAO Toolkit的官方文档。

这样，代码和配置文件中都增加了中文注释，便于理解和使用。