thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix export files

Browse Source
This commit is contained in:
Marcos Luciano
2023-09-12 15:44:13 -03:00
parent b2388fb3cb
commit b738644f6e
7 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
utils/export_yoloV5.py
View File

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
x = x[0] x = x[0]
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes return boxes, scores, classes

2
utils/export_yoloV6.py
View File

@@ -22,7 +22,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes return boxes, scores, classes

2
utils/export_yoloV7.py
View File

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes return boxes, scores, classes

2
utils/export_yoloV7_u6.py
View File

@@ -17,7 +17,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
x = x.transpose(1, 2) x = x.transpose(1, 2)
boxes = x[:, :, :4] 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() classes = classes.float()
return boxes, scores, classes return boxes, scores, classes

2
utils/export_yoloV8.py
View File

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
x = x.transpose(1, 2) x = x.transpose(1, 2)
boxes = x[:, :, :4] 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() classes = classes.float()
return boxes, scores, classes return boxes, scores, classes

2
utils/export_yolor.py
View File

@@ -15,7 +15,7 @@ class DeepStreamOutput(nn.Module):
x = x[0] x = x[0]
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes return boxes, scores, classes

2
utils/export_yolox.py
View File

@@ -17,7 +17,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes return boxes, scores, classes