From 659ad748c810bf0d6f5b5178e6aa1048780683ba Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 4 Jul 2020 17:13:43 -0700 Subject: [PATCH] update get_voc.sh --- Dockerfile | 2 +- data/get_voc.sh | 2 +- detect.py | 6 +++--- models/export.py | 12 ++++++++++++ 4 files changed, 17 insertions(+), 5 deletions(-) diff --git a/Dockerfile b/Dockerfile index 01551a0..bda78c0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,5 @@ # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch -FROM nvcr.io/nvidia/pytorch:20.06-py3 +FROM nvcr.io/nvidia/pytorch:20.03-py3 RUN pip install -U gsutil # Create working directory diff --git a/data/get_voc.sh b/data/get_voc.sh index 166370f..949e0cc 100644 --- a/data/get_voc.sh +++ b/data/get_voc.sh @@ -1,5 +1,5 @@ # PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/ -# Download command: bash yolov5/data/get_voc.sh +# Download command: bash ./data/get_voc.sh # Train command: python train.py --data voc.yaml # Dataset should be placed next to yolov5 folder: # /parent_folder diff --git a/detect.py b/detect.py index 2650c20..f9b12eb 100644 --- a/detect.py +++ b/detect.py @@ -102,7 +102,7 @@ def detect(save_img=False): if save_img or view_img: # Add bbox to image label = '%s %.2f' % (names[int(cls)], conf) - plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) + plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1) # Print time (inference + NMS) print('%sDone. (%.3fs)' % (s, t2 - t1)) @@ -139,10 +139,10 @@ def detect(save_img=False): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='model.pt path') + parser.add_argument('--weights', type=str, default='weights/yolov5m.pt', help='model.pt path') parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder - parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') + parser.add_argument('--img-size', type=int, default=1024, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)') diff --git a/models/export.py b/models/export.py index 54d7a12..c11c0a3 100644 --- a/models/export.py +++ b/models/export.py @@ -56,5 +56,17 @@ if __name__ == '__main__': except Exception as e: print('ONNX export failure: %s' % e) + # CoreML export + try: + import coremltools as ct + + print('\nStarting CoreML export with coremltools %s...' % ct.__version__) + model = ct.convert(ts, inputs=[ct.ImageType(name='images', shape=img.shape)]) # convert + f = opt.weights.replace('.pt', '.mlmodel') # filename + model.save(f) + print('CoreML export success, saved as %s' % f) + except Exception as e: + print('CoreML export failure: %s' % e) + # Finish print('\nExport complete. Visualize with https://github.com/lutzroeder/netron.')