From b738644f6e5c8918e9035647209af3814e406c17 Mon Sep 17 00:00:00 2001 From: Marcos Luciano Date: Tue, 12 Sep 2023 15:44:13 -0300 Subject: [PATCH] Fix export files --- utils/export_yoloV5.py | 2 +- utils/export_yoloV6.py | 2 +- utils/export_yoloV7.py | 2 +- utils/export_yoloV7_u6.py | 2 +- utils/export_yoloV8.py | 2 +- utils/export_yolor.py | 2 +- utils/export_yolox.py | 2 +- 7 files changed, 7 insertions(+), 7 deletions(-) diff --git a/utils/export_yoloV5.py b/utils/export_yoloV5.py index 7e7b0f5..08226b5 100644 --- a/utils/export_yoloV5.py +++ b/utils/export_yoloV5.py @@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module): x = x[0] boxes = x[:, :, :4] objectness = x[:, :, 4:5] - scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True) + scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True) scores *= objectness classes = classes.float() return boxes, scores, classes diff --git a/utils/export_yoloV6.py b/utils/export_yoloV6.py index abe09f9..3209ec8 100644 --- a/utils/export_yoloV6.py +++ b/utils/export_yoloV6.py @@ -22,7 +22,7 @@ class DeepStreamOutput(nn.Module): def forward(self, x): boxes = x[:, :, :4] objectness = x[:, :, 4:5] - scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True) + scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True) scores *= objectness classes = classes.float() return boxes, scores, classes diff --git a/utils/export_yoloV7.py b/utils/export_yoloV7.py index b87a030..ea5557f 100644 --- a/utils/export_yoloV7.py +++ b/utils/export_yoloV7.py @@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module): def forward(self, x): boxes = x[:, :, :4] objectness = x[:, :, 4:5] - scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True) + scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True) scores *= objectness classes = classes.float() return boxes, scores, classes diff --git a/utils/export_yoloV7_u6.py b/utils/export_yoloV7_u6.py index 0da10c2..5c1f89e 100644 --- a/utils/export_yoloV7_u6.py +++ b/utils/export_yoloV7_u6.py @@ -17,7 +17,7 @@ class DeepStreamOutput(nn.Module): def forward(self, x): x = x.transpose(1, 2) boxes = x[:, :, :4] - scores, classes = torch.max(x[:, :, 4:5], 2, keepdim=True) + scores, classes = torch.max(x[:, :, 4:], 2, keepdim=True) classes = classes.float() return boxes, scores, classes diff --git a/utils/export_yoloV8.py b/utils/export_yoloV8.py index df7f0f7..232b722 100644 --- a/utils/export_yoloV8.py +++ b/utils/export_yoloV8.py @@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module): def forward(self, x): x = x.transpose(1, 2) boxes = x[:, :, :4] - scores, classes = torch.max(x[:, :, 4:5], 2, keepdim=True) + scores, classes = torch.max(x[:, :, 4:], 2, keepdim=True) classes = classes.float() return boxes, scores, classes diff --git a/utils/export_yolor.py b/utils/export_yolor.py index e03a3e4..3175170 100644 --- a/utils/export_yolor.py +++ b/utils/export_yolor.py @@ -15,7 +15,7 @@ class DeepStreamOutput(nn.Module): x = x[0] boxes = x[:, :, :4] objectness = x[:, :, 4:5] - scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True) + scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True) scores *= objectness classes = classes.float() return boxes, scores, classes diff --git a/utils/export_yolox.py b/utils/export_yolox.py index 2b52113..a41c80e 100644 --- a/utils/export_yolox.py +++ b/utils/export_yolox.py @@ -17,7 +17,7 @@ class DeepStreamOutput(nn.Module): def forward(self, x): boxes = x[:, :, :4] objectness = x[:, :, 4:5] - scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True) + scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True) scores *= objectness classes = classes.float() return boxes, scores, classes