Compare commits

...

7 Commits

4 changed files with 243 additions and 13 deletions

9
.gitignore vendored
View File

@ -1,3 +1,10 @@
# tmp files
*__pycache__*
*.pt
*.mp4
*.jpg
*.mp4
*.t7
*.pth
*label*
result.txt
rank10_detection*.txt

110
README.md
View File

@ -2,27 +2,58 @@
## 程序余API介绍
'主要的程序接口API全部位于detect.py和param.py两个文件中。'
>主要的程序接口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)`方法进行(即扩充第一维度),实现的伪代码如下:
```python
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读取视频的方法进行即可:
```python
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
主要是包含一个类和*一段示例代码*
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 +64,60 @@
| 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)
### 参考链接
[YOLOv7官方实现-github](https://github.com/WongKinYiu/yolov7)
[YOLOv7论文-arxiv](https://arxiv.org/abs/2207.02696)

136
detector.py Normal file
View File

@ -0,0 +1,136 @@
import os
import cv2
import torch
import numpy as np
from torchvision import transforms
from models.experimental import attempt_load
from utils.general import check_img_size,non_max_suppression
import param
def IoU(box1, box2):
"""
:param box1: list in format [xmin1, ymin1, xmax1, ymax1]
:param box2: list in format [xmin2, ymin2, xamx2, ymax2]
:return: returns IoU ratio (intersection over union) of two boxes
"""
xmin1, ymin1, xmax1, ymax1 = box1
xmin2, ymin2, xmax2, ymax2 = box2[0],box2[1],box2[2],box2[3]
ymin2=ymin2-420
ymax2=ymax2-420
x_overlap = max(0, min(xmax1, xmax2) - max(xmin1, xmin2))
y_overlap = max(0, min(ymax1, ymax2) - max(ymin1, ymin2))
intersection = x_overlap * y_overlap
union = (xmax1 - xmin1) * (ymax1 - ymin1) + (xmax2 - xmin2) * (ymax2 - ymin2) - intersection
IOU=float(intersection) / union
return IOU
def cxywh2xyxy(cxywh):
xyxy = [0, 0, 0, 0]
xyxy[0] = int(cxywh[0] - cxywh[2] / 2)
xyxy[1] = int(cxywh[1] - cxywh[3] / 2)
xyxy[2] = int(cxywh[0] + cxywh[2] / 2)
xyxy[3] = int(cxywh[1] + cxywh[3] / 2)
xyxy = [(num if num > 0 else 0) for num in xyxy]
return xyxy
class YOLOv7(object):
def __init__(self,option:param.Parameters) -> None:
# load FP32 model
self.opt = option
self.model = attempt_load(weights=self.opt.weights, device=self.opt.device)
def detect(self,imgs:torch.Tensor) -> list:
# img.size should be (Batch,3,H,W),every value should be range of [0,1]
# Both H and W % 64 == 0
# Initialize
model = self.model
opt = self.opt
# Inference
with torch.no_grad(): # Calculating gradients would cause a GPU memory leak
pred = model(imgs, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
return pred
def get_img_test(path_src:str):
raw_img = cv2.imread(path_src)
raw_transform = transforms.Compose([transforms.ToPILImage(),
transforms.Resize((1080,1920)),
transforms.Pad((0, (1920-1080)//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)
print(type(img))
for info in pred_img:
x1, y1, x2, y2 = info[0], info[1], info[2], info[3]
if abs(x2-x1)> 32:
cut = img[(y1-420):(y2 - 420), x1:x2]
# cv2.rectangle(img, (x1, y1-420), (x2, y2-420), (255, 0, 0), 2)
cut = cv2.resize(cut, (128, 256))
results_len = len(results)
number_bbox = 0
for i in range(results_len):
id_bbox = results[i]
id = str(id_bbox).split(',')[0][1:]
xxyy = cxywh2xyxy(id_bbox[1:5])
if IoU(xxyy, info) > 0.7:
cv2.imwrite(f"./detector/cam{camera}/" + "d" + "_" + name.split(".")[0] + "_" + str(id) + ".jpg", cut)
break
else:
number_bbox += 1
if number_bbox == results_len:
cv2.imwrite(f"./detector/cam{camera}/" + "d" + "_" + name.split(".")[0] + "_" + "0" + ".jpg", cut)
# print(type(img))
# cv2.imshow("test", img)
cv2.waitKey(0)
# cv2.imwrite("result.png", img)
# example
if __name__ == '__main__':
my_option = param.Parameters()
model = YOLOv7(my_option) # load model in head
count_frame = -1
ret = None
frame = None
with torch.no_grad():
# for i in range(1,7):
camera = 4
with open(f"../UESTC_ReID_Dataset_V3/label/cam{camera}.txt", 'r') as flabel:
camera_path = f"../UESTC_ReID_Dataset_V3/images/img/cam{camera}/"
txt_data = flabel.readlines()
filename = os.listdir(camera_path)
results=[]
# filename = "011.jpg"
count = 0
for name in filename:
count += 1
if count % 12 == 0:
for line in txt_data:
data2 = line.split(',')
data2 = [int(data2) for data2 in data2[0:6]]
# print(data2[0])
if int(data2[0]) == int(name.split(".")[0]):
results.append(data2[1:6])
print(name)
img = get_img_test(f"../UESTC_ReID_Dataset_V3/images/img/cam{camera}/"+name) # 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(f"../UESTC_ReID_Dataset_V3/images/img/cam{camera}/"+name) # re-read for imshow
img_src = cv2.resize(img_src, (1920, 1080))
opencv_box_plot(img_src, result) # show the result in picture
results.clear()

View File

@ -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