嘿!您似乎在 United States,您想使用我们的 English 网站吗?
Switch to English site
Skip to main content

Pytorch深度学习框架X NVIDIA JetsonNano应用-YOLOv5辨识台湾实时路况

作者

嘉钧

难度

理论困难,实作普通

材料表

  1. NVIDIA Jestson Nano Developer Kit
  2. 64 GB SD卡
  3. 电源供应器
  4. 无线网卡(可以用网络线代替)

前景提要

当时YOLOv4出了没多久YOLOv5就悄然推出了,但是你可以发现v5其实不是v4的团队 (Alexey Bochkovskiy, Chien-Yao Wang, Hong-Yuan Mark Liao) 所做,更不是YOLO之父 (Joseph Redmon) 所做,而是一间名为Ultralytics LLC的公司所开发 ( 之前一直有发布关于YOLO转成PyTorch的Github ),在YOLOv5发布之前也没发布论文来左证YOLOv5,很多人对它的存在感到怀疑也闹出很大的风波。不过它是基于PyTorch实现,架构与YOLOv3、v4的DarkNet环境截然不同,对于再修改跟开发上比较简单,接着就跟我一起利用Jetson Nano来完成YOLOv5的实作吧!

如果你想了解更多可以看看Ultralytics LLC出面说明YOLOv4与YOLOv5差异-https://blog.roboflow.com/yolov4-versus-yolov5/,而Ultralytics也有推出基于YOLO的APP ( 仅限IOS ),现在也更新到YOLOv5了。

1_i_detection1_ff8706754bd4d3940fabdae383c8e0c5b256f55e.png

训练环境

先前都直接用原生的系统来装套件,有时候不同的项目会需要不同的版本,为了将其独立开来建议是使用虚拟环境来安装比较合适。今天会稍微介绍一下虚拟环境的部分,我在Windows上常会使用Anaconda而Jetson Nano上因为Anaconda不支持aarch64 (Arm64) 的核心所以要另外编译,非常麻烦!所以我直接使用Virtualenv ( 另一个轻便的虚拟环境套件)。

安装 virtualenv以及virtualenvwrapper

$ sudo pip3 install virtualenv virtualenvwrapper

修改环境变量:

$ nano ~/.bashrc
export WORKON_HOME=$HOME/.virtualenvs
export VIRTUALENVWRAPPER_PYTHON=/usr/bin/python3
export /usr/local/bin/virtualenvwrapper.sh

建置虚拟环境:

$ mkvirtualenv yolov5

开启虚拟环境:

$ workon yolov5

可以看到前面会多一个括号(env_name)就是你目前的环境名称:

2_virtualenv1_d74f444ff3ca13d77873f97ed6e5549cf439eff7.png

接下来先安装git,因为要下载YOLOv5

$ sudo apt-get install git-all -y

下载YOLOv5 Github

$ git clone https://github.com/ultralytics/yolov5.git

接着安装所需套件,可以先打开requirements.txt来看看所需套件,不管是Raspberry Pi 还是 Jetson系列,安装PyTorch或OpenCV都有特定的方式或来源,所以我在尝试别人的Github时都会分别开来安装。

3_required_txt1_2f245ed24e4e4056812b0479103c013e0574053e.png

像是这边可以看到Cython、Numpy都在安装PyTorch的时候会一起安装,而OpenCV因为原生就有所以用Link的方式就可以了,所以我们先来处理比较特别的OpenCV跟PyTorch。

4_nano_pytorch1_5556ce099eebfaf86e22457c78052bdfb8de7d78.png

这是官方提供的教学PyTorch for Jetson - version 1.6.0 now available,首先在YOLOv5提出的安装套件可以看到建议是1.6以上,目前只有JetPack4.4才能支持PyTorch 1.6哦!请特别注意自己的JetPack版本

$ sudo apt-cache show nvidia-jetpack

5_nano_jetpack1_a6e897bebccca25a0114e550fb3513415f8d18ef.png

安装PyTorch

$ wget https://nvidia.box.com/shared/static/9eptse6jyly1ggt9axbja2yrmj6pbarc.whl -O torch-1.6.0-cp36-cp36m-linux_aarch64.whl
$ sudo apt-get install python3-pip libopenblas-base libopenmpi-dev 
$ pip3 install Cython
$ pip3 install numpy torch-1.6.0-cp36-cp36m-linux_aarch64.whl

安装torchvision

$ sudo apt-get install libjpeg-dev zlib1g-dev
$ git clone --branch v0.7.0 https://github.com/pytorch/vision torchvision
$ cd torchvision
$ export BUILD_VERSION=0.7.0 
$ sudo python setup.py install

这部分大概10分钟内能搞定,可以透过导入函式库查看版本来确认是否安装成功:

6_nano_pytorch_check1_623e049e44666977bcd931f3faa7395d51e61fa2.png

接着要找到原生的OpenCV位置:

$ sudo find / -name cv2

7_nano_findcv1_b6127d7b472ae28245028ed6cfa65485e7826e08.png

寻找 .so 档案,大家应该都一样会在 /usr/lib/python3.6/dist-packages/cv2/python-3.6/ 里面:

8_nano_getcv1_5ddb91bb730a701f067e182e18b3379c87d56dcd.png

接着就要建立链接,使用ln 指令:

$ ln -s /usr/lib/python3.6/dist-packages/cv2/python-3.6/cv2.cpython-36m-aarch64-linux-gnu.so ~/.virtualenvs/{your env}/cv2.cpython-36m-aarch64-linux-gnu.so

9_nano_link_cv1_87a472b8115a855145c434b0d75cacd1cb07a10b.png

确认是否安装成功:

10_nano_check_cv1_8a3bf0b6de371d76358c3327ac73236bdbfbfa1f.png

剩下还没安装的套件整理一下会变成下面

$ pip3 install matplotlib pillow pyyaml tensorboard tqdm scipy

其中scipy可能原本就有了,由于它安装要好一阵子所以我建议如果有先执行看看,不行再安装。

执行YOLOv5的范例程序

使用Github范例程序 detect.py,范例程序主要会用到的自变量

--source

图片、影片目录或是 0 开启摄影机

--iou_thred

信心指数的阀值,虽然低一些可框出越多东西但准确度就不敢保证

--weights

权重,有s、m、l、x之分

我们可以下载训练好的weights,可以利用作者的download_weight.sh来下载,他需要用到util文件夹的程序所以我有移动到上一层目录,嫌麻烦的当然也可以直接到他的GoogleDrive下载。

$ cd weights/
$ mv download_weights.sh ..
$ cd ..
$ bash download_weights.sh

接着可以执行detec.py,我们先使用官网提供的范例图来测试:

$ python detect.py --source inference/images/ --weights weights/yolov5s.pt

运行结果如下:

11_run_yolo51_f0cef8c6c417449160eb20bc37f77e026423e76a.png

对两张图片进行推论,耗费时间为18秒:

12_1_bus_org1_e98a8475267f82103134c8a942612e95ba83bd07.jpg

12_2_bus_res2_d8626ac69e4ee58ee4a40d59b29bd250e87fa46f.jpg

13_1_zidane_org1_244323d057afd16e6e741fdb44d7f0cb207098f2.jpg

13_2_zidane_res1_83adc5eaec9ed424964e55f0995d0aecc295504f.jpg

还可以接上相机做实时影像侦测:

14_yolo5_camera1_03fa40683e07eda488a2c4ac294ac04070d01360.png

比较各模型差异

这边我们拿YOLOv3跟YOLOv5来做比较,可以注意到v3-tiny虽然秒数最少但运行结果不尽理想;然后v5-s目前看起来秒数少,框出来的物品多准确度也蛮高的;v5-l 、v3-spp准确度高但是也会框到一些错误的物品。

15_yolo35_compared1_e8e8242e89f451a362e9808f5baa84b38c8d01f0.png

拿YOLOv5来应用在实时路况影像

对于Jetson 系列的开发版相当多人会拿来做自驾车项目,而YOLO所训练的coco_2017数据集也能用于侦测人、车,所以我们先直接拿pre-trained model来实际运行看看,第一步是要取得到实时路况影像,这样类型的影像直接用手机路就太过时了,所以我们来玩点不一样的,我们可以到到下列这个网站获取「实时影像监视器

https://tw.live/

这个网站有各式各样的台湾路况可以查看,而这些都是实时影像。

16_taiwain_road_conditions1_430ae9711c4737f041bf83953123466d1d3a22e1.png

仔细看了一下西门町的路口监视器画面比较清晰也人多,所以最后我选择西门町的实时影像,接下来就要考虑一个大问题了~我该如何将直播影片给下载下来!

17_ximending1_f22d13459fd5b8b5320b125a25e1e8ffeafdb09f.png

其实你可以发现它从Youtube直播影片连动过来的,所以我写了这支程序用来撷取Youtube影像直播,主要利用pafy跟vlc来下载mp4影片,并且利用moviepy来剪辑预设秒数,首先先安装相关套件:

$ pip install pafy youtube-dl python-vlc moviepy

因为moviepy跟影片有关要安装相关的编码格式,Windows本身就有了但是Linux需要额外安装,透过pip安装在虚拟环境中是行不通的,目前安装在本身的环境:

$ sudo apt install ffmpeg

接着就是主要程序的部分:

def capture_video(opt):

    f_name = 'org.mp4'      # 下载影片名称
    o_name = opt.output     # 裁剪影片名称     
    sec = opt.second        # 欲保留秒数
    video = pafy.new(opt.url)   # 取得Youtube影片
    
    r_list = video.allstreams   # 取得串流来源列表

    print_div()
    for i,j in enumerate(r_list): print( '[ {} ] {} {}'.format(i,j.title,j))
    idx = input('\nChoose the resolution : ')
    
    if idx:
        ### 选择串流来源
        trg = r_list[int(idx)]
        print_div('您选择的分辨率是: %s'%(trg))
        
        ### 下载串流
        vlcInstance = vlc.Instance()
        player = vlcInstance.media_player_new()    # 创建一个新的MediaPlayer实例
        media = vlcInstance.media_new(trg.url)     # 创建一个新的Media实例
        media.get_mrl()
        media.add_option(f"sout=file/ts:{f_name}") # 将媒体直接储存 
        player.set_media(media)                    # 设置media_player将使用的媒体
        player.play()                              # 播放媒体
        time.sleep(1)                              # 等待信息跑完

        ### 进度条、撷取影片
        clock(sec)                                 # 播放 sec 秒 同时运行 进度条
        cut_video(f_name, sec, o_name)             # 裁切影片,因为停n秒长度不一定是n

        ### 关闭资源
        player.stop()                              # 关闭拨放器以及释放媒体

其余副函式,大部分是美观用,像是用来取得终端机窗口大小以及打印分隔符等,在clock的部分费了些心思写了类似tqdm的进度条,最后cut_video就是剪取影片,从第0秒到第n秒:

### 取得terminal窗口大小
def get_cmd_size():
    import shutil
    col, line = shutil.get_terminal_size()
    return col, line

### 打印分隔用
def print_div(text=None):
    col,line = get_cmd_size()
    col = col-1
    if text == None:
        print('-'*col)
    else:
        print('-'*col)
        print(text)
        print('-'*col)

### 计时、进度条
def clock(sec, sign = '#'):
    col, line = get_cmd_size()
    col_ = col - 42
    bar_bg = ' '*(col_)
    print_div()

    for i in range(sec+1):
        
        bar_idx = (col_//sec*(i+1))
        bar = ''
        for t in range(col_): bar += sign if t <= bar_idx else bar_bg[t]
        
        percent = int(100/sec*(i))
        end = '\n' if i==sec else '\r'
        print('Download Stream Video [{:02}/{:02}s] [{}] ({:02}%)'.format(i, sec, bar, percent), end=end)
        time.sleep(1)

### 撷取特定秒数并储存
def cut_video(name, sec, save_name):
    print_div()
    print('Cutting Video Used Moviepy\n')
    ffmpeg_extract_subclip(name, 0, sec, targetname=save_name)
    print_div(f'save file {save_name}')

为了使用更方便,我增加了argparse命令行选项,「-u」为Youtube链接;「-o」为输出影片名称;「-s」为输出影片秒数:

if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-u', '--url', help='youtube url')
    parser.add_argument('-o', '--output', type=str, default='sample.mp4' , help='save file path\name')
    parser.add_argument('-s', '--second',type=int, default=10 , help='video length')
    opt = parser.parse_args()
    capture_video(opt)

执行结果:

$ python capture_livestream.py -u 'https://www.youtube.com/watch?v=iVpOdRU0r9s&feature=emb_title&ab_channel=%E8%87%BA%E5%8C%97%E5%B8%82%E6%94%BF%E5%BA%9CTaipeiCityGovernment' -o test.mp4 -s 10

18_capture_stream_video2_904f0e42e24cc5be0c26967b7fbf772d3596f55b.png

接着可以直接执行范例程序来运行看看:

$ python detect.py --source test.mp4 --weights yolov5s.pt

19_run_sample_code_via_stream1_13c5b21919906761142ac7d7447abb2efc8447dd.png

结果如下:

20_streamvideo_result1_92df510440035ed007b81cdf0773c17b0a7ef33a.png

使用Jetson Nano运行平均0.165秒一帧,10秒的影片总共耗费83秒完成,这边提供运行完的影片给大家参考:

个人觉得这样的影片无法评估Nano效能是好是坏,所以我修改了一下范例程序将它变成实时影像辨识的方式,程序中稍微计算了FPS大概在5左右,一个顺畅的影片FPS至少要在30以上,所以可以看到有些许的卡顿,当然我用远程也有可能造成更多的Delay:

修改的内容相当简单,就是将原本要读取照片或影片的部分撷取出来,改成只有照片,并且在一开始读档的方式改成用OpenCV读取影像,最终修改后的程序如下:

import argparse
import os
import platform
import shutil
import time
from pathlib import Path
import numpy as np

import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random

from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords,
    xyxy2xywh, plot_one_box, strip_optimizer, set_logging)
from utils.torch_utils import select_device, load_classifier, time_synchronized

# Image process
###############################################################################

def img_preprocess(img0, img_size=640):
    
    # Padded resize
    img = letterbox(img0, new_shape=img_size)[0]

    # Convert
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
    img = np.ascontiguousarray(img)

    # cv2.imwrite(path + '.letterbox.jpg', 255 * img.transpose((1, 2, 0))[:, :, ::-1])  # save letterbox image
    return img, img0

# Get target size
###############################################################################

def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):
    # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
    shape = img.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better test mAP)
        r = min(r, 1.0)

    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, 64), np.mod(dh, 64)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return img, ratio, (dw, dh)

def print_div(text):

    print(text, '\n')
    print('='*40,'\n')

# Detect func
###############################################################################
def detect(save_img=False):
    print_div('INTIL')
    out, source, weights, view_img, save_txt, imgsz = \
        opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size

    # Initialize
    print_div('GET DEVICE')
    set_logging()
    device = select_device(opt.device)
    half = device.type != 'cpu'  # half precision only supported on CUDA
    
    # Load model
    print_div('LOAD MODEL')
    model = attempt_load(weights, map_location=device)  # load FP32 model
    imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size
    if half:
        model.half()  # to FP16

    # Second-stage classifier
    print_div('LOAD MODEL_CLASSIFIER')
    classify = False
    if classify:
        modelc = load_classifier(name='resnet101', n=2)  # initialize
        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model'])  # load weights
        modelc.to(device).eval()

    # Get names and colors
    print_div('SET LABEL COLOR')
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

    # Run inference
    ###############################################################################
    print_div("RUN INFERENCE")
    
    img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
    _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once

    video_path = source
    cap = cv2.VideoCapture(video_path)

    print_div('Start Play VIDEO')
    while cap.isOpened():
        
        ret, frame = cap.read()
        t0 = time.time()
        
        if not ret: 
            print_div('No Frame')
            break

        fps_t1 = time.time()

        img, img0 = img_preprocess(frame)   # img: Resize , img0:Orginal
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        t1 = time_synchronized()
        pred = model(img, augment=opt.augment)[0]

        # Apply NMS : 取得每项预测的数值
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        t2 = time_synchronized()

        # Apply Classifier : 取得该数值的LAbel
        if classify:
            pred = apply_classifier(pred, modelc, img, img0)
        
        # Draw Box
        for i, det in enumerate(pred):

            s = '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(img0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += '%g %ss, ' % (n, names[int(c)])  # add to string

                # Write results
                for *xyxy, conf, cls in reversed(det):

                    label = '%s %.2f' % (names[int(cls)], conf)
                    plot_one_box(xyxy, img0, label=label, color=colors[int(cls)], line_thickness=3)

        # Print Results(inference + NMS)
        print_div('%sDone. (%.3fs)' % (s, t2 - t1))

        # Draw Image
        x, y, w, h = (img0.shape[1]//4), 25, (img0.shape[1]//2), 30
        cv2.rectangle(img0, (x, 10),(x+w, y+h), (0,0,0), -1)
        
        rescale = 0.5
        re_img0 = (int(img0.shape[1]*rescale) ,int(img0.shape[0]*rescale))

        cv2.putText(img0, '{} | inference: {:.4f}s | fps: {:.4f}'.format(opt.weights[0], t2-t1, 1/(time.time()-t0)),(x+20, y+20),cv2.FONT_HERSHEY_SIMPLEX,1,(0,0,255),2)
        cv2.imshow('Stream_Detected', cv2.resize(img0, re_img0) )

        key = cv2.waitKey(1)
        if key == ord('q'): break

    # After break
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
    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('--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('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    opt = parser.parse_args()
    print(opt)

    with torch.no_grad():
        if opt.update:  # update all models (to fix SourceChangeWarning)
            for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
                detect()
                strip_optimizer(opt.weights)
        else:
            detect()

结语

YOLOv5纵使不是正统、最快的YOLO但是基于PyTorch实做的YOLO让我们修改更加的方便,以往在DarkNet上运行现在只要装好PyTorch基本就可以执行,档案大小也差非常多,边缘装置的负担也不会太大!大家可以去体验看看YOLOv5的方便性、轻便性,下一篇文章中我会教大家如何使用YOLOv5来训练自己的数据!

相关文章

在Jetson Nano (TX1/TX2)上使用Anaconda与PyTorch 1.1.0

https://zhuanlan.zhihu.com/p/64868319

YOLOv5 github

https://github.com/ultralytics/yolov5

CAVEDU Education is devoted into robotics education and maker movement since 2008, and is intensively active in teaching fundamental knowledge and skills. We had published many books for readers in all ages, topics including Deep Learning, edge computing, App Inventor, IoT and robotics. Please check CAVEDU's website for more information: http://www.cavedu.com, http://www.appinventor.tw
DesignSpark Electrical Logolinkedin