yolov3、yolov4等批量预测图片并保存(python)

说明

本批量预测图片代码很适用于深度学习目标检测当中，将代码主干添加到预测脚本当中，修改部分调用参数、读入文件夹、保存文件夹等都可适用，方法如下：

方法

一、按图片的文件夹获取图片路径

代码展示：predict.py

#-----------------------------------------------------------------------#
#   predict.py将单张图片预测、摄像头检测、FPS测试和目录遍历检测等功能
#   整合到了一个py文件中，通过指定mode进行模式的修改。
#-----------------------------------------------------------------------#
import time
import os
import cv2
import numpy as np
from PIL import Image
from io import BytesIO
from yolo import YOLO
import shutil
#from six.moves import urllib

if __name__ == "__main__":
    yolo = YOLO()
    #----------------------------------------------------------------------------------------------------------#
    #   mode用于指定测试的模式：
    #   'predict'表示单张图片预测，如果想对预测过程进行修改，如保存图片，截取对象等，可以先看下方详细的注释
    #   'video'表示视频检测，可调用摄像头或者视频进行检测，详情查看下方注释。
    #   'fps'表示测试fps，使用的图片是img里面的street.jpg，详情查看下方注释。
    #   'dir_predict'表示遍历文件夹进行检测并保存。默认遍历img文件夹，保存img_out文件夹，详情查看下方注释。
    #----------------------------------------------------------------------------------------------------------#
    mode = "dir_predict"
    #----------------------------------------------------------------------------------------------------------#
    #   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
    video_save_path = ""
    video_fps       = 25.0
    #-------------------------------------------------------------------------#
    #   test_interval用于指定测量fps的时候，图片检测的次数
    #   理论上test_interval越大，fps越准确。
    #-------------------------------------------------------------------------#
    test_interval   = 100
    #-------------------------------------------------------------------------#
    #   dir_origin_path指定了用于检测的图片的文件夹路径
    #   dir_save_path指定了检测完图片的保存路径
    #   dir_origin_path和dir_save_path仅在mode='dir_predict'时有效
    #-------------------------------------------------------------------------#
    #dir_origin_path = "img/"
    dir_txt_path ="G:/Deep Code/yolov4-pytorch-master/map_out/detection-results/"
    dir_origin_path = "C:/Users/lenovo/Desktop/VOCdevkit/VOC2012/JPEGImages/"
    dir_save_path   = "img_out/"
    if not os.path.exists(dir_save_path):
                    os.makedirs(dir_save_path)
    s=os.listdir(dir_origin_path)
    #----------------------------------------------------#
    #----------------批量预测-----------------------------#
    #----------------------------------------------------#
    if mode == "dir_predict":
        
        #----------------------------------------------------#
        #----------------先清空文件夹中的内容------------------#
        #----------------------------------------------------#
        #shutil.rmtree(dir_crop_save_path)
        #os.mkdir(dir_crop_save_path)
        shutil.rmtree(dir_save_path)
        os.mkdir(dir_save_path)
        while True:
            #img = input('Input image filename:')
        #with open(dir_txt_path)as f:
            for filename in os.listdir(r"C:/Users/lenovo/Desktop/VOCdevkit/VOC2012/JPEGImages/"): 
               print ( filename)
               filename=filename.strip('\n')
               image_path  = os.path.join('C:/Users/lenovo/Desktop/VOCdevkit/VOC2012/JPEGImages/',filename)
               image= Image.open(image_path)
               r_image = yolo.detect_image(image)
               
               #r_image.save(dir_save_path+os.sep+'%s.jpg')
               r_image.save(dir_save_path+filename)
            """
            for i in s:
                count=1
                document = os.path.join(dir_origin_path,i)
                #f = urllib.request.urlopen( dir_origin_path)
                #img = Image.open(document)
                image = Image.open(document)
                img=image.read()
                r_image = yolo.detect_image(img)
                #r_image.show()