diff --git a/README.md b/README.md index 5203d19..f320593 100644 --- a/README.md +++ b/README.md @@ -2,27 +2,31 @@ ## 程序余API介绍 -'主要的程序接口(API)全部位于detect.py和param.py两个文件中。' +>主要的程序接口(API)全部位于detect.py和param.py两个文件中。 ### 检测模块:detect.py -主要是包含一个类和*一段示例代码*: -0. `class YOLOv7(object)`:该类即为本程序提供的主要接口 +主要是包含*一个类*和*一段示例代码*: -1. `__init__(self,option:param.Parameters) -> None` +* `class YOLOv7(object)`:该类即为本程序提供的主要接口 -主要是创建一个已加载权重的模型(model)以供后续推断过程的使用,之所以将函数与推断模型分开的原因是权重加载较慢,这样做可以使得不必在每一次推断的过程中重复加载权重,这对提高实验程序的实时性有很大帮助。(尽管使用ONNX模型更加有利于提升速度,但很多时候完成课程或毕业设计这样已经足够了)。 -参数:`opt:param.Parameters`是位于param.py中的一个类,仅仅是为了避免使用全局变量同时满足各种情形下的传参操作才这么设计。 -返回值:无,但内部变量`self.model`是已加载参数的模型(其父类为torch.nn.Moudules)。 +* `__init__(self,option:param.Parameters) -> None` -1. `detect(self,imgs:torch.Tensor) -> list` +>主要是创建一个已加载权重的模型(model)以供后续推断过程的使用,之所以将函数与推断模型分开的原因是权重加载较慢,这样做可以使得不必在每一次推断的过程中重复加载权重,这对提高实验程序的实时性有很大帮助。(尽管使用ONNX模型更加有利于提升速度,但很多时候完成课程或毕业设计这样已经足够了)。 +> +>参数:`opt:param.Parameters`是位于param.py中的一个类,仅仅是为了避免使用全局变量同时满足各种情形下的传参操作才这么设计。 +> +>返回值:无,但内部变量`self.model`是已加载参数的模型(其父类为torch.nn.Moudules)。 -参数:`imgs`(类型为torch.Tensor),大小应为(B,3,H,W),B为照片的张数,当为单张照片时为B==1即可。3时RGB三通道,W、H分别为图片的高和宽,**要求为64的倍数**。 -返回值:一个包含所有图片的bouding_box的列表,列表中每个元素对应每一张图片,其类型是二维的torch.Tensor;最后一维长度为6,分别对应(左上对角点横坐标`x1`,左上对角点纵坐标`y1`,右下角坐标横坐标`x2`,右下角纵坐标`y2`,可信度`conf`,类别`cls`),第二维是对应单张图片中的bouding_box。 +* `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` +>主要是一个简单的类`Parameters` | 变量名称 | 说明 | | ------------ | ----------------------------------------------- | @@ -33,3 +37,54 @@ | classes | 选中的类,为`None`时代表全选 | | agnostic_nms | NMS算法选项,通常为`None`即可 | | augment | 传入YOLOv7模型的参数,通常为`None`即可 | + +## 示例 + +### 代码 + +```python +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) +``` + +```python +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 +``` + +```python +# 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 +``` + +### 输入图片 + +![avatar](test.png) + +### 输出图片 + +![avatar](result.png) diff --git a/param.py b/param.py index d97172e..e847e8b 100644 --- a/param.py +++ b/param.py @@ -6,7 +6,6 @@ class Parameters(object): self.conf_thres = 0.25 self.iou_thres = 0.45 self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - self.nosave = False self.classes = [0] self.agnostic_nms = None self.augment = None