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

Add YOLO-Face + Minor fixes

Browse Source
This commit is contained in:
Marcos Luciano
2023-09-09 20:32:25 -03:00
parent 1a9df997a4
commit 1445990136
35 changed files with 54 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
utils/export_yoloV6.py
View File

@@ -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:], 2, keepdim=True)
scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True)
scores *= objectness
classes = classes.float()
return boxes, scores, classes
@@ -109,12 +109,12 @@ def parse_args():
parser.add_argument('--opset', type=int, default=13, help='ONNX opset version')
parser.add_argument('--simplify', action='store_true', help='ONNX simplify model')
parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size')
parser.add_argument('--batch', type=int, default=1, help='Implicit batch-size')
parser.add_argument('--batch', type=int, default=1, help='Static batch-size')
args = parser.parse_args()
if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file')
if args.dynamic and args.batch > 1:
raise SystemExit('Cannot set dynamic batch-size and implicit batch-size at same time')
raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
return args