利用yoluv5推理对视频或者屏幕进行识别
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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('视频已提取完成,程序退出...') | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment