简化后的YOLOv7推断模块

程序余API介绍

主要的程序接口（API）全部位于detect.py和param.py两个文件中。

检测模块：detect.py

主要是包含一个类和一段示例代码：

• class YOLOv7(object):该类即为本程序提供的主要接口

• __init__(self,option:param.Parameters) -> None

主要是创建一个已加载权重的模型（model）以供后续推断过程的使用，之所以将函数与推断模型分开的原因是权重加载较慢，这样做可以使得不必在每一次推断的过程中重复加载权重，这对提高实验程序的实时性有很大帮助。（尽管使用ONNX模型更加有利于提升速度，但很多时候完成课程或毕业设计这样已经足够了）。

参数：opt:param.Parameters是位于param.py中的一个类，仅仅是为了避免使用全局变量同时满足各种情形下的传参操作才这么设计。

返回值：无，但内部变量self.model是已加载参数的模型（其父类为torch.nn.Moudules)。

• detect(self,imgs:torch.Tensor) -> list

参数：imgs(类型为torch.Tensor)，大小应为(B,3,H,W),B为照片的张数，当为单张照片时为B==1即可。3时RGB三通道，W、H分别为图片的高和宽，要求为64的倍数。

返回值:一个包含所有图片的bouding_box的列表，列表中每个元素对应每一张图片，其类型是二维的torch.Tensor；第二维是对应单张图片中的bouding_box;最后一维长度为6，分别对应(左上对角点横坐标x1，左上对角点纵坐标y1,右下角坐标横坐标x2，右下角纵坐标y2,可信度conf,类别cls)。注意元素类型都为float型,后期处理可能涉及到类型转换。

传参模块：param.py

主要是一个简单的类Parameters

变量名称	说明
weights	权重文件.pt文件的路径
conf_thres	置信度阈值，小于该置信度将被过滤
iou_thres	非极大值抑制(NMS)阈值，重叠比值超过该值将被合并
device	所部署到的设备.to(device)
classes	选中的类，为None时代表全选
agnostic_nms	NMS算法选项，通常为None即可
augment	传入YOLOv7模型的参数，通常为None即可

示例

代码

def get_img_test(path_src:str):                                         
    raw_img = cv2.imread(path_src)                                          # get image as numpy arrary
    raw_transform = transforms.Compose([transforms.ToPILImage(),            # process it as needed and transform it into torch.Tensor
                                        transforms.Resize((360,640)), 
                                        transforms.Pad((0,(640-360)//2)),       
                                        transforms.ToTensor()])
    return raw_transform(raw_img).unsqueeze(dim=0)

def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
    pred_img = pred_img.astype(np.uint)                       # transform its type into uint,pay attention to the confidence would be zero as percentage never large than 1.0
    print(type(img))
    for info in pred_img:
        x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
        cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)# recover the pixel bias because of padding
    print(type(img))
    cv2.imshow("test",img)
    cv2.waitKey(0)
    cv2.imwrite("result.png",img)                             # save the image

# example        
if __name__ == '__main__':
    my_option = param.Parameters()
    model = YOLOv7(my_option)                       # load model in head
    with torch.no_grad():
        img = get_img_test("test.png")              # get image as torch.Tensor with size of [1,1,640,640] 
        print(f"img size is {img.size()}")        
        result = model.detect(img)                  # get the sequence of result
        print(f"result is :{result}")
        result = result[0].detach().cpu().numpy()   # the single img is index 0
        img_src = cv2.imread("test.png")            # re-read for imshow
        img_src = cv2.resize(img_src,(640,360))
        opencv_box_plot(img_src,result)             # show the result in picture

输入图片

输出图片