ken4647
d6e0318f91
|
||
|---|---|---|
| models | ||
| utils | ||
| .gitignore | ||
| README.md | ||
| detect.py | ||
| param.py | ||
| result.png | ||
| test.png | ||
README.md
简化后的YOLOv7推断模块
程序余API介绍
主要的程序接口（API）全部位于detect.py和param.py两个文件中。在输入接口
detect(self,imgs:torch.Tensor) -> list的输入变量类型是torch的张量类型，可以通过opencv的cv2.imread()或者capture.read()读取得到numpy数组(array)类型，然后通过torch.Tensor()或者torchvision.transforms进行类型转换。图片尺寸必须是方形，即W=H,如果不是建议先通过resize和pad操作进行变换；如果是单张图片，输入模型前必须扩充维数，可以使用Tensor的unsqueeze(dim=0)方法进行（即扩充第一维度),实现的伪代码如下:
raw_img = cv2.imread(path_src) # 读取单张图片，path_src处填写图片路径
raw_transform = transforms.Compose([transforms.ToPILImage(),
transforms.Resize((360,640)),
transforms.Pad((0,(640-360)//2)),
transforms.ToTensor()])
return raw_transform(raw_img).unsqueeze(dim=0)
如果是读取视频只需按opencv读取视频的方法进行即可:
capture = cv2.VideoCapture(path_src) # 读取摄像头为cv2.VideoCapture(index),index是相机索引，通常为0即可；读取视频文件时path_src填写路径即可
raw_transform = transforms.Compose([transforms.ToPILImage(),
transforms.Resize((360,640)),
transforms.Pad((0,(640-360)//2)),
transforms.ToTensor()]) # 预先组合好的变换函数
ret,frame = capture.read()
while ret is not None:
frame_tensor = raw_transform(frame).unsqueeze(dim=0) # 由于输入的是单张图片，需要在dim=0进行维数扩充，由(C,H,W)到(1,C,H,W)，总尺寸大小其实不会发生改变
# your coder for detection
# ...
ret,frame = capture.read()
capture.release()
检测模块：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
输入图片
输出图片
