Fix logger error in DeepStream 6.0 / 6.0.1 + Change output classes format + Fixes

utils/export_damoyolo.py

@@ -18,6 +18,7 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         boxes = x[1]
         scores, classes = torch.max(x[0], 2, keepdim=True)
+        classes = classes.float()
         return boxes, scores, classes
 
 

utils/export_ppyoloe.py

@@ -19,7 +19,7 @@ class DeepStreamOutput(nn.Layer):
         boxes = x['bbox']
         x['bbox_num'] = x['bbox_num'].transpose([0, 2, 1])
         scores = paddle.max(x['bbox_num'], 2, keepdim=True)
-        classes = paddle.argmax(x['bbox_num'], 2, keepdim=True)
+        classes = paddle.cast(paddle.argmax(x['bbox_num'], 2, keepdim=True), dtype='float32')
         return boxes, scores, classes
 
 

utils/export_yoloV5.py

@@ -20,6 +20,7 @@ class DeepStreamOutput(nn.Module):
         objectness = x[:, :, 4:5]
         scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
         scores *= objectness
+        classes = classes.float()
         return boxes, scores, classes
 
 

utils/export_yoloV6.py

@@ -24,6 +24,7 @@ class DeepStreamOutput(nn.Module):
         objectness = x[:, :, 4:5]
         scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
         scores *= objectness
+        classes = classes.float()
         return boxes, scores, classes
 
 

utils/export_yoloV7.py

@@ -20,6 +20,7 @@ class DeepStreamOutput(nn.Module):
         objectness = x[:, :, 4:5]
         scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
         scores *= objectness
+        classes = classes.float()
         return boxes, scores, classes
 
 

utils/export_yoloV7_u6.py

@@ -18,6 +18,7 @@ class DeepStreamOutput(nn.Module):
         x = x.transpose(1, 2)
         boxes = x[:, :, :4]
         scores, classes = torch.max(x[:, :, 4:], 2, keepdim=True)
+        classes = classes.float()
         return boxes, scores, classes
 
 

utils/export_yoloV8.py

@@ -19,6 +19,7 @@ class DeepStreamOutput(nn.Module):
         x = x.transpose(1, 2)
         boxes = x[:, :, :4]
         scores, classes = torch.max(x[:, :, 4:], 2, keepdim=True)
+        classes = classes.float()
         return boxes, scores, classes
 
 

utils/export_yolonas.py

@@ -15,6 +15,7 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         boxes = x[0]
         scores, classes = torch.max(x[1], 2, keepdim=True)
+        classes = classes.float()
         return boxes, scores, classes
 
 

utils/export_yolor.py

@@ -17,6 +17,7 @@ class DeepStreamOutput(nn.Module):
         objectness = x[:, :, 4:5]
         scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
         scores *= objectness
+        classes = classes.float()
         return boxes, scores, classes
 
 

utils/export_yolox.py

@@ -19,6 +19,7 @@ class DeepStreamOutput(nn.Module):
         objectness = x[:, :, 4:5]
         scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
         scores *= objectness
+        classes = classes.float()
         return boxes, scores, classes