Pure onnx inference, exporting weights in PyTorch format, inference with PyTorch weights

demo.py

@@ -22,7 +22,7 @@
 """hyper parameters"""
 use_cuda = True
 
-def detect_cv2(cfgfile, weightfile, imgfile):
+def detect_cv2(cfgfile, weightfile, imgfile, namesfile=None):
     import cv2
     m = Darknet(cfgfile)
 
@@ -33,13 +33,14 @@ def detect_cv2(cfgfile, weightfile, imgfile):
     if use_cuda:
         m.cuda()
 
-    num_classes = m.num_classes
-    if num_classes == 20:
-        namesfile = 'data/voc.names'
-    elif num_classes == 80:
-        namesfile = 'data/coco.names'
-    else:
-        namesfile = 'data/x.names'
+    if namesfile is None:
+        num_classes = m.num_classes
+        if num_classes == 20:
+            namesfile = 'data/voc.names'
+        elif num_classes == 80:
+            namesfile = 'data/coco.names'
+        else:
+            namesfile = 'data/x.names'
     class_names = load_class_names(namesfile)
 
     img = cv2.imread(imgfile)
@@ -48,15 +49,15 @@ def detect_cv2(cfgfile, weightfile, imgfile):
 
     for i in range(2):
         start = time.time()
-        boxes = do_detect(m, sized, 0.4, 0.6, use_cuda)
+        boxes = do_detect(m, sized, 0.15, 0.2, use_cuda)
         finish = time.time()
         if i == 1:
             print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
 
-    plot_boxes_cv2(img, boxes[0], savename='predictions.jpg', class_names=class_names)
+    plot_boxes_cv2(img, boxes[0], savename='predictions_demopy.jpg', class_names=class_names)
 
 
-def detect_cv2_camera(cfgfile, weightfile):
+def detect_cv2_camera(cfgfile, weightfile, namesfile=None):
     import cv2
     m = Darknet(cfgfile)
 
@@ -73,13 +74,14 @@ def detect_cv2_camera(cfgfile, weightfile):
     cap.set(4, 720)
     print("Starting the YOLO loop...")
 
-    num_classes = m.num_classes
-    if num_classes == 20:
-        namesfile = 'data/voc.names'
-    elif num_classes == 80:
-        namesfile = 'data/coco.names'
-    else:
-        namesfile = 'data/x.names'
+    if namesfile is None:
+        num_classes = m.num_classes
+        if num_classes == 20:
+            namesfile = 'data/voc.names'
+        elif num_classes == 80:
+            namesfile = 'data/coco.names'
+        else:
+            namesfile = 'data/x.names'
     class_names = load_class_names(namesfile)
 
     while True:
@@ -100,7 +102,7 @@ def detect_cv2_camera(cfgfile, weightfile):
     cap.release()
 
 
-def detect_skimage(cfgfile, weightfile, imgfile):
+def detect_skimage(cfgfile, weightfile, imgfile, namesfile=None):
     from skimage import io
     from skimage.transform import resize
     m = Darknet(cfgfile)
@@ -112,26 +114,27 @@ def detect_skimage(cfgfile, weightfile, imgfile):
     if use_cuda:
         m.cuda()
 
-    num_classes = m.num_classes
-    if num_classes == 20:
-        namesfile = 'data/voc.names'
-    elif num_classes == 80:
-        namesfile = 'data/coco.names'
-    else:
-        namesfile = 'data/x.names'
+    if namesfile is None:
+        num_classes = m.num_classes
+        if num_classes == 20:
+            namesfile = 'data/voc.names'
+        elif num_classes == 80:
+            namesfile = 'data/coco.names'
+        else:
+            namesfile = 'data/x.names'
     class_names = load_class_names(namesfile)
 
     img = io.imread(imgfile)
     sized = resize(img, (m.width, m.height)) * 255
 
     for i in range(2):
         start = time.time()
-        boxes = do_detect(m, sized, 0.4, 0.4, use_cuda)
+        boxes = do_detect(m, sized, 0.15, 0.15, use_cuda)
         finish = time.time()
         if i == 1:
             print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
 
-    plot_boxes_cv2(img, boxes, savename='predictions.jpg', class_names=class_names)
+    plot_boxes_cv2(img, boxes, savename='predictions_demopy.jpg', class_names=class_names)
 
 
 def get_args():
@@ -144,6 +147,8 @@ def get_args():
     parser.add_argument('-imgfile', type=str,
                         default='./data/mscoco2017/train2017/190109_180343_00154162.jpg',
                         help='path of your image file.', dest='imgfile')
+    parser.add_argument('-class_names', type=str, default=None,
+                        help='path to file containing class names')
     args = parser.parse_args()
 
     return args
@@ -152,9 +157,9 @@ def get_args():
 if __name__ == '__main__':
     args = get_args()
     if args.imgfile:
-        detect_cv2(args.cfgfile, args.weightfile, args.imgfile)
-        # detect_imges(args.cfgfile, args.weightfile)
-        # detect_cv2(args.cfgfile, args.weightfile, args.imgfile)
-        # detect_skimage(args.cfgfile, args.weightfile, args.imgfile)
+        detect_cv2(args.cfgfile, args.weightfile, args.imgfile, args.class_names)
+        # detect_imges(args.cfgfile, args.weightfile, args.class_names)
+        # detect_cv2(args.cfgfile, args.weightfile, args.imgfile, args.class_names)
+        # detect_skimage(args.cfgfile, args.weightfile, args.imgfile, args.class_names)
     else:
-        detect_cv2_camera(args.cfgfile, args.weightfile)
+        detect_cv2_camera(args.cfgfile, args.weightfile, args.class_names)

onnx_inference_demo.py

@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+'''
+@Time          : 20/04/25 15:49
+@Author        : huguanghao
+@File          : onnx_inference_demo.py
+@Noice         : Infer with onnx weights and network structure.
+@Modificattion :
+    @Author    : Huxwell
+    @Time      : 21/04/15 9:58
+    @Detail    : Instead of darknet weights and cfg, use the onnx ones (i.e to test for fidelity)
+'''
+# import sys
+# import time
+# from PIL import Image, ImageDraw
+# from models.tiny_yolo import TinyYoloNet
+from tool.utils import *
+from tool.torch_utils import *
+from tool.darknet2pytorch import Darknet
+import onnxruntime
+import argparse
+
+"""hyper parameters"""
+use_cuda = True
+
+def detect_cv2(onnx, imgfile, namesfile):
+    import cv2
+    class_names = load_class_names(namesfile)
+
+    session = onnxruntime.InferenceSession(onnx, None)
+    input_name = session.get_inputs()[0].name
+    input_shape = session.get_inputs()[0].shape
+    print(input_name, input_shape)
+    input_h = input_shape[2]
+    input_w = input_shape[3]
+    network_width, network_height = input_shape[3], input_shape[2]
+
+    img = cv2.imread(imgfile)
+    sized = cv2.resize(img, (network_width, network_height))
+    sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
+    sized = sized.astype('float32') / 255.
+    sized = sized.transpose(2, 0, 1)
+    sized = sized.reshape(*input_shape)
+
+    for i in range(2):
+        start = time.time()
+        boxes = do_detect_onnx(session, sized, 0.15, 0.2, use_cuda, input_name)
+        finish = time.time()
+        if i == 1:
+            print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
+
+    plot_boxes_cv2(img, boxes[0], savename='predictions_onnx_weights.jpg', class_names=class_names)
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Test your image or video by trained model.')
+    parser.add_argument('-onnx', type=str,
+                        default='yolov4_1_3_416_416_static.onnx',
+                        help='path to onnx file', dest='onnx')
+    parser.add_argument('-imgfile', type=str,
+                        default='/data/mscoco2017/train2017/190109_180343_00154162.jpg',
+                        help='path of your image file.', dest='imgfile')
+    parser.add_argument('-class_names', type=str,
+                        help='path to file containing class names')
+    args = parser.parse_args()
+
+    return args
+
+
+if __name__ == '__main__':
+    args = get_args()
+    detect_cv2(args.onnx, args.imgfile, args.class_names)

pytorch_weights_demo.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+'''
+@Time          : 20/04/25 15:49
+@Author        : huguanghao
+@File          : pytorch_weights_demo.py
+@Noice         : Infer with PyTorch weights.
+@Modificattion :
+    @Author    : Huxwell
+    @Time      : 21/04/14 18:15
+    @Detail    : Instead of darknet weights, use the converted ones (i.e to test for fidelity)
+'''
+# import sys
+# import time
+# from PIL import Image, ImageDraw
+# from models.tiny_yolo import TinyYoloNet
+from tool.utils import *
+from tool.torch_utils import *
+from tool.darknet2pytorch import Darknet
+import argparse
+
+"""hyper parameters"""
+use_cuda = True
+
+
+def detect_cv2(cfgfile, torch_weights, imgfile, namesfile=None):
+    import cv2
+    m = Darknet(cfgfile)
+
+    m.print_network()
+    # m.load_weights(weightfile)
+    m.load_state_dict(torch.load(torch_weights))
+    print('Loading weights from %s... Done!' % (torch_weights))
+
+    if use_cuda:
+        m.cuda()
+
+    if namesfile is None:
+        num_classes = m.num_classes
+        if num_classes == 20:
+            namesfile = 'data/voc.names'
+        elif num_classes == 80:
+            namesfile = 'data/coco.names'
+        else:
+            namesfile = 'data/x.names'
+    class_names = load_class_names(namesfile)
+
+    img = cv2.imread(imgfile)
+    sized = cv2.resize(img, (m.width, m.height))
+    sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
+
+    for i in range(2):
+        start = time.time()
+        boxes = do_detect(m, sized, 0.15, 0.2, use_cuda)
+        finish = time.time()
+        if i == 1:
+            print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
+
+    plot_boxes_cv2(img, boxes[0], savename='predictions_torch_weights.jpg', class_names=class_names)
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Test your image or video by trained model.')
+    parser.add_argument('-cfgfile', type=str, default='./cfg/yolov4.cfg',
+                        help='path of cfg file', dest='cfgfile')
+    parser.add_argument('-torch_weights', type=str,
+                        default='weights.pth',
+                        help='path of trained model.', dest='torch_weights')
+    parser.add_argument('-imgfile', type=str,
+                        default='./data/mscoco2017/train2017/190109_180343_00154162.jpg',
+                        help='path of your image file.', dest='imgfile')
+    parser.add_argument('-class_names', type=str, default=None,
+                        help='path to file containing class names')
+    args = parser.parse_args()
+
+    return args
+
+
+if __name__ == '__main__':
+    args = get_args()
+    detect_cv2(args.cfgfile, args.torch_weights, args.imgfile, args.class_names)

save_pytorch_weights.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+'''
+@Time          : 20/04/25 15:49
+@Author        : huguanghao
+@File          : demo.py
+@Noice         : Save weights in PyTorch format.
+@Modificattion :
+    @Author    : Huxwell
+    @Time      : 21/04/14 18:15
+    @Detail    : 
+'''
+# import sys
+# import time
+# from PIL import Image, ImageDraw
+# from models.tiny_yolo import TinyYoloNet
+from tool.utils import *
+from tool.torch_utils import *
+from tool.darknet2pytorch import Darknet
+import argparse
+
+"""hyper parameters"""
+use_cuda = True
+
+def save_model_pytorch(cfgfile, weightfile, imgfile, output, namesfile=None):
+    import cv2
+    m = Darknet(cfgfile)
+
+    m.print_network()
+    m.load_weights(weightfile)
+    print('Loading weights from %s... Done!' % (weightfile))
+
+    if use_cuda:
+        m.cuda()
+
+    if namesfile is None:
+        num_classes = m.num_classes
+        if num_classes == 20:
+            namesfile = 'data/voc.names'
+        elif num_classes == 80:
+            namesfile = 'data/coco.names'
+        else:
+            namesfile = 'data/x.names'
+    class_names = load_class_names(namesfile)
+
+    img = cv2.imread(imgfile)
+    sized = cv2.resize(img, (m.width, m.height))
+    sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
+
+    for i in range(2):
+        start = time.time()
+        boxes = do_detect(m, sized, 0.15, 0.2, use_cuda)
+        finish = time.time()
+        if i == 1:
+            print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
+
+    plot_boxes_cv2(img, boxes[0], savename='predictions_while_saving_model.jpg', class_names=class_names)
+    torch.save(m.state_dict(), output)
+    print("Saved weights to ", output)
+
+def get_args():
+    parser = argparse.ArgumentParser('Test your image or video by trained model.')
+    parser.add_argument('-cfgfile', type=str, default='./cfg/yolov4.cfg',
+                        help='path of cfg file', dest='cfgfile')
+    parser.add_argument('-weightfile', type=str,
+                        default='./checkpoints/Yolov4_epoch1.pth',
+                        help='path of trained model.', dest='weightfile')
+    parser.add_argument('-imgfile', type=str,
+                        default='./data/mscoco2017/train2017/190109_180343_00154162.jpg',
+                        help='path of your image file.', dest='imgfile')
+    parser.add_argument('-output', type=str,
+                        default='converted_from_darknet.pth',
+                        help='path to save weights.', dest='output')
+    parser.add_argument('-class_names', type=str, default=None,
+                        help='path to file containing class names')
+    args = parser.parse_args()
+
+    return args
+
+
+if __name__ == '__main__':
+    args = get_args()
+    if args.imgfile:
+        save_model_pytorch(args.cfgfile, args.weightfile, args.imgfile, args.output, args.class_names)