import cv2
import os
from os.path import join
from tqdm import tqdm
import numpy as np
import copy
import xlwings as xw
class dataItem:
def __init__(self, name, recognition, unRecognition):
self.name = name # 名称
self.recognition = recognition # 识别数量
self.unRecognition = unRecognition # 未识别数量
self.total = recognition + unRecognition # 总数
self.recognitionRate = '{:.2f}'.format(recognition / (recognition + unRecognition)) # 识别率
def excel(data_list):
new_path = join(path, 'test.xlsx')
app = xw.App(visible=False, add_book=False)
# 不显示Excel消息框
app.display_alerts = False
# 关闭屏幕更新,可加快宏的执行速度
app.screen_updating = False
isExist = True
if not os.path.exists(new_path):
wb = app.books.add()
isExist = False
else:
wb = app.books.open(new_path)
sheet = wb.sheets["sheet1"]
if not isExist:
sheet.range('A1').value = '名称'
sheet.range('B1').value = '识别数量'
sheet.range('C1').value = '未识别数量'
sheet.range('D1').value = '总数'
sheet.range('E1').value = '识别率'
# 读取行数
rows = sheet.used_range.last_cell.row
for data in data_list:
cur_row = rows + 1
sheet.range(cur_row, 1).value = data.name
sheet.range(cur_row, 2).value = data.recognition
sheet.range(cur_row, 3).value = data.unRecognition
sheet.range(cur_row, 4).value = data.total
sheet.range(cur_row, 5).value = data.recognitionRate
wb.save(new_path)
wb.close()
# 退出excel程序，
app.quit()
def detect_img(file_path):
file_path_list.clear()
filter_files_dir(file_path, file_path_list, '.jpg')
lable = os.path.split(file_path)[1]
cout = 0
recognition = 0
UnRecognition = 0
data_list = []
for file in file_path_list:
file_name = os.path.basename(file)
file_name = file_name.split('.')[0]
img = cv2.imread(file)
result = label_seg(img)
dir_path = join(file_path, str(result))
if not os.path.exists(dir_path):
os.mkdir(dir_path)
new_path = join(dir_path, file_name + '.jpg')
cv2.imwrite(new_path, img)
cout = cout + 1
print(cout, new_path)
# 统计
if result == 1:
recognition = recognition + 1
else:
UnRecognition = UnRecognition + 1
data = dataItem(lable, recognition, UnRecognition)
data_list.append(data)
excel(data_list)
print(lable, " 完成！")
def detect_video(file):
capture = cv2.VideoCapture(file)
p_path = os.path.abspath(os.path.join(file, ".."))
lable = os.path.split(os.path.split(file)[0])[-1]
cout = 0
recognition = 0
UnRecognition = 0
data_list = []
if capture.isOpened():
while True:
ret, img = capture.read() # img 就是一帧图片
if ret:
result = label_seg(img)
dir_path = join(p_path, str(result))
if not os.path.exists(dir_path):
os.mkdir(dir_path)
new_path = join(dir_path, str(cout) + '.png')
cv2.imwrite(new_path, img)
cout = cout + 1
print(cout, new_path)
# 统计
if result == 1:
recognition = recognition + 1
else:
UnRecognition = UnRecognition + 1
else:
data = dataItem(lable, recognition, UnRecognition)
data_list.append(data)
excel(data_list)
print("视频播放完成！")
break
else:
print('视频打开失败！')
capture.release()
def label_seg(img):
# src_img = copy.deepcopy(img)
# 在彩色图像的情况下，解码图像将以b g r顺序存储通道。
grid_RGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 从RGB色彩空间转换到HSV色彩空间
grid_HSV = cv2.cvtColor(grid_RGB, cv2.COLOR_RGB2HSV)
lower1 = np.array([1, 95, 245])
upper1 = np.array([5, 110, 255])
label_mark1 = cv2.inRange(grid_HSV, lower1, upper1) # mask1 为二值图像
lower2 = np.array([0, 195, 245])
upper2 = np.array([5, 210, 255])
label_mark2 = cv2.inRange(grid_HSV, lower2, upper2) # mask1 为二值图像
lower3 = np.array([175, 190, 245])
upper3 = np.array([185, 210, 255])
label_mark3 = cv2.inRange(grid_HSV, lower3, upper3) # mask1 为二值图像
# 11.181818181818187 225.3012048192771 249.0
# 11.336405529953907 224.02834008097165 247.0
lower4 = np.array([8, 210, 240])
upper4 = np.array([13, 250, 255])
label_mark4 = cv2.inRange(grid_HSV, lower4, upper4) # mask1 为二值图像
# 71.8840579710145 247.8169014084507 213.0
# 72.14634146341463 247.74881516587675 211.0
# 72.0 247.74881516587675 211.0
lower5 = np.array([65, 240, 205])
upper5 = np.array([80, 255, 220])
label_mark5 = cv2.inRange(grid_HSV, lower5, upper5) # mask1 为二值图像
label_mark = label_mark1 + label_mark2 + label_mark3 + label_mark4 + label_mark5
contours, hierarchy = cv2.findContours(label_mark, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
contours = sorted(contours, key=cv2.contourArea, reverse=True)
if len(contours) == 0:
return 0
arc = cv2.arcLength(contours[0], True)
if len(contours[0]) <= 5 and arc < 50:
return 0
# cv2.polylines(img, [contours[0][:, 0]], True, (255, 0, 0), 1)
return 1
# 获取当前路径下的指定mp4文件路径
def filter_files_dir(path, file_path_list, file_type):
filelist = os.listdir(path) # 该文件夹下所有的文件（包括文件夹）
for files in tqdm(filelist, ''): # 遍历所有文件
olddir = os.path.join(path, files) # 原来的文件路径
if os.path.isdir(olddir): # 如果是文件夹则跳过
filter_files_dir(olddir, file_path_list, file_type)
elif files[-4:] == file_type:
img_path = join(path, files)
file_path_list.append(img_path)
# time.sleep(0.000001)
def pre_video():
# 删除文件
filter_files_dir(path, file_path_list, '.png')
for file in file_path_list:
os.remove(file)
file_path_list.clear()
filter_files_dir(path, file_path_list, '.avi')
for file in file_path_list:
detect_video(file)
def pre_image():
# 删除文件
filter_files_dir(path, file_path_list, '.png')
for file in file_path_list:
os.remove(file)
file_path_list.clear()
filelist = os.listdir(path) # 该文件夹下所有的文件（包括文件夹）
for file in filelist:
dst_dir = os.path.join(path, file) # 原来的文件路径
detect_img(dst_dir)
path = r'runs\detect240319x'
# path = r'D:\SJZhang\yolov5\runs\detect1'
file_path_list = []
pre_image()
# path = r'D:\AOPENG_SVN\ALPHA\Alex\Aopeng.Alex\Aopeng.Alex.ImageNavigation\alex-stent\Test\EM_out\detect\exp\0\36.jpg'
# img = cv2.imread(path)
# label_seg(img)
浙公网安备 33010602011771号