研究の序文
この論文では、ツールに依存せずに Yolov7-OBB アルゴリズムの TensorRT 高速推論を実装しますtorch2trt。
ソースコードのダウンロード
https://github.com/Egrt/yolov7-obb
気に入ったらスターを注文できます。
TensorRT をインストールする
1. TensorRT の概要
公式ウェブサイトのリンク: https://developer.nvidia.com/tensorrt
NVIDIA ® TensorRT™ は、トレーニングされた深層学習モデルを最適化し、高性能の推論を可能にする SDK です。TensorRT には、トレーニングされた深層学習モデル用の深層学習推論オプティマイザーと、実行用のランタイムが含まれています。選択したフレームワークで深層学習モデルをトレーニングした後、TensorRT を使用すると、より高いスループットとより低いレイテンシでモデルを実行できるようになります。
公式のTensorRT紹介によると、これはトレーニング済みモデルTensorRT用の SDK であり、これを通じて NVIDIA デバイス上で高性能の推論を実行できます。トレーニング済みモデルにどのような最適化が行われるかについては、次の図に示すように、公式 Web サイトの図を参照してください。TensorRTTensorRT
まとめると、主に以下の6点になります。
Reduced Precision: モデルの推論速度を向上させるために、モデルをINT8または のデータ型に量子化しますFP16(精度を維持するか、わずかに低下することを前提としています)。Layer and Tensor Fusion: 複数のレイヤー構造 (水平方向と垂直方向を含む) を融合することで、GPU メモリと帯域幅を最適化します。Kernel Auto-Tuning: 現在使用されている GPU プラットフォームに応じて、最適なデータ レイヤーとアルゴリズムを選択します。Dynamic Tensor Memory: メモリ フットプリントを最小限に抑え、テンソル用にメモリを効率的に再利用します。Multi-Stream Execution: スケーラブルな設計を使用して、複数の入力ストリームを並列処理します。Time Fusion: 動的に生成されたカーネルを使用して、時間ステップとともに変化する RNN ネットワークを最適化します。
2. TensorRT をダウンロードする
公式 Web サイトhttps://developer.nvidia.com/nvidia-tensorrt-8x-download
にアクセスして、対応するバージョンを見つけます。ここでは選択します。
zip アーカイブをダウンロードして解凍します。
3. TensorRT のインストール
次の手順を任意の順序で実行します。
- TensorRT-8.4.3.1\bin の内容を C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3\bin にコピーします。
- TensorRTのインクルードフォルダをCUDAのインクルードフォルダにコピーする
- TensorRT-8.4.3.1\lib フォルダー内の lib ファイルを CUDA の lib フォルダーにコピーし、dll ファイルを CUDA の bin フォルダーにコピーします。
以下の図に示すように、pip install xxx.whl を使用して TensorRT-8.4.3.1 をフォルダーにインストールします。
Python を使用して、インストールが成功したかどうかを確認します。
TensorRTライブラリをインストールした後、モデル変換pycudaを開始できます。Yolov7-OBB
モデル変換
pthまずモデルをONNXモデルに変換し、次にONNXモデルをengineモデルに変換する必要があります。
ONNX モデルのエクスポート
まず、 yolov7-obb または yolov7-tiny-obb のプロジェクト ルート ディレクトリでそれを見つけpredict.py、その中でモードを変更しますexport_onnx。エクスポートされた onnx ファイルは次の場所にあります。model_data/models.onnx
エンジンモデルのエクスポート
次に、次のコードを使用してONNXモデルをengineモデルに変換します。
'''
Description:
Author: Egrt
Date: 2023-03-29 12:12:45
LastEditors: Egrt
LastEditTime: 2023-05-26 20:24:28
'''
import tensorrt as trt
def build_engine(onnx_file_path, engine_file_path, half=False):
"""Takes an ONNX file and creates a TensorRT engine to run inference with"""
logger = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(logger)
config = builder.create_builder_config()
config.max_workspace_size = 4 * 1 << 30
flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
network = builder.create_network(flag)
parser = trt.OnnxParser(network, logger)
if not parser.parse_from_file(str(onnx_file_path)):
raise RuntimeError(f'failed to load ONNX file: {
onnx_file_path}')
half &= builder.platform_has_fast_fp16
if half:
config.set_flag(trt.BuilderFlag.FP16)
with builder.build_engine(network, config) as engine, open(engine_file_path, 'wb') as t:
t.write(engine.serialize())
return engine_file_path
if __name__ =="__main__":
onnx_path = 'model_data/yolov7_obb.onnx'
engine_path = 'model_data/yolov7_obb.engine'
build_engine(onnx_path, engine_path, True)
エンジン モデルを取得したら、numpy前処理と後処理を使用して予測を実現する必要があります。
歪みのない画像のサイズ変更
まず、絵の歪みのなさを実現する必要がありますがresize、Bガイド版ではImageクラス描画画像とcv2読み出し画像を使用しており、2種類の画像を変換するため前処理や読み込み時間の増加につながります。後処理がfps大幅に軽減されます。したがって、この記事はすべて実装に基づいており、cv2歪みのない画像のresize修正は次のようになります。
#---------------------------------------------------#
# Image版本
#---------------------------------------------------#
def resize_image(image, size, letterbox_image):
iw, ih = image.size
w, h = size
if letterbox_image:
scale = min(w/iw, h/ih)
nw = int(iw*scale)
nh = int(ih*scale)
image = image.resize((nw,nh), Image.BICUBIC)
new_image = Image.new('RGB', size, (128,128,128))
new_image.paste(image, ((w-nw)//2, (h-nh)//2))
else:
new_image = image.resize((w, h), Image.BICUBIC)
return new_image
#---------------------------------------------------#
# cv2版本
#---------------------------------------------------#
def resize_image(image, size, letterbox_image):
ih, iw, _ = image.shape
h, w = size
if letterbox_image:
scale = min(w/iw, h/ih)
nw = int(iw * scale)
nh = int(ih * scale)
image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_CUBIC)
new_image = np.zeros((h, w, 3), dtype=np.uint8)
new_image[...] = 128
new_image[(h-nh)//2:(h-nh)//2+nh, (w-nw)//2:(w-nw)//2+nw, :] = image
else:
new_image = cv2.resize(image, (w, h), interpolation=cv2.INTER_CUBIC)
return new_image
resize_imageこの関数は、画像、希望のサイズ、および画像を操作するかどうかを示すブール値を受け取りますletterboxing。letterbox_imageの場合、True関数はアスペクト比を維持しながら画像のサイズを変更し、灰色の背景を持つ新しい画像の上に貼り付けて、空のスペースを埋めます。letterbox_image「はい」の場合False、関数は単に画像を希望のサイズに変更します。この関数はcv2、ライブラリを使用して画像をサイズ変更して貼り付けます。具体的には、 メソッドを使用してresize画像のサイズを変更し、ゼロと配列スライス メソッドを使用して灰色の背景を持つ新しい画像を作成し、サイズ変更した画像を新しい画像に貼り付けます。
デコード
デコード部分については、オリジナル版では torch を使用して実装していますが、この記事では numpy を使用して実装していますtorch.repeat()。np.title()
非最大抑制
非最大値抑制のオリジナル版はcudaでコンパイルしたライブラリを使って実装されていますが、インストールの手間が省けて推論部分で利用できます なお、cv2.dnn.NMSBoxesRotated(bboxes, scores, conf_thres, nms_thres)この関数で入力される回転ボックスのフォーマットはトレーニング バージョンとは少し異なります。
cv2.dnn.NMSBoxesRotatedこれは、OpenCV回転フレームの非最大値抑制 (NMS) を実行する機能です。従来の NMS とは異なり、回転ボックスをサポートしており、回転物体の検出に使用できます。
この関数の入力パラメータには次のものがあります:
bboxes: (N, 5) の形状を持つ numpy 配列、各行は (x, y, w, h, angle) の 5 つの値を含む回転ボックスを表します; : 形状(N
scores, ) の各要素は、対応するスピン ボックスのスコアを表します。
score_threshold: スコアのしきい値、このしきい値未満のスピン ボックスは無視されます。
nms_threshold: NMS のしきい値、このしきい値を超えるスピン ボックスは抑制されます。
eta: a非最大抑制アルゴリズムの交差比率のしきい値を調整するために使用される拡張係数:
selected_indices保存された回転ボックスのインデックスを格納するために使用される形状 (N,) を持つ数多の配列。
この関数の出力は、形状 (K) の numpy 配列です。ここで、K は保持されているスピン ボックスの数であり、
各要素は対応するスピン ボックスのインデックスを表します。出力配列の要素の順序は入力配列の要素の順序と異なる場合があるため、出力配列の順序に従って入力配列をフィルタリングする必要があることに注意してください。
元の実装:
#------------------------------------------#
# 使用官方自带的非极大抑制会速度更快一些！
# 筛选出一定区域内，属于同一种类得分最大的框
#------------------------------------------#
_, keep = obb_nms(
detections_class[:, :5],
detections_class[:, 5] * detections_class[:, 6],
nms_thres
)
max_detections = detections_class[keep]
cv2の実装
#------------------------------------------#
# 使用cv2.dnn.NMSBoxesRotated进行非极大抑制
#------------------------------------------#
bboxes = [[[bbox[0], bbox[1]], [bbox[2], bbox[3]], bbox[4]* 180 / np.pi] for bbox in detections_class[:, :5]]
scores = [float(score) for score in detections_class[:, 5] * detections_class[:, 6]]
indices = cv2.dnn.NMSBoxesRotated(bboxes, scores, conf_thres, nms_thres)
max_detections = detections_class[indices.flatten()]
完全な tensorrt 推論コードは次のとおりです。
'''
Author: [egrt]
Date: 2023-03-26 09:39:21
LastEditors: Egrt
LastEditTime: 2023-03-29 14:06:57
Description:
'''
import colorsys
import numpy as np
import time
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import cv2
def resize_image(image, size, letterbox_image):
ih, iw = image.shape[:2]
h, w = size
if letterbox_image:
scale = min(w/iw, h/ih)
nw = int(iw*scale)
nh = int(ih*scale)
image = cv2.resize(image, (nw,nh), interpolation=cv2.INTER_CUBIC)
new_image = 128 * np.ones((h, w, 3), dtype=np.uint8)
new_image[(h-nh)//2:(h-nh)//2+nh, (w-nw)//2:(w-nw)//2+nw, :] = image
else:
new_image = cv2.resize(image, (w, h), interpolation=cv2.INTER_CUBIC)
return new_image
def preprocess_input(image):
image /= 255.0
return image
class DecodeBox():
def __init__(self, anchors, num_classes, input_shape, anchors_mask = [[6,7,8], [3,4,5], [0,1,2]]):
super(DecodeBox, self).__init__()
self.anchors = anchors
self.num_classes = num_classes
self.bbox_attrs = 5 + num_classes
self.input_shape = input_shape
#-----------------------------------------------------------#
# 13x13的特征层对应的anchor是[142, 110],[192, 243],[459, 401]
# 26x26的特征层对应的anchor是[36, 75],[76, 55],[72, 146]
# 52x52的特征层对应的anchor是[12, 16],[19, 36],[40, 28]
#-----------------------------------------------------------#
self.anchors_mask = anchors_mask
def decode_box(self, inputs):
outputs = []
for i, input in enumerate(inputs):
#-----------------------------------------------#
# 输入的input一共有三个，他们的shape分别是
# batch_size = 1
# batch_size, 3 * (5 + 1 + 80), 20, 20
# batch_size, 255, 40, 40
# batch_size, 255, 80, 80
#-----------------------------------------------#
batch_size = input.shape[0]
input_height = input.shape[2]
input_width = input.shape[3]
#-----------------------------------------------#
# 输入为640x640时
# stride_h = stride_w = 32、16、8
#-----------------------------------------------#
stride_h = self.input_shape[0] / input_height
stride_w = self.input_shape[1] / input_width
#-------------------------------------------------#
# 此时获得的scaled_anchors大小是相对于特征层的
#-------------------------------------------------#
scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors[self.anchors_mask[i]]]
#-----------------------------------------------#
# 输入的input一共有三个，他们的shape分别是
# batch_size, 3, 20, 20, 85
# batch_size, 3, 40, 40, 85
# batch_size, 3, 80, 80, 85
#-----------------------------------------------#
prediction = input.reshape(batch_size, len(self.anchors_mask[i]), self.bbox_attrs, input_height, input_width)
prediction = np.transpose(prediction, (0, 1, 3, 4, 2))
#-----------------------------------------------#
# 先验框的中心位置的调整参数
#-----------------------------------------------#
x = 1 / (1 + np.exp(-prediction[..., 0]))
y = 1 / (1 + np.exp(-prediction[..., 1]))
#-----------------------------------------------#
# 先验框的宽高调整参数
#-----------------------------------------------#
w = 1 / (1 + np.exp(-prediction[..., 2]))
h = 1 / (1 + np.exp(-prediction[..., 3]))
#-----------------------------------------------#
# 获得置信度，是否有物体
#-----------------------------------------------#
conf = 1 / (1 + np.exp(-prediction[..., 4]))
#-----------------------------------------------#
# 种类置信度
#-----------------------------------------------#
pred_cls = 1 / (1 + np.exp(-prediction[..., 5:]))
#----------------------------------------------------------#
# 生成网格，先验框中心，网格左上角
# batch_size,3,20,20
#----------------------------------------------------------#
grid_x = np.linspace(0, input_width - 1, input_width)
grid_x = np.tile(grid_x, (input_height, 1))
grid_x = np.tile(grid_x, (batch_size * len(self.anchors_mask[i]), 1, 1)).reshape(x.shape)
grid_y = np.linspace(0, input_height - 1, input_height)
grid_y = np.tile(grid_y, (input_width, 1)).T
grid_y = np.tile(grid_y, (batch_size * len(self.anchors_mask[i]), 1, 1)).reshape(y.shape)
scaled_anchors = np.array(scaled_anchors)
anchor_w = scaled_anchors[:, 0:1]
anchor_h = scaled_anchors[:, 1:2]
anchor_w = np.tile(anchor_w, (batch_size, 1)).reshape(1, -1, 1)
anchor_w = np.tile(anchor_w, (1, 1, input_height * input_width)).reshape(w.shape)
anchor_h = np.tile(anchor_h, (batch_size, 1)).reshape(1, -1, 1)
anchor_h = np.tile(anchor_h, (1, 1, input_height * input_width)).reshape(h.shape)
#----------------------------------------------------------#
# 利用预测结果对先验框进行调整
# 首先调整先验框的中心，从先验框中心向右下角偏移
# 再调整先验框的宽高。
# x 0 ~ 1 => 0 ~ 2 => -0.5, 1.5 => 负责一定范围的目标的预测
# y 0 ~ 1 => 0 ~ 2 => -0.5, 1.5 => 负责一定范围的目标的预测
# w 0 ~ 1 => 0 ~ 2 => 0 ~ 4 => 先验框的宽高调节范围为0~4倍
# h 0 ~ 1 => 0 ~ 2 => 0 ~ 4 => 先验框的宽高调节范围为0~4倍
#----------------------------------------------------------#
pred_boxes = np.zeros(prediction[..., :4].shape, dtype='float32')
pred_boxes[..., 0] = x * 2. - 0.5 + grid_x
pred_boxes[..., 1] = y * 2. - 0.5 + grid_y
pred_boxes[..., 2] = (w * 2) ** 2 * anchor_w
pred_boxes[..., 3] = (h * 2) ** 2 * anchor_h
#----------------------------------------------------------#
# 将输出结果归一化成小数的形式
#----------------------------------------------------------#
_scale = np.array([input_width, input_height, input_width, input_height]).astype('float32')
output = np.concatenate((pred_boxes.reshape(batch_size, -1, 4) / _scale, conf.reshape(batch_size, -1, 1),
pred_cls.reshape(batch_size, -1, self.num_classes)), -1)
output = np.concatenate((pred_boxes.reshape(batch_size, -1, 4) / _scale, conf.reshape(batch_size, -1, 1),
pred_cls.reshape(batch_size, -1, self.num_classes)), -1)
outputs.append(output)
return outputs
def non_max_suppression(self, prediction, num_classes, input_shape, image_shape, letterbox_image, conf_thres=0.5, nms_thres=0.4):
#----------------------------------------------------------#
# prediction [batch_size, num_anchors, 85]
#----------------------------------------------------------#
output = [None for _ in range(len(prediction))]
for i, image_pred in enumerate(prediction):
#----------------------------------------------------------#
# 对种类预测部分取max。
# class_conf [num_anchors, 1] 种类置信度
# class_pred [num_anchors, 1] 种类
#----------------------------------------------------------#
class_conf = np.max(image_pred[:, 5:5 + num_classes], axis=1, keepdims=True)
class_pred = np.argmax(image_pred[:, 5:5 + num_classes], axis=1)
class_pred = np.expand_dims(class_pred, axis=1)
#----------------------------------------------------------#
# 利用置信度进行第一轮筛选
#----------------------------------------------------------#
conf_mask = (image_pred[:, 4] * class_conf[:, 0] >= conf_thres).squeeze()
#----------------------------------------------------------#
# 根据置信度进行预测结果的筛选
#----------------------------------------------------------#
image_pred = image_pred[conf_mask]
class_conf = class_conf[conf_mask]
class_pred = class_pred[conf_mask]
if not image_pred.shape[0]:
continue
#-------------------------------------------------------------------------#
# detections [num_anchors, 8]
# 8的内容为：x, y, w, h, angle, obj_conf, class_conf, class_pred
#-------------------------------------------------------------------------#
detections = np.concatenate((image_pred[:, :5], class_conf, class_pred), 1)
#------------------------------------------#
# 获得预测结果中包含的所有种类
#------------------------------------------#
unique_labels = np.unique(detections[:, -1])
for c in unique_labels:
#------------------------------------------#
# 获得某一类得分筛选后全部的预测结果
#------------------------------------------#
detections_class = detections[detections[:, -1] == c]
#------------------------------------------#
# 使用cv2.dnn.NMSBoxesRotated进行非极大抑制
#------------------------------------------#
bboxes = [[bbox[0], bbox[1], bbox[2], bbox[3]] for bbox in detections_class[:, :4]]
scores = [float(score) for score in detections_class[:, 4] * detections_class[:, 5]]
indices = cv2.dnn.NMSBoxes(bboxes, scores, conf_thres, nms_thres)
max_detections = detections_class[indices.flatten()]
# Add max detections to outputs
output[i] = max_detections if output[i] is None else np.concatenate((output[i], max_detections))
if output[i] is not None:
output[i][:, :4] = self.yolo_correct_boxes(output[i], input_shape, image_shape, letterbox_image)
return output
def yolo_correct_boxes(self, output, input_shape, image_shape, letterbox_image):
#-----------------------------------------------------------------#
# 把y轴放前面是因为方便预测框和图像的宽高进行相乘
#-----------------------------------------------------------------#
box_xy = output[..., 0:2]
box_wh = output[..., 2:4]
box_yx = box_xy[..., ::-1]
box_hw = box_wh[..., ::-1]
input_shape = np.array(input_shape)
image_shape = np.array(image_shape)
if letterbox_image:
#-----------------------------------------------------------------#
# 这里求出来的offset是图像有效区域相对于图像左上角的偏移情况
# new_shape指的是宽高缩放情况
#-----------------------------------------------------------------#
new_shape = np.round(image_shape * np.min(input_shape/image_shape))
offset = (input_shape - new_shape)/2./input_shape
scale = input_shape/new_shape
box_yx = (box_yx - offset) * scale
box_hw *= scale
box_xy = box_yx[:, ::-1]
box_hw = box_wh[:, ::-1]
boxes = np.concatenate([box_xy, box_wh], axis=-1)
boxes[:, [0, 2]] *= image_shape[1]
boxes[:, [1, 3]] *= image_shape[0]
return boxes
class YOLO(object):
_defaults = {
#--------------------------------------------------------------------------#
# 使用自己训练好的模型进行预测一定要修改model_path和classes_path！
# model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
#
# 训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
# 验证集损失较低不代表mAP较高，仅代表该权值在验证集上泛化性能较好。
# 如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
#--------------------------------------------------------------------------#
"model_path" : 'model_data/yolov7-tiny.engine',
#---------------------------------------------------------------------#
# 输入图片的大小，必须为32的倍数。
#---------------------------------------------------------------------#
"input_shape" : [640, 640],
#---------------------------------------------------------------------#
# 只有得分大于置信度的预测框会被保留下来
#---------------------------------------------------------------------#
"confidence" : 0.5,
#---------------------------------------------------------------------#
# 非极大抑制所用到的nms_iou大小
#---------------------------------------------------------------------#
"nms_iou" : 0.3,
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化YOLO
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
for name, value in kwargs.items():
setattr(self, name, value)
self._defaults[name] = value
#---------------------------------------------------#
# 获得种类和先验框的数量
#---------------------------------------------------#
self.class_names = ['car', 'truck', 'van', 'bus', 'pedestrian']
self.num_classes = len(self.class_names)
self.anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
self.anchors = np.array([[ 12, 16],[ 19, 36],[ 40, 28],
[ 36, 75],[ 76, 55],[ 72, 146],
[142, 110],[192, 243],[459, 401]])
self.num_anchors = 9
self.bbox_util = DecodeBox(self.anchors, self.num_classes, (self.input_shape[0], self.input_shape[1]), self.anchors_mask)
#---------------------------------------------------#
# 画框设置不同的颜色
#---------------------------------------------------#
hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
self.generate()
#---------------------------------------------------#
# 生成模型
#---------------------------------------------------#
def generate(self):
#---------------------------------------------------#
# 建立yolo模型，载入yolo模型的权重
#---------------------------------------------------#
engine = self.load_engine(self.model_path)
self.context = engine.create_execution_context()
self.inputs, self.outputs, self.bindings = [], [], []
self.stream = cuda.Stream()
for binding in engine:
size = engine.get_binding_shape(binding)
dtype = trt.nptype(engine.get_binding_dtype(binding))
host_mem = np.empty(size, dtype=dtype)
host_mem = np.ascontiguousarray(host_mem)
device_mem = cuda.mem_alloc(host_mem.nbytes)
self.bindings.append(int(device_mem))
if engine.binding_is_input(binding):
self.inputs.append({
'host': host_mem, 'device': device_mem})
else:
self.outputs.append({
'host': host_mem, 'device': device_mem})
def load_engine(self, engine_path):
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read())
def forward(self, img):
self.inputs[0]['host'] = np.ravel(img)
# transfer data to the gpu
for inp in self.inputs:
cuda.memcpy_htod_async(inp['device'], inp['host'], self.stream)
# run inference
self.context.execute_async_v2(
bindings=self.bindings,
stream_handle=self.stream.handle)
# fetch outputs from gpu
for out in self.outputs:
cuda.memcpy_dtoh_async(out['host'], out['device'], self.stream)
# synchronize stream
self.stream.synchronize()
return [out['host'] for out in self.outputs]
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image):
#---------------------------------------------------#
# 计算输入图片的高和宽
#---------------------------------------------------#
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_data = resize_image(image, (self.input_shape[1], self.input_shape[0]), True)
#---------------------------------------------------------#
# 添加上batch_size维度
# h, w, 3 => 3, h, w => 1, 3, h, w
#---------------------------------------------------------#
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
outputs = self.forward(image_data)[::-1]
outputs = self.bbox_util.decode_box(outputs)
#---------------------------------------------------------#
# 将预测框进行堆叠，然后进行非极大抑制
#---------------------------------------------------------#
results = self.bbox_util.non_max_suppression(np.concatenate(outputs, axis=1), self.num_classes, self.input_shape,
image_shape, True, conf_thres = self.confidence, nms_thres = self.nms_iou)
if results[0] is None:
return image
top_label = np.array(results[0][:, 6], dtype = 'int32')
top_conf = results[0][:, 4] * results[0][:, 5]
top_boxes = results[0][:, :4]
#---------------------------------------------------------#
# 图像绘制
#---------------------------------------------------------#
for i, c in list(enumerate(top_label)):
predicted_class = self.class_names[int(c)]
box = top_boxes[i]
score = top_conf[i]
xc, yc, w, h = box.astype('int32')
left, top, right, bottom = int(xc-w/2), int(yc-h/2), int(xc+w/2), int(yc+h/2)
cv2.rectangle(image, (left, top), (right, bottom), self.colors[int(c)], thickness=2)
label = '{} {:.2f}'.format(predicted_class, score)
cv2.putText(image, label, (left, top-10), cv2.FONT_HERSHEY_SIMPLEX, 1, self.colors[int(c)], thickness=1)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
return image
if __name__=='__main__':
yolo = YOLO()
#----------------------------------------------------------------------------------------------------------#
# mode用于指定测试的模式：
# 'predict' 表示单张图片预测，如果想对预测过程进行修改，如保存图片，截取对象等，可以先看下方详细的注释
# 'video' 表示视频检测，可调用摄像头或者视频进行检测，详情查看下方注释。
# 'fps' 表示测试fps，使用的图片是img里面的street.jpg，详情查看下方注释。
# 'dir_predict' 表示遍历文件夹进行检测并保存。默认遍历img文件夹，保存img_out文件夹，详情查看下方注释。
#----------------------------------------------------------------------------------------------------------#
mode = "video"
#----------------------------------------------------------------------------------------------------------#
# video_path 用于指定视频的路径，当video_path=0时表示检测摄像头
# 想要检测视频，则设置如video_path = "xxx.mp4"即可，代表读取出根目录下的xxx.mp4文件。
# video_save_path 表示视频保存的路径，当video_save_path=""时表示不保存
# 想要保存视频，则设置如video_save_path = "yyy.mp4"即可，代表保存为根目录下的yyy.mp4文件。
# video_fps 用于保存的视频的fps
#
# video_path、video_save_path和video_fps仅在mode='video'时有效
# 保存视频时需要ctrl+c退出或者运行到最后一帧才会完成完整的保存步骤。
#----------------------------------------------------------------------------------------------------------#
video_path = "D:/Notebook/traffic-ai/dataset/video-03.mp4"
video_save_path = ""
video_fps = 25.0
#----------------------------------------------------------------------------------------------------------#
# test_interval 用于指定测量fps的时候，图片检测的次数。理论上test_interval越大，fps越准确。
# fps_image_path 用于指定测试的fps图片
#
# test_interval和fps_image_path仅在mode='fps'有效
#----------------------------------------------------------------------------------------------------------#
test_interval = 100
fps_image_path = "img/test.jpg"
#-------------------------------------------------------------------------#
# dir_origin_path 指定了用于检测的图片的文件夹路径
# dir_save_path 指定了检测完图片的保存路径
#
# dir_origin_path和dir_save_path仅在mode='dir_predict'时有效
#-------------------------------------------------------------------------#
dir_origin_path = "img/"
dir_save_path = "img_out/"
if mode == "predict":
'''
1、如果想要进行检测完的图片的保存，利用r_image.save("img.jpg")即可保存，直接在predict.py里进行修改即可。
2、如果想要获得预测框的坐标，可以进入yolo.detect_image函数，在绘图部分读取top，left，bottom，right这四个值。
3、如果想要利用预测框截取下目标，可以进入yolo.detect_image函数，在绘图部分利用获取到的top，left，bottom，right这四个值
在原图上利用矩阵的方式进行截取。
4、如果想要在预测图上写额外的字，比如检测到的特定目标的数量，可以进入yolo.detect_image函数，在绘图部分对predicted_class进行判断，
比如判断if predicted_class == 'car': 即可判断当前目标是否为车，然后记录数量即可。利用draw.text即可写字。
'''
while True:
img = input('Input image filename:')
try:
image = cv2.imread(img)
except:
print('Open Error! Try again!')
continue
else:
r_image = yolo.detect_image(image)
cv2.imshow('result', r_image)
c = cv2.waitKey(0)
elif mode == "video":
capture = cv2.VideoCapture(video_path)
if video_save_path!="":
fourcc = cv2.VideoWriter_fourcc(*'XVID')
size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
out = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)
ref, frame = capture.read()
if not ref:
raise ValueError("未能正确读取摄像头（视频），请注意是否正确安装摄像头（是否正确填写视频路径）。")
fps = 0.0
while(True):
t1 = time.time()
# 读取某一帧
ref, frame = capture.read()
if not ref:
break
# 进行检测
frame = yolo.detect_image(frame)
fps = ( fps + (1./(time.time()-t1)) ) / 2
print("fps= %.2f"%(fps))
frame = cv2.putText(frame, "fps= %.2f"%(fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
cv2.imshow("video",frame)
c= cv2.waitKey(1) & 0xff
if video_save_path!="":
out.write(frame)
if c==27:
capture.release()
break
print("Video Detection Done!")
capture.release()
if video_save_path!="":
print("Save processed video to the path :" + video_save_path)
out.release()
cv2.destroyAllWindows()
elif mode == "fps":
img = cv2.imread(fps_image_path)
tact_time = yolo.get_FPS(img, test_interval)
print(str(tact_time) + ' seconds, ' + str(1/tact_time) + 'FPS, @batch_size 1')
elif mode == "dir_predict":
import os
from tqdm import tqdm
img_names = os.listdir(dir_origin_path)
for img_name in tqdm(img_names):
if img_name.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
image_path = os.path.join(dir_origin_path, img_name)
image = cv2.imread(image_path)
r_image = yolo.detect_image(image)
if not os.path.exists(dir_save_path):
os.makedirs(dir_save_path)
r_image.save(os.path.join(dir_save_path, img_name.replace(".jpg", ".png")), quality=95, subsampling=0)
else:
raise AssertionError("Please specify the correct mode: 'predict', 'video', 'fps', 'dir_predict'.")
結果の比較
最終モデルの推論速度は次のように比較されます。
変換前の fps は53、変換後の fps は です120。速度の向上は著しく、mAP は 98% から 97.75% に減少しましたが、影響はほとんどありません。
| モデル名 | トレーニングセット | テストセット | fps | 地図 |
|---|---|---|---|---|
| Yolov7-Tiny-OBB | UAV-RODトレイン | UAV-ROD 値 | 53 | 98.00% |
| Yolov7-Tiny-OBB-TRT | UAV-RODトレイン | UAV-ROD 値 | 120 | 97.75% |
