上传主要代码

Pipfile

@@ -0,0 +1,19 @@
+[[source]]
+name = "pypi"
+url = "https://pypi.org/simple"
+verify_ssl = true
+
+[dev-packages]
+
+[packages]
+streamlit = ">0.49.0"
+opencv-python = "*"
+numpy = "*"
+torchvision = "0.9.1"
+torch = "1.8.1"
+Pillow = "8.2.0"
+pyyaml = "6.0"
+matplotlib = "*"
+opencv-python-headless = "4.5.2.52"
+av = "*"
+streamlit-webrtc = "0.36.1"

app.py

@@ -0,0 +1,400 @@
+"""Create an Object Detection Web App using PyTorch and Streamlit."""
+# import libraries
+from PIL import Image
+from torchvision import models, transforms
+import torch
+import streamlit as st
+from yolo import YOLO
+import os
+import urllib
+import numpy as np
+from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
+import av
+# 设置网页的icon
+st.set_page_config(page_title='Gesture Detector', page_icon='✌',
+                   layout='centered', initial_sidebar_state='expanded')
+
+RTC_CONFIGURATION = RTCConfiguration(
+    {
+      "RTCIceServer": [{
+        "urls": ["stun:stun.l.google.com:19302"],
+        "username": "pikachu",
+        "credential": "1234",
+      }]
+    }
+)
+def main():
+    # Render the readme as markdown using st.markdown.
+    readme_text = st.markdown(open("instructions.md",encoding='utf-8').read())
+
+    
+    # Once we have the dependencies, add a selector for the app mode on the sidebar.
+    st.sidebar.title("What to do")
+    app_mode = st.sidebar.selectbox("Choose the app mode",
+        ["Show instructions", "Run the app", "Show the source code"])
+    if app_mode == "Show instructions":
+        st.sidebar.success('To continue select "Run the app".')
+    elif app_mode == "Show the source code":
+        readme_text.empty()
+        st.code(open("app.py",encoding='utf-8').read())
+    elif app_mode == "Run the app":
+        # Download external dependencies.
+        for filename in EXTERNAL_DEPENDENCIES.keys():
+            download_file(filename)
+
+        readme_text.empty()
+        run_the_app()
+
+# External files to download.
+EXTERNAL_DEPENDENCIES = {
+    "yolov4_tiny.pth": {
+        "url": "https://github.com/Dreaming-future/my_weights/releases/download/v1.3/yolov4_tiny.pth",
+        "size": 23631189 
+    },
+    "yolov4_SE.pth": {
+        "url": "https://github.com/Dreaming-future/my_weights/releases/download/v1.3/yolov4_SE.pth",
+        "size": 23806027
+    },
+    "yolov4_CBAM.pth":{
+        "url": "https://github.com/Dreaming-future/my_weights/releases/download/v1.3/yolov4_CBAM.pth",
+        "size": 23981478
+    },
+    "yolov4_ECA.pth":{
+        "url": "https://github.com/Dreaming-future/my_weights/releases/download/v1.3/yolov4_ECA.pth",
+        "size": 23632688
+    },
+    "yolov4_weights_ep150_608.pth":{
+        "url": "https://github.com/Dreaming-future/my_weights/releases/download/v1.3/yolov4_weights_ep150_608.pth",
+        "size": 256423031
+    },
+    "yolov4_weights_ep150_416.pth":{
+        "url": "https://github.com/Dreaming-future/my_weights/releases/download/v1.3/yolov4_weights_ep150_416.pth",
+        "size": 256423031
+    },
+}
+
+
+# This file downloader demonstrates Streamlit animation.
+def download_file(file_path):
+    # Don't download the file twice. (If possible, verify the download using the file length.)
+    if os.path.exists(file_path):
+        if "size" not in EXTERNAL_DEPENDENCIES[file_path]:
+            return
+        elif os.path.getsize(file_path) == EXTERNAL_DEPENDENCIES[file_path]["size"]:
+            return
+    # print(os.path.getsize(file_path))
+    # These are handles to two visual elements to animate.
+    weights_warning, progress_bar = None, None
+    try:
+        weights_warning = st.warning("Downloading %s..." % file_path)
+        progress_bar = st.progress(0)
+        with open(file_path, "wb") as output_file:
+            with urllib.request.urlopen(EXTERNAL_DEPENDENCIES[file_path]["url"]) as response:
+                length = int(response.info()["Content-Length"])
+                counter = 0.0
+                MEGABYTES = 2.0 ** 20.0
+                while True:
+                    data = response.read(8192)
+                    if not data:
+                        break
+                    counter += len(data)
+                    output_file.write(data)
+
+                    # We perform animation by overwriting the elements.
+                    weights_warning.warning("Downloading %s... (%6.2f/%6.2f MB)" %
+                        (file_path, counter / MEGABYTES, length / MEGABYTES))
+                    progress_bar.progress(min(counter / length, 1.0))
+    except Exception as e:
+        print(e)
+    # Finally, we remove these visual elements by calling .empty().
+    finally:
+        if weights_warning is not None:
+            weights_warning.empty()
+        if progress_bar is not None:
+            progress_bar.empty()
+
+# This is the main app app itself, which appears when the user selects "Run the app".
+def run_the_app():    
+    class Config():
+        def __init__(self, weights = 'yolov4_tiny.pth', tiny = True, phi = 0, shape = 416,nms_iou = 0.3, confidence = 0.5):
+            self.weights = weights
+            self.tiny = tiny
+            self.phi = phi
+            self.cuda = False
+            self.shape = shape
+            self.confidence = confidence
+            self.nms_iou = nms_iou
+    # set title of app
+    st.markdown('<h1 align="center">✌ Gesture Detection</h1>',
+                unsafe_allow_html=True)
+    st.sidebar.markdown("# Gesture Detection on?")
+    activities = ["Example","Image", "Camera", "FPS", "Heatmap","Real Time", "Video"]
+    choice = st.sidebar.selectbox("Choose among the given options:", activities)
+    phi = st.sidebar.selectbox("yolov4-tiny 使用的自注意力模式:",('0tiny','1SE','2CABM','3ECA'))
+    print("")
+
+    tiny = st.sidebar.checkbox('是否使用 yolov4 tiny 模型')
+    if not tiny:
+        shape = st.sidebar.selectbox("Choose shape to Input:", [416,608])
+    conf,nms = object_detector_ui()
+    @st.cache
+    def get_yolo(tiny,phi,conf,nms,shape=416):
+        weights = 'yolov4_tiny.pth'
+        if tiny:
+            if phi == '0tiny':
+                weights = 'yolov4_tiny.pth'
+            elif phi == '1SE':
+                weights = 'yolov4_SE.pth'
+            elif phi == '2CABM':
+                weights = 'yolov4_CBAM.pth'
+            elif phi == '3ECA':
+                weights = 'yolov4_ECA.pth'
+        else:
+            if shape == 608:
+                weights = 'yolov4_weights_ep150_608.pth'
+            elif shape == 416:
+                weights = 'yolov4_weights_ep150_416.pth'
+        opt = Config(weights = weights, tiny = tiny , phi = int(phi[0]), shape = shape,nms_iou = nms, confidence = conf)
+        yolo = YOLO(opt)
+        return yolo
+    
+    if tiny:
+        yolo = get_yolo(tiny, phi, conf, nms)
+        st.write("YOLOV4 tiny 模型加载完毕")
+    else:
+        yolo = get_yolo(tiny, phi, conf, nms, shape)
+        st.write("YOLOV4 模型加载完毕")
+    
+    if choice == 'Image':
+        detect_image(yolo)
+    elif choice =='Camera':
+        detect_camera(yolo)
+    elif choice == 'FPS':
+        detect_fps(yolo)
+    elif choice == "Heatmap":
+        detect_heatmap(yolo)
+    elif choice == "Example":
+        detect_example(yolo)
+    elif choice == "Real Time":
+        detect_realtime(yolo)
+    elif choice == "Video":
+        detect_video(yolo)
+        
+
+
+# This sidebar UI lets the user select parameters for the YOLO object detector.
+def object_detector_ui():
+    st.sidebar.markdown("# Model")
+    confidence_threshold = st.sidebar.slider("Confidence threshold", 0.0, 1.0, 0.5, 0.01)
+    overlap_threshold = st.sidebar.slider("Overlap threshold", 0.0, 1.0, 0.3, 0.01)
+    return confidence_threshold, overlap_threshold
+
+def predict(image,yolo):
+    """Return predictions.
+
+    Parameters
+    ----------
+    :param image: uploaded image
+    :type image: jpg
+    :rtype: list
+    :return: none
+    """
+    crop            = False
+    count           = False
+    try:
+        # image = Image.open(image)
+        r_image = yolo.detect_image(image, crop = crop, count=count)
+        transform = transforms.Compose([transforms.ToTensor()])        
+        result = transform(r_image)
+        st.image(result.permute(1,2,0).numpy(), caption = 'Processed Image.', use_column_width = True)
+    except Exception as e:
+        print(e)
+
+def fps(image,yolo):
+    test_interval = 50
+    tact_time = yolo.get_FPS(image, test_interval)
+    st.write(str(tact_time) + ' seconds, ', str(1/tact_time),'FPS, @batch_size 1')
+    return tact_time
+    # print(str(tact_time) + ' seconds, ' + str(1/tact_time) + 'FPS, @batch_size 1')
+
+
+def detect_image(yolo):
+    # enable users to upload images for the model to make predictions
+    file_up = st.file_uploader("Upload an image", type = ["jpg","png","jpeg"])
+    classes = ["up","down","left","right","front","back","clockwise","anticlockwise"]
+    class_to_idx = {cls: idx for (idx, cls) in enumerate(classes)}
+    st.sidebar.markdown("See the model preformance and play with it")
+    if file_up is not None:
+        with st.spinner(text='Preparing Image'):
+            # display image that user uploaded
+            image = Image.open(file_up)
+            st.image(image, caption = 'Uploaded Image.', use_column_width = True)
+            st.balloons()
+            detect = st.button("开始检测Image")
+            if detect:
+                st.write("")
+                st.write("Just a second ...")
+                predict(image,yolo)
+                st.balloons()
+
+
+
+def detect_camera(yolo):
+    picture = st.camera_input("Take a picture")
+    if picture:
+        filters_to_funcs = {
+            "No filter": predict,
+            "Heatmap": heatmap,
+            "FPS": fps,
+        }
+        filters = st.selectbox("...and now, apply a filter!", filters_to_funcs.keys())
+        image = Image.open(picture)
+        with st.spinner(text='Preparing Image'):
+            filters_to_funcs[filters](image,yolo)
+            st.balloons()
+
+def detect_fps(yolo):
+    file_up = st.file_uploader("Upload an image", type = ["jpg","png","jpeg"])
+    classes = ["up","down","left","right","front","back","clockwise","anticlockwise"]
+    class_to_idx = {cls: idx for (idx, cls) in enumerate(classes)}
+    st.sidebar.markdown("See the model preformance and play with it")
+    if file_up is not None:
+        # display image that user uploaded
+        image = Image.open(file_up)
+        st.image(image, caption = 'Uploaded Image.', use_column_width = True)
+        st.balloons()
+        detect = st.button("开始检测 FPS")
+        if detect:
+            with st.spinner(text='Preparing Image'):
+                st.write("")
+                st.write("Just a second ...")
+                tact_time = fps(image,yolo)
+                # st.write(str(tact_time) + ' seconds, ', str(1/tact_time),'FPS, @batch_size 1')
+                st.balloons()
+
+def heatmap(image,yolo):
+    heatmap_save_path = "heatmap_vision.png"
+    yolo.detect_heatmap(image, heatmap_save_path)
+    img = Image.open(heatmap_save_path)
+    transform = transforms.Compose([transforms.ToTensor()])        
+    result = transform(img)
+    st.image(result.permute(1,2,0).numpy(), caption = 'Processed Image.', use_column_width = True)
+
+def detect_heatmap(yolo):
+    file_up = st.file_uploader("Upload an image", type = ["jpg","png","jpeg"])
+    classes = ["up","down","left","right","front","back","clockwise","anticlockwise"]
+    class_to_idx = {cls: idx for (idx, cls) in enumerate(classes)}
+    st.sidebar.markdown("See the model preformance and play with it")
+    if file_up is not None:
+        # display image that user uploaded
+        image = Image.open(file_up)
+        st.image(image, caption = 'Uploaded Image.', use_column_width = True)
+        st.balloons()
+        detect = st.button("开始检测 heatmap")
+        if detect:
+            with st.spinner(text='Preparing Heatmap'):
+                st.write("")
+                st.write("Just a second ...")
+                heatmap(image,yolo)
+                st.balloons()
+
+def detect_example(yolo):
+    st.sidebar.title("Choose an Image as a example")
+    images = os.listdir('./img')
+    images.sort()
+    image = st.sidebar.selectbox("Image Name", images)
+    st.sidebar.markdown("See the model preformance and play with it")
+    image = Image.open(os.path.join('img',image))
+    st.image(image, caption = 'Choose Image.', use_column_width = True)
+    st.balloons()
+    detect = st.button("开始检测Image")
+    if detect:
+        st.write("")
+        st.write("Just a second ...")
+        predict(image,yolo)
+        st.balloons()
+
+def detect_realtime(yolo):
+
+    class VideoProcessor:
+        def recv(self, frame):
+            img = frame.to_ndarray(format="bgr24")
+            img = Image.fromarray(img)
+            crop            = False
+            count           = False
+            r_image = yolo.detect_image(img, crop = crop, count=count)
+            transform = transforms.Compose([transforms.ToTensor()])        
+            result = transform(r_image)
+            result = result.permute(1,2,0).numpy()
+            result = (result * 255).astype(np.uint8)
+            return av.VideoFrame.from_ndarray(result, format="bgr24")
+       
+    webrtc_ctx = webrtc_streamer(
+        key="example",
+        mode=WebRtcMode.SENDRECV,
+        rtc_configuration=RTC_CONFIGURATION,
+        media_stream_constraints={"video": True, "audio": False},
+        async_processing=False,
+        video_processor_factory=VideoProcessor
+    )
+
+import cv2
+import time
+def detect_video(yolo):
+    file_up = st.file_uploader("Upload a video", type = ["mp4"])
+    print(file_up)
+    classes = ["up","down","left","right","front","back","clockwise","anticlockwise"]
+    
+    if file_up is not None:
+        video_path = 'video.mp4'
+        st.video(file_up)
+        with open(video_path, 'wb') as f:
+            f.write(file_up.read())       
+        detect = st.button("开始检测 Video")
+        
+        if detect: 
+            video_save_path = 'video2.mp4'
+            # display image that user uploaded
+            capture = cv2.VideoCapture(video_path)
+            
+            video_fps = st.slider("Video FPS", 5, 30, int(capture.get(cv2.CAP_PROP_FPS)), 1)
+            fourcc  = cv2.VideoWriter_fourcc(*'XVID')
+            size    = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+            out     = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)
+
+
+            
+            while(True):
+                # 读取某一帧
+                ref, frame = capture.read()
+                if not ref:
+                    break
+                # 转变成Image
+                # frame = Image.fromarray(np.uint8(frame))
+                # 格式转变，BGRtoRGB
+                frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
+                # 转变成Image
+                frame = Image.fromarray(np.uint8(frame))
+                # 进行检测
+                frame = np.array(yolo.detect_image(frame))
+                # RGBtoBGR满足opencv显示格式
+                frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
+
+                # print("fps= %.2f"%(fps))
+                # frame = cv2.putText(frame, "fps= %.2f"%(fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+                out.write(frame)
+                
+            out.release()
+            capture.release()
+            print("Save processed video to the path :" + video_save_path)
+            
+            with open(video_save_path, "rb") as file:
+                btn = st.download_button(
+                        label="Download Video",
+                        data=file,
+                        file_name="video.mp4",
+                    )
+            st.balloons()
+
+if __name__ == "__main__":
+    main()

get_yaml.py

@@ -0,0 +1,14 @@
+import os
+import sys
+import yaml
+
+def get_config():
+    yaml_path = 'model_data/gesture.yaml'
+    f = open(yaml_path,'r',encoding='utf-8')
+    config = yaml.load(f,Loader =yaml.FullLoader)
+    f.close()
+    return config
+
+if __name__ == "__main__":
+    config = get_config()
+    print(config)

instructions.md

@@ -0,0 +1,9 @@
+# ✌ Gesture Detection
+
+
+这是一个基于无人机视觉图像手势识别控制系统，选择了YOLOv4模型进行训练
+
+ **YOLOv4 = CSPDarknet53（主干） + SPP** **附加模块（颈** **） +** **PANet** **路径聚合（颈** **） + YOLOv3（头部）**
+
+![img](https://pdf.cdn.readpaper.com/parsed/fetch_target/699143cdb334ecfc63caf8192472490c_0_Figure_1.png)
+

packages.txt

@@ -0,0 +1,2 @@
+freeglut3-dev
+libgtk2.0-dev

predict.py

@@ -0,0 +1,194 @@
+#-----------------------------------------------------------------------#
+#   predict.py将单张图片预测、摄像头检测、FPS测试和目录遍历检测等功能
+#   整合到了一个py文件中，通过指定mode进行模式的修改。
+#-----------------------------------------------------------------------#
+import time
+import yaml
+import cv2
+import numpy as np
+from PIL import Image
+from get_yaml import get_config
+from yolo import YOLO
+import argparse
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--weights',type=str,default='model_data/yolotiny_SE_ep100.pth',help='initial weights path')
+    parser.add_argument('--tiny',action='store_true',help='使用yolotiny模型')
+    parser.add_argument('--phi',type=int,default=1,help='yolov4tiny注意力机制类型')
+    parser.add_argument('--mode',type=str,choices=['dir_predict', 'video', 'fps','predict','heatmap','export_onnx'],default="dir_predict",help='预测的模式')
+    parser.add_argument('--cuda',action='store_true',help='表示是否使用GPU')
+    parser.add_argument('--shape',type=int,default=416,help='输入图像的shape')
+    parser.add_argument('--video',type=str,default='',help='需要检测的视频文件')
+    parser.add_argument('--save-video',type=str,default='',help='保存视频的位置')
+    parser.add_argument('--confidence',type=float,default=0.5,help='只有得分大于置信度的预测框会被保留下来')
+    parser.add_argument('--nms_iou',type=float,default=0.3,help='非极大抑制所用到的nms_iou大小')
+    opt = parser.parse_args()
+    print(opt)
+
+    # 配置文件
+    config = get_config()
+    yolo = YOLO(opt)
+
+    #----------------------------------------------------------------------------------------------------------#
+    #   mode用于指定测试的模式：
+    #   'predict'           表示单张图片预测，如果想对预测过程进行修改，如保存图片，截取对象等，可以先看下方详细的注释
+    #   'video'             表示视频检测，可调用摄像头或者视频进行检测，详情查看下方注释。
+    #   'fps'               表示测试fps，使用的图片是img里面的street.jpg，详情查看下方注释。
+    #   'dir_predict'       表示遍历文件夹进行检测并保存。默认遍历img文件夹，保存img_out文件夹，详情查看下方注释。
+    #   'heatmap'           表示进行预测结果的热力图可视化，详情查看下方注释。
+    #   'export_onnx'       表示将模型导出为onnx，需要pytorch1.7.1以上。
+    #----------------------------------------------------------------------------------------------------------#
+    mode = opt.mode
+    #-------------------------------------------------------------------------#
+    #   crop                指定了是否在单张图片预测后对目标进行截取
+    #   count               指定了是否进行目标的计数
+    #   crop、count仅在mode='predict'时有效
+    #-------------------------------------------------------------------------#
+    crop            = False
+    count           = False
+    #----------------------------------------------------------------------------------------------------------#
+    #   video_path          用于指定视频的路径，当video_path=0时表示检测摄像头
+    #                       想要检测视频，则设置如video_path = "xxx.mp4"即可，代表读取出根目录下的xxx.mp4文件。
+    #   video_save_path     表示视频保存的路径，当video_save_path=""时表示不保存
+    #                       想要保存视频，则设置如video_save_path = "yyy.mp4"即可，代表保存为根目录下的yyy.mp4文件。
+    #   video_fps           用于保存的视频的fps
+    #
+    #   video_path、video_save_path和video_fps仅在mode='video'时有效
+    #   保存视频时需要ctrl+c退出或者运行到最后一帧才会完成完整的保存步骤。
+    #----------------------------------------------------------------------------------------------------------#
+    video_path      = 0 if opt.video == '' else opt.video
+    video_save_path = opt.save_video
+    video_fps       = 25.0
+    #----------------------------------------------------------------------------------------------------------#
+    #   test_interval       用于指定测量fps的时候，图片检测的次数。理论上test_interval越大，fps越准确。
+    #   fps_image_path      用于指定测试的fps图片
+    #   
+    #   test_interval和fps_image_path仅在mode='fps'有效
+    #----------------------------------------------------------------------------------------------------------#
+    test_interval   = 100
+    fps_image_path  = "img/up.jpg"
+    #-------------------------------------------------------------------------#
+    #   dir_origin_path     指定了用于检测的图片的文件夹路径
+    #   dir_save_path       指定了检测完图片的保存路径
+    #   
+    #   dir_origin_path和dir_save_path���在mode='dir_predict'时有效
+    #-------------------------------------------------------------------------#
+    dir_origin_path = "img/"
+    dir_save_path   = "img_out/"
+    #-------------------------------------------------------------------------#
+    #   heatmap_save_path   热力图的保存路径，默认保存在model_data下
+    #   
+    #   heatmap_save_path仅在mode='heatmap'有效
+    #-------------------------------------------------------------------------#
+    heatmap_save_path = "model_data/heatmap_vision.png"
+    #-------------------------------------------------------------------------#
+    #   simplify            使用Simplify onnx
+    #   onnx_save_path      指定了onnx的保存路径
+    #-------------------------------------------------------------------------#
+    simplify        = True
+    onnx_save_path  = "model_data/models.onnx"
+
+    if mode == "predict":
+        '''
+        1、如果想要进行检测完的图片的保存，利用r_image.save("img.jpg")即可保存，直接在predict.py里进行修改即可。 
+        2、如果想要获得预测框的坐标，可以进入yolo.detect_image函数，在绘图部分读取top，left，bottom，right这四个值。
+        3、如果想要利用预测框截取下目标，可以进入yolo.detect_image函数，在绘图部分利用获取到的top，left，bottom，right这四个值
+        在原图上利用矩阵的方式进行截取。
+        4、如果想要在预测图上写额外的字，比如检测到的特定目标的数量，可以进入yolo.detect_image函数，在绘图部分对predicted_class进行判断，
+        比如判断if predicted_class == 'car': 即可判断当前目标是否为车，然后记录数量即可。利用draw.text即可写字。
+        '''
+        while True:
+            img = input('Input image filename:')
+            try:
+                image = Image.open(img)
+            except:
+                print('Open Error! Try again!')
+                continue
+            else:
+                r_image = yolo.detect_image(image, crop = crop, count=count)
+                r_image.show()
+                r_image.save(dir_save_path + 'img_result.jpg')
+
+    elif mode == "video":
+        capture = cv2.VideoCapture(video_path)
+        if video_save_path != '':
+            fourcc  = cv2.VideoWriter_fourcc(*'XVID')
+            size    = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+            out     = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)
+
+        ref, frame = capture.read()
+        if not ref:
+            raise ValueError("未能正确读取摄像头（视频），请注意是否正确安装摄像头（是否正确填写视频路径）。")
+
+        fps = 0.0
+        while(True):
+            t1 = time.time()
+            # 读取某一帧
+            ref, frame = capture.read()
+            if not ref:
+                break
+            # 格式转变，BGRtoRGB
+            frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
+            # 转变成Image
+            frame = Image.fromarray(np.uint8(frame))
+            # 进行检测
+            frame = np.array(yolo.detect_image(frame))
+            # RGBtoBGR满足opencv显示格式
+            frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
+            
+            fps  = ( fps + (1./(time.time()-t1)) ) / 2
+            print("fps= %.2f"%(fps))
+            frame = cv2.putText(frame, "fps= %.2f"%(fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+            
+            cv2.imshow("video",frame)
+            c= cv2.waitKey(1) & 0xff 
+            if video_save_path != '':
+                out.write(frame)
+
+            if c==27:
+                capture.release()
+                break
+
+        print("Video Detection Done!")
+        capture.release()
+        if video_save_path != '':
+            print("Save processed video to the path :" + video_save_path)
+            out.release()
+        cv2.destroyAllWindows()
+        
+    elif mode == "fps":
+        img = Image.open(fps_image_path)
+        tact_time = yolo.get_FPS(img, test_interval)
+        print(str(tact_time) + ' seconds, ' + str(1/tact_time) + 'FPS, @batch_size 1')
+
+    elif mode == "dir_predict":
+        import os
+
+        from tqdm import tqdm
+
+        img_names = os.listdir(dir_origin_path)
+        for img_name in tqdm(img_names):
+            if img_name.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
+                image_path  = os.path.join(dir_origin_path, img_name)
+                image       = Image.open(image_path)
+                r_image     = yolo.detect_image(image)
+                if not os.path.exists(dir_save_path):
+                    os.makedirs(dir_save_path)
+                r_image.save(os.path.join(dir_save_path, img_name.replace(".jpg", ".png")), quality=95, subsampling=0)
+
+    elif mode == "heatmap":
+        while True:
+            img = input('Input image filename:')
+            try:
+                image = Image.open(img)
+            except:
+                print('Open Error! Try again!')
+                continue
+            else:
+                yolo.detect_heatmap(image, heatmap_save_path)
+                
+    elif mode == "export_onnx":
+        yolo.convert_to_onnx(simplify, onnx_save_path)
+        
+    else:
+        raise AssertionError("Please specify the correct mode: 'predict', 'video', 'fps', 'heatmap', 'export_onnx', 'dir_predict'.")

requirements.txt

@@ -0,0 +1,15 @@
+scipy
+numpy
+matplotlib
+opencv_python
+torch==1.8.1
+torchvision==0.9.1
+tqdm==4.60.0
+Pillow==8.2.0
+h5py==2.10.0
+tensorboard
+pyyaml==6.0
+torchinfo
+labelimg==1.8.6
+streamlit==1.8.1
+opencv-python-headless==4.5.2.52

yolo.py

@@ -0,0 +1,422 @@
+import colorsys
+import os
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import ImageDraw, ImageFont
+
+from nets.yolo import YoloBody
+from nets.yolo_tiny import YoloBodytiny
+from utils.utils import (cvtColor, get_anchors, get_classes, preprocess_input,
+                         resize_image)
+from utils.utils_bbox import DecodeBox
+from get_yaml import get_config
+import argparse
+'''
+训练自己的数据集必看注释！
+'''
+class YOLO(object):
+    # 配置文件
+    config = get_config()
+    _defaults = {
+        #--------------------------------------------------------------------------#
+        #   使用自己训练好的模型进行预测一定要修改model_path和classes_path！
+        #   model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
+        #
+        #   训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
+        #   验证集损失较低不代表mAP较高，仅代表该权值在验证集上泛化性能较好。
+        #   如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
+        #--------------------------------------------------------------------------#
+        "class_names"       : config['classes'],
+        "num_classes"       : config['nc'],
+        #---------------------------------------------------------------------#
+        #   anchors_path代表先验框对应的txt文件，一般不修改。
+        #   anchors_mask用于帮助代码找到对应的先验框，一般不修改。
+        #---------------------------------------------------------------------#
+        "anchors_path"      : 'model_data/yolo_anchors.txt',
+        "anchors_mask"      : [[6, 7, 8], [3, 4, 5], [0, 1, 2]],
+        #---------------------------------------------------------------------#
+        #   只有得分大于置信度的预测框会被保留下来
+        #---------------------------------------------------------------------#
+        "confidence"        : 0.5, # 0.5,
+        #---------------------------------------------------------------------#
+        #   非极大抑制所用到的nms_iou大小
+        #---------------------------------------------------------------------#
+        "nms_iou"           : 0.3, # 0.3,
+        #---------------------------------------------------------------------#
+        #   该变量用于控制是否使用letterbox_image对输入图像进行不失真的resize，
+        #   在多次测试后，发现关闭letterbox_image直接resize的效果更好
+        #---------------------------------------------------------------------#
+        "letterbox_image"   : config['letterbox_image'], # False,
+    }
+
+
+
+    @classmethod
+    def get_defaults(cls, n):
+        if n in cls._defaults:
+            return cls._defaults[n]
+        else:
+            return "Unrecognized attribute name '" + n + "'"
+
+    #---------------------------------------------------#
+    #   初始化YOLO
+    #---------------------------------------------------#
+    def __init__(self, opt, **kwargs):
+        self.__dict__.update(self._defaults)
+        for name, value in kwargs.items():
+            setattr(self, name, value)
+        self.phi = opt.phi
+        self.tiny = opt.tiny
+        self.cuda = opt.cuda
+        self.input_shape = [opt.shape,opt.shape]
+        self.model_path = opt.weights
+        self.phi = opt.phi
+        self.confidence = opt.confidence
+        self.nms_iou = opt.nms_iou
+        if self.tiny:
+            self.anchors_mask = [[3,4,5], [1,2,3]]
+            self.anchors_path = 'model_data/yolotiny_anchors.txt'
+        #---------------------------------------------------#
+        #   获得种类和先验框的数量
+        #---------------------------------------------------#
+        # self.class_names, self.num_classes  = get_classes(self.classes_path)
+        self.anchors, self.num_anchors      = get_anchors(self.anchors_path)
+        self.bbox_util                      = DecodeBox(self.anchors, self.num_classes, (self.input_shape[0], self.input_shape[1]), self.anchors_mask)
+
+        #---------------------------------------------------#
+        #   画框设置不同的颜色
+        #---------------------------------------------------#
+        hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
+        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
+        self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
+        self.generate()
+
+    #---------------------------------------------------#
+    #   生成模型
+    #---------------------------------------------------#
+    def generate(self, onnx=False):
+        #---------------------------------------------------#
+        #   建立yolo模型，载入yolo模型的权重
+        #---------------------------------------------------#
+
+        if not self.tiny:
+            self.net    = YoloBody(self.anchors_mask, self.num_classes)
+        elif self.tiny:
+            self.net    = YoloBodytiny(self.anchors_mask, self.num_classes, self.phi)
+        device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.net.load_state_dict(torch.load(self.model_path, map_location=device))
+        self.net    = self.net.eval()
+
+        print('{} model, anchors, and classes loaded.'.format(self.model_path))
+        if not onnx:
+            if self.cuda:
+                self.net = nn.DataParallel(self.net)
+                self.net = self.net.cuda()
+
+    #---------------------------------------------------#
+    #   检测图片
+    #---------------------------------------------------#
+    def detect_image(self, image, crop = False, count = False):
+        #---------------------------------------------------#
+        #   计算输入图片的高和宽
+        #---------------------------------------------------#
+        image_shape = np.array(np.shape(image)[0:2])
+        #---------------------------------------------------------#
+        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
+        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+        #---------------------------------------------------------#
+        image       = cvtColor(image)
+        #---------------------------------------------------------#
+        #   给图像增加灰条，实现不失真的resize
+        #   也可以直接resize进行识别
+        #---------------------------------------------------------#
+        image_data  = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
+        #---------------------------------------------------------#
+        #   添加上batch_size维度
+        #---------------------------------------------------------#
+        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
+
+        with torch.no_grad():
+            images = torch.from_numpy(image_data)
+            if self.cuda:
+                images = images.cuda()
+            #---------------------------------------------------------#
+            #   将图像输入网络当中进行预测！
+            #---------------------------------------------------------#
+            outputs = self.net(images)
+            outputs = self.bbox_util.decode_box(outputs)
+            #---------------------------------------------------------#
+            #   将预测框进行堆叠，然后进行非极大抑制
+            #---------------------------------------------------------#
+            results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, 
+                        image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
+                                                    
+            if results[0] is None: 
+                return image
+
+            top_label   = np.array(results[0][:, 6], dtype = 'int32')
+            top_conf    = results[0][:, 4] * results[0][:, 5]
+            top_boxes   = results[0][:, :4]
+        #---------------------------------------------------------#
+        #   设置字体与边框厚度
+        #---------------------------------------------------------#
+        font        = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
+        thickness   = int(max((image.size[0] + image.size[1]) // np.mean(self.input_shape), 1))
+        #---------------------------------------------------------#
+        #   计数
+        #---------------------------------------------------------#
+        if count:
+            print("top_label:", top_label)
+            classes_nums    = np.zeros([self.num_classes])
+            for i in range(self.num_classes):
+                num = np.sum(top_label == i)
+                if num > 0:
+                    print(self.class_names[i], " : ", num)
+                classes_nums[i] = num
+            print("classes_nums:", classes_nums)
+        #---------------------------------------------------------#
+        #   是否进行目标的裁剪
+        #---------------------------------------------------------#
+        if crop:
+            for i, c in list(enumerate(top_label)):
+                top, left, bottom, right = top_boxes[i]
+                top     = max(0, np.floor(top).astype('int32'))
+                left    = max(0, np.floor(left).astype('int32'))
+                bottom  = min(image.size[1], np.floor(bottom).astype('int32'))
+                right   = min(image.size[0], np.floor(right).astype('int32'))
+                
+                dir_save_path = "img_crop"
+                if not os.path.exists(dir_save_path):
+                    os.makedirs(dir_save_path)
+                crop_image = image.crop([left, top, right, bottom])
+                crop_image.save(os.path.join(dir_save_path, "crop_" + str(i) + ".png"), quality=95, subsampling=0)
+                print("save crop_" + str(i) + ".png to " + dir_save_path)
+        #---------------------------------------------------------#
+        #   图像绘制
+        #---------------------------------------------------------#
+        for i, c in list(enumerate(top_label)):
+            predicted_class = self.class_names[int(c)]
+            box             = top_boxes[i]
+            score           = top_conf[i]
+
+            top, left, bottom, right = box
+
+            top     = max(0, np.floor(top).astype('int32'))
+            left    = max(0, np.floor(left).astype('int32'))
+            bottom  = min(image.size[1], np.floor(bottom).astype('int32'))
+            right   = min(image.size[0], np.floor(right).astype('int32'))
+
+            label = '{} {:.2f}'.format(predicted_class, score)
+            draw = ImageDraw.Draw(image)
+            label_size = draw.textsize(label, font)
+            label = label.encode('utf-8')
+            print(label, top, left, bottom, right)
+            
+            if top - label_size[1] >= 0:
+                text_origin = np.array([left, top - label_size[1]])
+            else:
+                text_origin = np.array([left, top + 1])
+
+            for i in range(thickness):
+                draw.rectangle([left + i, top + i, right - i, bottom - i], outline=self.colors[c])
+            draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=self.colors[c])
+            draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
+            del draw
+
+        return image
+
+    def get_FPS(self, image, test_interval):
+        image_shape = np.array(np.shape(image)[0:2])
+        #---------------------------------------------------------#
+        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
+        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+        #---------------------------------------------------------#
+        image       = cvtColor(image)
+        #---------------------------------------------------------#
+        #   给图像增加灰条，实现不失真的resize
+        #   也可以直接resize进行识别
+        #---------------------------------------------------------#
+        image_data  = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
+        #---------------------------------------------------------#
+        #   添加上batch_size维度
+        #---------------------------------------------------------#
+        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
+
+        with torch.no_grad():
+            images = torch.from_numpy(image_data)
+            if self.cuda:
+                images = images.cuda()
+            #---------------------------------------------------------#
+            #   将图像输入网络当中进行预测！
+            #---------------------------------------------------------#
+            outputs = self.net(images)
+            outputs = self.bbox_util.decode_box(outputs)
+            #---------------------------------------------------------#
+            #   将预测框进行堆叠，然后进行非极大抑制
+            #---------------------------------------------------------#
+            results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, 
+                        image_shape, self.letterbox_image, conf_thres=self.confidence, nms_thres=self.nms_iou)
+                                                    
+        t1 = time.time()
+        for _ in range(test_interval):
+            with torch.no_grad():
+                #---------------------------------------------------------#
+                #   将图像输入网络当中进行预测！
+                #---------------------------------------------------------#
+                outputs = self.net(images)
+                outputs = self.bbox_util.decode_box(outputs)
+                #---------------------------------------------------------#
+                #   将预测框进行堆叠，然后进行非极大抑制
+                #---------------------------------------------------------#
+                results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, 
+                            image_shape, self.letterbox_image, conf_thres=self.confidence, nms_thres=self.nms_iou)
+                            
+        t2 = time.time()
+        tact_time = (t2 - t1) / test_interval
+        return tact_time
+
+    def detect_heatmap(self, image, heatmap_save_path):
+        import cv2
+        import matplotlib.pyplot as plt
+        def sigmoid(x):
+            y = 1.0 / (1.0 + np.exp(-x))
+            return y
+        #---------------------------------------------------------#
+        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
+        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+        #---------------------------------------------------------#
+        image       = cvtColor(image)
+        #---------------------------------------------------------#
+        #   给图像增加灰条，实现不失真的resize
+        #   也可以直接resize进行识别
+        #---------------------------------------------------------#
+        image_data  = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
+        #---------------------------------------------------------#
+        #   添加上batch_size维度
+        #---------------------------------------------------------#
+        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
+
+        with torch.no_grad():
+            images = torch.from_numpy(image_data)
+            if self.cuda:
+                images = images.cuda()
+            #---------------------------------------------------------#
+            #   将图像输入网络当中进行预测！
+            #---------------------------------------------------------#
+            outputs = self.net(images)
+        
+        plt.imshow(image, alpha=1)
+        plt.axis('off')
+        mask    = np.zeros((image.size[1], image.size[0]))
+        for sub_output in outputs:
+            sub_output = sub_output.cpu().numpy()
+            b, c, h, w = np.shape(sub_output)
+            sub_output = np.transpose(np.reshape(sub_output, [b, 3, -1, h, w]), [0, 3, 4, 1, 2])[0]
+            score      = np.max(sigmoid(sub_output[..., 4]), -1)
+            score      = cv2.resize(score, (image.size[0], image.size[1]))
+            normed_score    = (score * 255).astype('uint8')
+            mask            = np.maximum(mask, normed_score)
+            
+        plt.imshow(mask, alpha=0.5, interpolation='nearest', cmap="jet")
+
+        plt.axis('off')
+        plt.subplots_adjust(top=1, bottom=0, right=1,  left=0, hspace=0, wspace=0)
+        plt.margins(0, 0)
+        plt.savefig(heatmap_save_path, dpi=200, bbox_inches='tight', pad_inches = -0.1)
+        print("Save to the " + heatmap_save_path)
+        plt.show()
+
+    def convert_to_onnx(self, simplify, model_path):
+        import onnx
+        self.generate(onnx=True)
+
+        im                  = torch.zeros(1, 3, *self.input_shape).to('cpu')  # image size(1, 3, 512, 512) BCHW
+        input_layer_names   = ["images"]
+        output_layer_names  = ["output"]
+        
+        # Export the model
+        print(f'Starting export with onnx {onnx.__version__}.')
+        torch.onnx.export(self.net,
+                        im,
+                        f               = model_path,
+                        verbose         = False,
+                        opset_version   = 12,
+                        training        = torch.onnx.TrainingMode.EVAL,
+                        do_constant_folding = True,
+                        input_names     = input_layer_names,
+                        output_names    = output_layer_names,
+                        dynamic_axes    = None)
+
+        # Checks
+        model_onnx = onnx.load(model_path)  # load onnx model
+        onnx.checker.check_model(model_onnx)  # check onnx model
+
+        # Simplify onnx
+        if simplify:
+            import onnxsim
+            print(f'Simplifying with onnx-simplifier {onnxsim.__version__}.')
+            model_onnx, check = onnxsim.simplify(
+                model_onnx,
+                dynamic_input_shape=False,
+                input_shapes=None)
+            assert check, 'assert check failed'
+            onnx.save(model_onnx, model_path)
+
+        print('Onnx model save as {}'.format(model_path))
+
+    def get_map_txt(self, image_id, image, class_names, map_out_path):
+        f = open(os.path.join(map_out_path, "detection-results/"+image_id+".txt"),"w") 
+        image_shape = np.array(np.shape(image)[0:2])
+        #---------------------------------------------------------#
+        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
+        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+        #---------------------------------------------------------#
+        image       = cvtColor(image)
+        #---------------------------------------------------------#
+        #   给图像增加灰条，实现不失真的resize
+        #   也可以直接resize进行识别
+        #---------------------------------------------------------#
+        image_data  = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
+        #---------------------------------------------------------#
+        #   添加上batch_size维度
+        #---------------------------------------------------------#
+        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
+
+        with torch.no_grad():
+            images = torch.from_numpy(image_data)
+            if self.cuda:
+                images = images.cuda()
+            #---------------------------------------------------------#
+            #   将图像输入网络当中进行预测！
+            #---------------------------------------------------------#
+            outputs = self.net(images)
+            outputs = self.bbox_util.decode_box(outputs)
+            #---------------------------------------------------------#
+            #   将预测框进行堆叠，然后进行非极大抑制
+            #---------------------------------------------------------#
+            results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape, 
+                        image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
+                                                    
+            if results[0] is None: 
+                return 
+
+            top_label   = np.array(results[0][:, 6], dtype = 'int32')
+            top_conf    = results[0][:, 4] * results[0][:, 5]
+            top_boxes   = results[0][:, :4]
+
+        for i, c in list(enumerate(top_label)):
+            predicted_class = self.class_names[int(c)]
+            box             = top_boxes[i]
+            score           = str(top_conf[i])
+
+            top, left, bottom, right = box
+            if predicted_class not in class_names:
+                continue
+
+            f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))
+
+        f.close()
+        return