diff --git a/demo.py b/demo.py index a8988a59..19c62c8d 100644 --- a/demo.py +++ b/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,13 +114,14 @@ 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) @@ -126,12 +129,12 @@ def detect_skimage(cfgfile, weightfile, imgfile): 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) diff --git a/onnx_inference_demo.py b/onnx_inference_demo.py new file mode 100644 index 00000000..810695cb --- /dev/null +++ b/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) \ No newline at end of file diff --git a/pytorch_weights_demo.py b/pytorch_weights_demo.py new file mode 100644 index 00000000..6f2e7445 --- /dev/null +++ b/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) diff --git a/save_pytorch_weights.py b/save_pytorch_weights.py new file mode 100644 index 00000000..b6e6407b --- /dev/null +++ b/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) \ No newline at end of file diff --git a/tool/torch_utils.py b/tool/torch_utils.py index 68cbfb27..dfbf5866 100644 --- a/tool/torch_utils.py +++ b/tool/torch_utils.py @@ -102,3 +102,33 @@ def do_detect(model, img, conf_thresh, nms_thresh, use_cuda=1): return utils.post_processing(img, conf_thresh, nms_thresh, output) +def do_detect_onnx(session, img, conf_thresh, nms_thresh, use_cuda, input_name): + # model.eval() + t0 = time.time() + + # if type(img) == np.ndarray and len(img.shape) == 3: # cv2 image + # # img = torch.from_numpy(img.transpose(2, 0, 1)).float().div(255.0).unsqueeze(0) + # # elif type(img) == np.ndarray and len(img.shape) == 4: + # # img = torch.from_numpy(img.transpose(0, 3, 1, 2)).float().div(255.0) + # else: + # print("unknow image type") + # exit(-1) + + # if use_cuda: + # img = img.cuda() + # img = torch.autograd.Variable(img) + + t1 = time.time() + + # output = model(img) + raw_result = session.run([], {input_name: img}) + output = raw_result + + t2 = time.time() + + print('-----------------------------------') + print(' Preprocess : %f' % (t1 - t0)) + print(' Model Inference : %f' % (t2 - t1)) + print('-----------------------------------') + + return utils.post_processing(img, conf_thresh, nms_thresh, output) \ No newline at end of file diff --git a/tool/utils.py b/tool/utils.py index 6593c35a..3b19da1e 100644 --- a/tool/utils.py +++ b/tool/utils.py @@ -133,7 +133,8 @@ def get_color(c, x, max_val): blue = get_color(0, offset, classes) if color is None: rgb = (red, green, blue) - img = cv2.putText(img, class_names[cls_id], (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1.2, rgb, 1) + text = class_names[cls_id] + " " + str(cls_conf) + img = cv2.putText(img, text, (x1-25, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.4, rgb, 1) img = cv2.rectangle(img, (x1, y1), (x2, y2), rgb, 1) if savename: print("save plot results to %s" % savename)