import cv2
import torch
import numpy as np
from yolov7.utils import non_max_suppression, letterbox, scale_coords, plot_one_box
class yolov7:
def __init__(self, model_path, iou_thres, conf_thres, device, names):
device = self.select_device(device)
print(device)
if model_path.endswith('pt'):
model = torch.jit.load(model_path).to(device)
elif model_path.endswith('onnx'):
try:
import onnxruntime as ort
except:
raise 'please install onnxruntime.'
providers = ['CUDAExecutionProvider'] if device.type != 'cpu' else ['CPUExecutionProvider']
model = ort.InferenceSession(model_path, providers=providers)
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
self.__dict__.update(locals())
导入了一些自定义的工具函数和类:non_max_suppression、letterbox、scale_coords和plot_one_box。这些函数和类用于执行非最大抑制、图像letterbox处理(调整图像尺寸后填充)、坐标缩放和绘制边界框等功能。
接着,定义了一个名为yolov7
的类。该类的构造函数__init__
接收以下参数:model_path
(模型文件路径,可以是.pt
或.onnx
格式)、iou_thres
(IoU阈值,用于非最大抑制的阈值)、conf_thres
(置信度阈值,用于目标检测结果的筛选)、device
(设备类型,可以是CPU或GPU)、names
(类别名称列表)。
__init__
函数首先通过调用select_device
方法选择设备,并打印出所选择的设备。然后,根据model_path
的后缀来加载模型,如果是.pt
格式,则使用torch.jit.load
加载模型并将其发送到选择的设备;如果是.onnx
格式,则使用onnxruntime.InferenceSession
加载模型,同时根据设备类型选择使用的执行提供者。
def __call__(self, data):
if type(data) is str:
image = cv2.imdecode(np.fromfile(data, np.uint8), cv2.IMREAD_COLOR)
else:
image = data
im = self.processing(image)
if self.model_path.endswith('pt'):
result = self.model(im)[0]
elif self.model_path.endswith('onnx'):
result = self.model.run([i.name for i in self.model.get_outputs()], {'images':im})[0]
image = self.post_processing(result, im, image)
return image
def processing(self, img):
image = letterbox(img, auto=False)[0]
image = image.transpose((2, 0, 1))[::-1]
image = np.expand_dims(image, 0)
image = np.ascontiguousarray(image)
image = np.array(image, dtype=np.float32)
image /= 255
if self.model_path.endswith('pt'):
im = torch.from_numpy(image).float().to(self.device)
elif self.model_path.endswith('onnx'):
im = image
return im
def post_processing(self, result, im=None, img=None):
if self.model_path.endswith('pt'):
result = non_max_suppression(result, conf_thres=self.conf_thres, iou_thres=self.iou_thres)[0]
result[:, :4] = scale_coords(im.shape[2:], result[:, :4], img.shape)
for *xyxy, conf, cls in result:
label = f'{self.names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, img, label=label, color=self.colors[int(cls)], line_thickness=1)
elif self.model_path.endswith('onnx'):
result = result[:, 1:]
ratio, dwdh = letterbox(img, auto=False)[1:]
result[:, :4] -= np.array(dwdh * 2)
result[:, :4] /= ratio
for *xyxy, cls, conf in result:
label = f'{self.names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, img, label=label, color=self.colors[int(cls)], line_thickness=1)
return img
def select_device(self, device):
if device == -1:
return torch.device('cpu')
else:
os.environ['CUDA_VISIBLE_DEVICES'] = str(device)
assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
device = torch.device('cuda:0')
return device
if __name__ == '__main__':
# read cfg
with open('model.yaml') as f:
cfg = yaml.load(f, Loader=yaml.SafeLoader)
# print cfg
print(cfg)
# init
yolo = yolov7(**cfg)
image_path = '1.jpg'
# inference
image = yolo(image_path)
cv2.imshow('pic', image)
cv2.waitKey(0)
cv2.destroyAllWindows()
定义了一个colors
列表,其中包含了每个类别的随机颜色,用于绘制边界框时的类别标签颜色。
__call__
方法用于实际执行推断。它接收一个输入数据data
,可以是图像文件的路径字符串或者是OpenCV图像数组。
在__call__
方法内部,首先根据输入数据的类型,使用OpenCV的imdecode
函数将图像文件解码为OpenCV图像数组。然后,调用processing
方法对图像进行预处理,将其转换为模型所需的格式。
接着,根据model_path
的后缀,使用加载的模型对输入数据进行推断。如果是.pt
格式,则使用模型进行一次前向传播,得到输出结果;如果是.onnx
格式,则使用模型的run
方法得到输出结果。
最后,调用post_processing
方法对输出结果进行后处理,将检测到的边界框绘制在原始图像上,并返回绘制完成的图像。
processing
方法用于对图像进行预处理。它首先使用letterbox
函数将图像调整为模型所需的输入尺寸,并做一些其他必要的转换操作,如转置、扩展维度和归一化。最后,根据模型的类型和格式,将图像转换为相应的torch.Tensor
或numpy.ndarray
格式,并返回预处理后的图像。
post_processing
方法用于对模型的输出结果进行后处理。首先,根据模型的类型和格式,对输出结果进行一些调整和计算,如非最大抑制的阈值过滤、坐标缩放等。然后,利用plot_one_box
函数将结果绘制在原始图像上,并返回绘制完成的图像。
select_device
方法用于选择设备。如果device
为-1,则选择CPU设备;否则,通过设置环境变量CUDA_VISIBLE_DEVICES
来选择GPU设备,并确保选择的设备可用。
最后,代码利用上述定义的yolov7
类实例化了一个对象,并进行了简单的测试。它首先加载了一个配置文件(model.yaml
),然后根据配置初始化了yolov7
对象。接着,指定一个图像路径,通过调用yolo
对象对图像进行推断,并将结果显示在窗口中。