End2end (#61)

* export end2end onnx model

* fixbug

* add web demo (#58)

* Update README.md

* main code

update yolov7-tiny deploy cfg

* main code

update yolov7-tiny training cfg

* main code

@liguagua752109150 https://github.com/WongKinYiu/yolov7/issues/33#issuecomment-1178669212

* main code

@albertfaromatics https://github.com/WongKinYiu/yolov7/issues/35#issuecomment-1178800685

* main code

update link

* main code

add custom hyp

* main code

update default activation function

* main code

update path

* main figure

add more tasks

* main code

update readme

* main code

update reparameterization

* Update README.md

* main code

update readme

* main code

update aux training

* main code

update aux training

* main code

update aux training

* main figure

update yolov7 prediction

* main code

update readme

* main code

rename

* main code

rename

* main code

rename

* main code

rename

* main code

update readme

* main code

update visualization

* main code

fix gain for train_aux

* main code

update loss

* main code

update instance segmentation demo

* main code

update keypoint detection demo

* main code

update pose demo

* main code

update pose

* main code

update pose

* main code

update pose

* main code

update pose

* main code

update trace

* Update README.md

* main code

fix ciou

* main code

fix nan of aux training https://github.com/WongKinYiu/yolov7/issues/250#issue-1312356380 @hudingding

* support onnx to tensorrt convert (#114)

* fuse IDetect (#148)

* Fixes #199 (#203)

* minor fix

* resolve conflict

* resolve conflict

* resolve conflict

* resolve conflict

* resolve conflict

* resolve

* resolve

* resolve

* resolve

Co-authored-by: AK391 <81195143+AK391@users.noreply.github.com>
Co-authored-by: Alexey <AlexeyAB@users.noreply.github.com>
Co-authored-by: Kin-Yiu, Wong <102582011@cc.ncu.edu.tw>
Co-authored-by: linghu8812 <36389436+linghu8812@users.noreply.github.com>
Co-authored-by: Alexander <84590713+SashaAlderson@users.noreply.github.com>
Co-authored-by: Ben Raymond <ben@theraymonds.org>
Co-authored-by: AlexeyAB84 <alexeyab84@gmail.com>

models/common.py

@@ -236,7 +236,7 @@ class Res(nn.Module):
 class ResX(Res):
     # ResNet bottleneck
     def __init__(self, c1, c2, shortcut=True, g=32, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
-        super().__init__(c1, c2, shortcu, g, e)
+        super().__init__(c1, c2, shortcut, g, e)
         c_ = int(c2 * e)  # hidden channels
 
 

models/yolo.py

@@ -5,7 +5,7 @@ from copy import deepcopy
 
 sys.path.append('./')  # to run '$ python *.py' files in subdirectories
 logger = logging.getLogger(__name__)
-
+import torch
 from models.common import *
 from models.experimental import *
 from utils.autoanchor import check_anchor_order
@@ -23,7 +23,7 @@ except ImportError:
 class Detect(nn.Module):
     stride = None  # strides computed during build
     export = False  # onnx export
-
+    include_nms = False 
     def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
         super(Detect, self).__init__()
         self.nc = nc  # number of classes
@@ -48,7 +48,6 @@ class Detect(nn.Module):
             if not self.training:  # inference
                 if self.grid[i].shape[2:4] != x[i].shape[2:4]:
                     self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
-
                 y = x[i].sigmoid()
                 if not torch.onnx.is_in_onnx_export():
                     y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
@@ -59,13 +58,28 @@ class Detect(nn.Module):
                     y = torch.cat((xy, wh, y[..., 4:]), -1)
                 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)
 
     @staticmethod
     def _make_grid(nx=20, ny=20):
         yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
         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 IDetect(nn.Module):
     stride = None  # strides computed during build