AMXZzzz/图片视频推理识别.py

## 图片视频推理识别.py
# QQ2716490456
# 随风而息
# ————————————————
import os
import win32api
import win32gui
import win32con
import mss
import cv2
import numpy as np
from utils.general import  non_max_suppression,scale_coords, xyxy2xywh
from utils.datasets import letterbox
import argparse
import torch
from models.experimental import attempt_load
import keyboard


'''
注意事项：选择识别模式后，需要修改路径的话(单个视频的路径修改就去单个视频识别，多个视频就去多个视频识别，屏幕识别就去屏幕识别)。
        路径前的r是屏蔽转义符，
'''
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default=r'C:\Users\Banditek\Desktop\yolov5-6.0\weights\CF-600img.pt', help='模型地址,请使用绝对路径')
parser.add_argument('--imgsz', type=int, default=640, help='和你训练模型时imgsz一样')
'''
        ！！！！！！！！！！！！！！ 模式选择(必选)！！！！！！！！！！！！！
'''
parser.add_argument('--choose', type=int, default=0, help='模式选择:  0：单个视频识别，  1：多个视频识别，  2：屏幕识别')
'''
0：单个视频识别
'''
parser.add_argument('--video_path', type=str, default=r'D:\obs\video\297510698-1-80.flv', help='输入路径：单个<视频文件>的绝对路径')
parser.add_argument('--output_path', type=str, default=r'D:/obs/xml', help='输出路径单个视频文件的图片输出路径')
'''
1：多个视频识别
'''
parser.add_argument('--videos_path', type=str, default=r'D:\obs\xml', help='输入路径：多个视频<文件夹>的绝对路径')
parser.add_argument('--outputs_path', type=str, default=r'D:/obs/img', help='输出路径：多个视频文件的图片输出路径')
'''
2：屏幕识别
'''
parser.add_argument('--region', type=tuple, default=(1.0, 1.0), help='屏幕检测范围；分别为横向和竖向，(1.0, 1.0)表示全屏检测，越低检测范围越小(始终保持屏幕中心为中心)')
parser.add_argument('--show-window', type=bool, default=True, help='是否显示实时检测窗口(新版里改进了效率。若为True，不要去点右上角的X！)')
parser.add_argument('--resize-window', type=float, default=1/3, help='缩放实时检测窗口大小')
parser.add_argument('--top-most', type=bool, default=False, help='是否保持实时检测窗口置顶，')
parser.add_argument('--mss_output_path', type=str, default=r'D:/obs/img', help='输出路径屏幕识别的图片输出路径')
args = parser.parse_args()

'''mss截图'''
sct = mss.mss()
def grab_screen_mss(monitor):
    try:
        return cv2.cvtColor(np.array(sct.grab(monitor)), cv2.COLOR_BGRA2BGR)
    except SyntaxError as e:
        print('错误1')

def get_parameters():
    try:
        x, y = get_screen_size().values()
        return 0, 0, x, y
    except SyntaxError as e:
        print('错误2')
def get_screen_size():
    try:
        wide = win32api.GetSystemMetrics(0)
        high = win32api.GetSystemMetrics(1)
        return {"wide": wide, "high": high}
    except SyntaxError as e:
        print('错误3')

'''model'''

def load_model(args):
    device = torch.device('cuda')  # 定义device 选择用gpu还是cpu
    #weights = r'C:\Users\Banditek\Desktop\yolov5-6.0\weights\CF-600img.pt'  # 训练好的模型地址
    imgsz = 640  # 训练时的分辨率是多少，这里就填多少
    model = attempt_load(args.weights, map_location=device)      #导入模型
    model.half()  # to FP16
    model(torch.zeros(1,3,imgsz,imgsz).to(device).type_as(next(model.parameters()))) #模型在cuda上初始化
    return model

''''''

def wangluo(image,aims):
    '''进行推理'''
    # 1.6ms
    img = letterbox(image, img_size, stride=stride)[0]  # 将截图image传入处理程序
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR 转 RGB, to 3x416x416    对图片进行格式转换
    img = np.ascontiguousarray(img)

    # 0.4ms
    img = torch.from_numpy(img).to(device)  # 4跟随推理程序的64行
    img = img.half() if half else img.float()  # 4
    img /= 255.0  # 4跟随66，对3个rgb（0.256）进行规划，对图片进行归化处理
    if img.ndimension() == 3:  # 4
        img = img[None]  # 4与img = img.unsqueeze(0)一模一样

        # 截图和对图片进行归化用了50-70ms左右
    '''----调用yolov5推理----'''
    # 50ms左右

    pred = model(img, augment='store_true')[0]  # 初始化数据
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False)  # 4跟随75，传入参数进行推理
    # tis2 = time.perf_counter()
    # print(tis2 - tis1)
    # print(pred)                                                     #输出检测到的图片数据

    '''----对推理的结果进行解密----'''
    # 2ms
    # aims = []  # 定义一个变量列表，用于循环清空
    for i, det in enumerate(pred):  # 此for式对推理出来的图片进行解密
        gn = torch.tensor(image.shape)[[1, 0, 1, 0]]
        if len(det):
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], image.shape).round()
            for *xyxy, conf, cls in reversed(det):
                # bbox;(tag,x_center,y_center,x_width,y_width)
                '''
                0 hat
                '''
                xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                line = (cls, *xywh)  # label format
                aim = ('%g ' * len(line)).rstrip() % line
                aim = aim.split(' ')  # 曾经重大bug，split书写错误
                #print(aim)
                aims.append(aim)  # 将标签数据写入aims列表
    return aims

'''视频识别'''
def video2image(video_input,output_path):

    count = 0  # 计数用，分割的图片按照count来命名
    # 提取视频图片
    times = 0  # 用来记录帧
    frame_frequency = 10  # 提取视频的频率，每frameFrequency帧提取一张图片，提取完整视频帧设置为1

    cap = cv2.VideoCapture(video_input)  # 读取视频文件
    frameall = int(cap.get(7))  # 获取总帧数
    print('本次识别视频一共有',frameall,'帧')
    print('开始提取', video_input, '视频的图片')

    while True:
        aims = []
        times += 1
        res, image = cap.read()  # 读出图片。res表示是否读取到图片，image表示读取到的每一帧图片，每调用一次读取下一帧数
        if times >= frameall:
            print('————视频图片已提取结束')
            break
        '''进行推理'''
        wangluo(image,aims)
        '''推理结束，判断是否有识别出的数据'''
        if len(aims):
            # if times % frame_frequency == 0:
            img_name = str(count).zfill(6) + '.jpg'  # 图片计数，6位上限
            cv2.imwrite(output_path + os.sep + img_name, image)  # 存储图片
            count += 1
            # 输出提示
            if times % 100 == 0:
                #print(output_path + os.sep + img_name)
                print('视频提取进度:',round((times/frameall)*100,1),'%','\t','————{ 在运行框中，按< L >跳过本次识别','如需停止请杀死程序 }————')
        if keyboard.is_pressed('l'):
            print('————已跳过本次识别...')
            break
    cap.release()

def fun_mss(mss_output_path):
    '''循环函数'''
    count = 0  # 计数用，分割的图片按照count来命名

    while True:
        aims = []
        img0 = grab_screen_mss(monitor)
        img0 = cv2.resize(img0, (len_x, len_y))
        wangluo(img0,aims)
        if len(aims):
            img_name = str(count).zfill(6) + '.jpg'  # 图片计数，6位上限
            cv2.imwrite(mss_output_path + os.sep + img_name, img0)  # 存储图片
            count += 1
            print('已识别图片{} 并存入文件夹'.format(img_name),'————{ 在运行框中，按< L >退出屏幕识别 }————')
            if args.show_window:
                for i, det in enumerate(aims):
                    tag, x_center, y_center, width, height = det
                    x_center, width = len_x * float(x_center), len_x * float(width)
                    y_center, height = len_y * float(y_center), len_y * float(height)
                    top_left = (int(x_center - width / 2.), int(y_center - height / 2.))
                    bottom_right = (int(x_center + width / 2.), int(y_center + height / 2.))
                    cv2.rectangle(img0, top_left, bottom_right, (0, 255, 0), thickness=3)
        if args.show_window:
            cv2.namedWindow('csgo-detect', cv2.WINDOW_NORMAL)
            cv2.resizeWindow('csgo-detect', int(len_x * args.resize_window), int(len_y * args.resize_window))
            cv2.imshow('csgo-detect', img0)
            if args.top_most:
                hwnd = win32gui.FindWindow(None, 'csgo-detect')
                CVRECT = cv2.getWindowImageRect('csgo-detect')
                win32gui.SetWindowPos(hwnd, win32con.HWND_TOPMOST, 0, 0, 0, 0,
                                      win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)
        if cv2.waitKey(1) & 0xff == ord('l'):  # 2如果检测到按到‘l’，就退出此窗口，l可以自定义
            cv2.destroyAllWindows()
            break

top_x, top_y, x, y = get_parameters()  # x， y截取显示器的分辨率大小
len_x, len_y = int(x * args.region[0]), int(y * args.region[1])  # 原生分辨率*检测比例，横向对比#截图范围的左上角坐标
top_x, top_y = int(top_x + x // 2 * (1. - args.region[0])), int(
    top_y + y // 2 * (1. - args.region[1]))  ##截图范围的右下角坐标
monitor = {'left': top_x, 'top': top_y, 'width': len_x,
           'height': len_y}  # 用mss截取检测图片的分辨率大小，横向对比top_x = 0,  top_y = 0，len_x = 1920,  len_y = 1080

'''变量定义'''
img_size = 640
model = load_model(args)
stride = int(model.stride.max())
device = 'cuda' if torch.cuda.is_available() else ''  # 4定义device
half = device != 'cpu'
conf_thres = 0.25
iou_thres = 0.05
imgsz = args.imgsz

if __name__ == '__main__':

    '''单个视频文件逐帧推理截图'''
    if args.choose == 0:
        # 视频路径   # 图片输出路径
        video_input = args.video_path                     #r'D:\obs\apex\2022-01-11 21-25-45.mp4'
        output_path = args.output_path                                     #'D:/obs/img'
        # 输出文件夹不存在，则创建输出文件夹
        if not os.path.exists(output_path):         #不存在则创建输出文件夹
            os.makedirs(output_path)
        print('模式 0:———— 单个视频提取')
        video2image(video_input,output_path)        #run
        print('视频已提取完成，程序退出...')

        '''多个视频文件逐帧推理截图'''
    elif args.choose == 1 :
        #多个视频的文件夹
        videos_path = args.videos_path
        outputs_path = args.outputs_path
        lst1 = os.walk(videos_path)         #查询该文件夹下的所有文件
        if not os.path.exists(outputs_path):
            os.makedirs(outputs_path)
        print('模式 1:———— 多个视频提取')
        for i in lst1:              #遍历lst1所查询到的文件，i得到的是元组类型，有3个数据
            lst = i[2]              #获取元组中的索引为2的数据 0:是主目录 ，1：是子目录，若没有主目录则返回空列表， 2：是所有目录下的文件
            print('-' * 100)
            print('此文件夹搜索到以下视频文件：',lst)
            print('-' * 100)
            for file in lst:    #遍历lst的文件信息
                videos_input = os.path.join(videos_path, file)  # 将主目录和文件拼接，若有子目录则在videos_path后面添加子目录信息
                video2image(videos_input, outputs_path)         #run
        print('视频已提取完成，程序退出...')

        '''mss屏幕识别'''
    elif args.choose == 2:
        mss_output_path = args.mss_output_path          #输出图片目录
        print('模式 2：———— 屏幕识别，按< L >键退出识别')
        fun_mss(mss_output_path)
        print('视频已提取完成，程序退出...')
	# QQ2716490456
	# 随风而息
	# ————————————————
	import os
	import win32api
	import win32gui
	import win32con
	import mss
	import cv2
	import numpy as np
	from utils.general import non_max_suppression,scale_coords, xyxy2xywh
	from utils.datasets import letterbox
	import argparse
	import torch
	from models.experimental import attempt_load
	import keyboard


	'''
	注意事项：选择识别模式后，需要修改路径的话(单个视频的路径修改就去单个视频识别，多个视频就去多个视频识别，屏幕识别就去屏幕识别)。
	路径前的r是屏蔽转义符，
	'''
	parser = argparse.ArgumentParser()
	parser.add_argument('--weights', type=str, default=r'C:\Users\Banditek\Desktop\yolov5-6.0\weights\CF-600img.pt', help='模型地址,请使用绝对路径')
	parser.add_argument('--imgsz', type=int, default=640, help='和你训练模型时imgsz一样')
	'''
	！！！！！！！！！！！！！！模式选择(必选)！！！！！！！！！！！！！
	'''
	parser.add_argument('--choose', type=int, default=0, help='模式选择: 0：单个视频识别， 1：多个视频识别， 2：屏幕识别')
	'''
	0：单个视频识别
	'''
	parser.add_argument('--video_path', type=str, default=r'D:\obs\video\297510698-1-80.flv', help='输入路径：单个<视频文件>的绝对路径')
	parser.add_argument('--output_path', type=str, default=r'D:/obs/xml', help='输出路径单个视频文件的图片输出路径')
	'''
	1：多个视频识别
	'''
	parser.add_argument('--videos_path', type=str, default=r'D:\obs\xml', help='输入路径：多个视频<文件夹>的绝对路径')
	parser.add_argument('--outputs_path', type=str, default=r'D:/obs/img', help='输出路径：多个视频文件的图片输出路径')
	'''
	2：屏幕识别
	'''
	parser.add_argument('--region', type=tuple, default=(1.0, 1.0), help='屏幕检测范围；分别为横向和竖向，(1.0, 1.0)表示全屏检测，越低检测范围越小(始终保持屏幕中心为中心)')
	parser.add_argument('--show-window', type=bool, default=True, help='是否显示实时检测窗口(新版里改进了效率。若为True，不要去点右上角的X！)')
	parser.add_argument('--resize-window', type=float, default=1/3, help='缩放实时检测窗口大小')
	parser.add_argument('--top-most', type=bool, default=False, help='是否保持实时检测窗口置顶，')
	parser.add_argument('--mss_output_path', type=str, default=r'D:/obs/img', help='输出路径屏幕识别的图片输出路径')
	args = parser.parse_args()

	'''mss截图'''
	sct = mss.mss()
	def grab_screen_mss(monitor):
	try:
	return cv2.cvtColor(np.array(sct.grab(monitor)), cv2.COLOR_BGRA2BGR)
	except SyntaxError as e:
	print('错误1')

	def get_parameters():
	try:
	x, y = get_screen_size().values()
	return 0, 0, x, y
	except SyntaxError as e:
	print('错误2')
	def get_screen_size():
	try:
	wide = win32api.GetSystemMetrics(0)
	high = win32api.GetSystemMetrics(1)
	return {"wide": wide, "high": high}
	except SyntaxError as e:
	print('错误3')

	'''model'''

	def load_model(args):
	device = torch.device('cuda') # 定义device 选择用gpu还是cpu
	#weights = r'C:\Users\Banditek\Desktop\yolov5-6.0\weights\CF-600img.pt' # 训练好的模型地址
	imgsz = 640 # 训练时的分辨率是多少，这里就填多少
	model = attempt_load(args.weights, map_location=device) #导入模型
	model.half() # to FP16
	model(torch.zeros(1,3,imgsz,imgsz).to(device).type_as(next(model.parameters()))) #模型在cuda上初始化
	return model

	''''''

	def wangluo(image,aims):
	'''进行推理'''
	# 1.6ms
	img = letterbox(image, img_size, stride=stride)[0] # 将截图image传入处理程序
	img = img[:, :, ::-1].transpose(2, 0, 1) # BGR 转 RGB, to 3x416x416 对图片进行格式转换
	img = np.ascontiguousarray(img)

	# 0.4ms
	img = torch.from_numpy(img).to(device) # 4跟随推理程序的64行
	img = img.half() if half else img.float() # 4
	img /= 255.0 # 4跟随66，对3个rgb（0.256）进行规划，对图片进行归化处理
	if img.ndimension() == 3: # 4
	img = img[None] # 4与img = img.unsqueeze(0)一模一样

	# 截图和对图片进行归化用了50-70ms左右
	'''----调用yolov5推理----'''
	# 50ms左右

	pred = model(img, augment='store_true')[0] # 初始化数据
	pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False) # 4跟随75，传入参数进行推理
	# tis2 = time.perf_counter()
	# print(tis2 - tis1)
	# print(pred) #输出检测到的图片数据

	'''----对推理的结果进行解密----'''
	# 2ms
	# aims = [] # 定义一个变量列表，用于循环清空
	for i, det in enumerate(pred): # 此for式对推理出来的图片进行解密
	gn = torch.tensor(image.shape)[[1, 0, 1, 0]]
	if len(det):
	det[:, :4] = scale_coords(img.shape[2:], det[:, :4], image.shape).round()
	for *xyxy, conf, cls in reversed(det):
	# bbox;(tag,x_center,y_center,x_width,y_width)
	'''
	0 hat
	'''
	xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
	line = (cls, *xywh) # label format
	aim = ('%g ' * len(line)).rstrip() % line
	aim = aim.split(' ') # 曾经重大bug，split书写错误
	#print(aim)
	aims.append(aim) # 将标签数据写入aims列表
	return aims

	'''视频识别'''
	def video2image(video_input,output_path):

	count = 0 # 计数用，分割的图片按照count来命名
	# 提取视频图片
	times = 0 # 用来记录帧
	frame_frequency = 10 # 提取视频的频率，每frameFrequency帧提取一张图片，提取完整视频帧设置为1

	cap = cv2.VideoCapture(video_input) # 读取视频文件
	frameall = int(cap.get(7)) # 获取总帧数
	print('本次识别视频一共有',frameall,'帧')
	print('开始提取', video_input, '视频的图片')

	while True:
	aims = []
	times += 1
	res, image = cap.read() # 读出图片。res表示是否读取到图片，image表示读取到的每一帧图片，每调用一次读取下一帧数
	if times >= frameall:
	print('————视频图片已提取结束')
	break
	'''进行推理'''
	wangluo(image,aims)
	'''推理结束，判断是否有识别出的数据'''
	if len(aims):
	# if times % frame_frequency == 0:
	img_name = str(count).zfill(6) + '.jpg' # 图片计数，6位上限
	cv2.imwrite(output_path + os.sep + img_name, image) # 存储图片
	count += 1
	# 输出提示
	if times % 100 == 0:
	#print(output_path + os.sep + img_name)
	print('视频提取进度:',round((times/frameall)*100,1),'%','\t','————{ 在运行框中，按< L >跳过本次识别','如需停止请杀死程序 }————')
	if keyboard.is_pressed('l'):
	print('————已跳过本次识别...')
	break
	cap.release()

	def fun_mss(mss_output_path):
	'''循环函数'''
	count = 0 # 计数用，分割的图片按照count来命名

	while True:
	aims = []
	img0 = grab_screen_mss(monitor)
	img0 = cv2.resize(img0, (len_x, len_y))
	wangluo(img0,aims)
	if len(aims):
	img_name = str(count).zfill(6) + '.jpg' # 图片计数，6位上限
	cv2.imwrite(mss_output_path + os.sep + img_name, img0) # 存储图片
	count += 1
	print('已识别图片{} 并存入文件夹'.format(img_name),'————{ 在运行框中，按< L >退出屏幕识别 }————')
	if args.show_window:
	for i, det in enumerate(aims):
	tag, x_center, y_center, width, height = det
	x_center, width = len_x * float(x_center), len_x * float(width)
	y_center, height = len_y * float(y_center), len_y * float(height)
	top_left = (int(x_center - width / 2.), int(y_center - height / 2.))
	bottom_right = (int(x_center + width / 2.), int(y_center + height / 2.))
	cv2.rectangle(img0, top_left, bottom_right, (0, 255, 0), thickness=3)
	if args.show_window:
	cv2.namedWindow('csgo-detect', cv2.WINDOW_NORMAL)
	cv2.resizeWindow('csgo-detect', int(len_x * args.resize_window), int(len_y * args.resize_window))
	cv2.imshow('csgo-detect', img0)
	if args.top_most:
	hwnd = win32gui.FindWindow(None, 'csgo-detect')
	CVRECT = cv2.getWindowImageRect('csgo-detect')
	win32gui.SetWindowPos(hwnd, win32con.HWND_TOPMOST, 0, 0, 0, 0,
	win32con.SWP_NOMOVE \| win32con.SWP_NOSIZE)
	if cv2.waitKey(1) & 0xff == ord('l'): # 2如果检测到按到‘l’，就退出此窗口，l可以自定义
	cv2.destroyAllWindows()
	break

	top_x, top_y, x, y = get_parameters() # x， y截取显示器的分辨率大小
	len_x, len_y = int(x * args.region[0]), int(y * args.region[1]) # 原生分辨率*检测比例，横向对比#截图范围的左上角坐标
	top_x, top_y = int(top_x + x // 2 * (1. - args.region[0])), int(
	top_y + y // 2 * (1. - args.region[1])) ##截图范围的右下角坐标
	monitor = {'left': top_x, 'top': top_y, 'width': len_x,
	'height': len_y} # 用mss截取检测图片的分辨率大小，横向对比top_x = 0, top_y = 0，len_x = 1920, len_y = 1080

	'''变量定义'''
	img_size = 640
	model = load_model(args)
	stride = int(model.stride.max())
	device = 'cuda' if torch.cuda.is_available() else '' # 4定义device
	half = device != 'cpu'
	conf_thres = 0.25
	iou_thres = 0.05
	imgsz = args.imgsz

	if __name__ == '__main__':

	'''单个视频文件逐帧推理截图'''
	if args.choose == 0:
	# 视频路径 # 图片输出路径
	video_input = args.video_path #r'D:\obs\apex\2022-01-11 21-25-45.mp4'
	output_path = args.output_path #'D:/obs/img'
	# 输出文件夹不存在，则创建输出文件夹
	if not os.path.exists(output_path): #不存在则创建输出文件夹
	os.makedirs(output_path)
	print('模式 0:———— 单个视频提取')
	video2image(video_input,output_path) #run
	print('视频已提取完成，程序退出...')

	'''多个视频文件逐帧推理截图'''
	elif args.choose == 1 :
	#多个视频的文件夹
	videos_path = args.videos_path
	outputs_path = args.outputs_path
	lst1 = os.walk(videos_path) #查询该文件夹下的所有文件
	if not os.path.exists(outputs_path):
	os.makedirs(outputs_path)
	print('模式 1:———— 多个视频提取')
	for i in lst1: #遍历lst1所查询到的文件，i得到的是元组类型，有3个数据
	lst = i[2] #获取元组中的索引为2的数据 0:是主目录，1：是子目录，若没有主目录则返回空列表， 2：是所有目录下的文件
	print('-' * 100)
	print('此文件夹搜索到以下视频文件：',lst)
	print('-' * 100)
	for file in lst: #遍历lst的文件信息
	videos_input = os.path.join(videos_path, file) # 将主目录和文件拼接，若有子目录则在videos_path后面添加子目录信息
	video2image(videos_input, outputs_path) #run
	print('视频已提取完成，程序退出...')

	'''mss屏幕识别'''
	elif args.choose == 2:
	mss_output_path = args.mss_output_path #输出图片目录
	print('模式 2：———— 屏幕识别，按< L >键退出识别')
	fun_mss(mss_output_path)
	print('视频已提取完成，程序退出...')