apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

minor fix

Browse Source
This commit is contained in:
AlexeyAB84 2022-07-22 17:25:11 +03:00
parent e5928b095a
commit a5b853c3fa
2 changed files with 26 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
README.md
View File

@ -153,20 +153,26 @@ python detect.py --weights yolov7.pt --conf 0.25 --img-size 640 --source inferen
## Export ## Export
Tested with: Python 3.7.13 and Pytorch 1.12.0+cu113
Pytorch to ONNX, use `--include-nms` flag for the end-to-end ONNX model with `EfficientNMS`. Pytorch -> ONNX -> TensorRT -> Detection on TensorRT in Python <a href="https://colab.research.google.com/gist/AlexeyAB/fcb47ae544cf284eb24d8ad8e880d45c/yolov7trtlinaom.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
**Pytorch to ONNX**, use `--include-nms` flag for the end-to-end ONNX model with `EfficientNMS`
```shell ```shell
wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt
python export.py --weights yolov7-tiny.pt --grid --include-nms python export.py --weights yolov7-tiny.pt --grid --include-nms
``` ```
ONNX to TensorRT **ONNX to TensorRT**
```shell ```shell
git clone https://github.com/Linaom1214/tensorrt-python.git git clone https://github.com/Linaom1214/tensorrt-python.git
cd tensorrt-python cd tensorrt-python
python export.py -o yolov7-tiny.onnx -e yolov7-tiny-nms.trt -p fp16 python export.py -o yolov7-tiny.onnx -e yolov7-tiny-nms.trt -p fp16
``` ```
Tested with: Python 3.7.13, Pytorch 1.12.0+cu113
## Citation ## Citation
``` ```

18
models/yolo.py
View File

@ -84,6 +84,7 @@ class Detect(nn.Module):
class IDetect(nn.Module): class IDetect(nn.Module):
stride = None # strides computed during build stride = None # strides computed during build
export = False # onnx export export = False # onnx export
include_nms = False
def __init__(self, nc=80, anchors=(), ch=()): # detection layer def __init__(self, nc=80, anchors=(), ch=()): # detection layer
super(IDetect, self).__init__() super(IDetect, self).__init__()
@ -139,7 +140,10 @@ class IDetect(nn.Module):
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
z.append(y.view(bs, -1, self.no)) z.append(y.view(bs, -1, self.no))
return x if self.training else (torch.cat(z, 1), x) if self.include_nms:
z = self.convert(z)
return x if self.training else (z, ) if self.include_nms else (torch.cat(z, 1), x)
def fuse(self): def fuse(self):
print("IDetect.fuse") print("IDetect.fuse")
@ -160,6 +164,18 @@ class IDetect(nn.Module):
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
def convert(self, z):
z = torch.cat(z, 1)
box = z[:, :, :4]
conf = z[:, :, 4:5]
score = z[:, :, 5:]
score *= conf
convert_matrix = torch.tensor([[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]],
dtype=torch.float32,
device=z.device)
box @= convert_matrix
return (box, score)
class IKeypoint(nn.Module): class IKeypoint(nn.Module):
stride = None # strides computed during build stride = None # strides computed during build