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

4
utils/export_damoyolo.py
View File

@@ -107,14 +107,14 @@ def parse_args():
parser.add_argument('--opset', type=int, default=11, 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 not os.path.isfile(args.config):
raise SystemExit('Invalid config 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

4
utils/export_ppyoloe.py
View File

@@ -89,12 +89,12 @@ def parse_args():
parser.add_argument('--opset', type=int, default=11, 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('\nInvalid weights file')
if args.dynamic and args.batch > 1:
raise SystemExit('\nCannot set dynamic batch-size and implicit batch-size at same time')
raise SystemExit('\nCannot set dynamic batch-size and static batch-size at same time')
elif args.dynamic:
args.batch = None
return args

6
utils/export_yoloV5.py
View File

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
x = x[0]
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
@@ -106,12 +106,12 @@ def parse_args():
parser.add_argument('--opset', type=int, default=17, 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

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

6
utils/export_yoloV7.py
View File

@@ -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:], 2, keepdim=True)
scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True)
scores *= objectness
classes = classes.float()
return boxes, scores, classes
@@ -110,12 +110,12 @@ def parse_args():
parser.add_argument('--opset', type=int, default=12, 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

6
utils/export_yoloV7_u6.py
View File

@@ -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:], 2, keepdim=True)
scores, classes = torch.max(x[:, :, 4:5], 2, keepdim=True)
classes = classes.float()
return boxes, scores, classes
@@ -100,12 +100,12 @@ def parse_args():
parser.add_argument('--opset', type=int, default=12, 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

6
utils/export_yoloV8.py
View File

@@ -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:], 2, keepdim=True)
scores, classes = torch.max(x[:, :, 4:5], 2, keepdim=True)
classes = classes.float()
return boxes, scores, classes
@@ -108,12 +108,12 @@ def parse_args():
parser.add_argument('--opset', type=int, default=16, 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

4
utils/export_yolonas.py
View File

@@ -89,14 +89,14 @@ def parse_args():
parser.add_argument('--opset', type=int, default=14, 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 args.model == '':
raise SystemExit('Invalid model name')
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

6
utils/export_yolor.py
View File

@@ -15,7 +15,7 @@ class DeepStreamOutput(nn.Module):
x = x[0]
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
@@ -123,12 +123,12 @@ def parse_args():
parser.add_argument('--opset', type=int, default=12, 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

6
utils/export_yolox.py
View File

@@ -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:], 2, keepdim=True)
scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True)
scores *= objectness
classes = classes.float()
return boxes, scores, classes
@@ -96,14 +96,14 @@ def parse_args():
parser.add_argument('--opset', type=int, default=11, 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 not os.path.isfile(args.exp):
raise SystemExit('Invalid exp 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