import os,shutil,copy,time,copy
os.environ["OPENCV_IO_MAX_IMAGE_PIXELS"] = pow(2,40).__str__()
import asyncio
from numpy.core.shape_base import block
# from sqlalchemy.sql.expression import false, true
from torch._C import wait
import websockets,json,functools
from yolov5_fz import YOLOv5
import torch,cv2
import torch.multiprocessing as mp
import numpy as np
from utils.datasets import LoadImages,LoadStreams,LoadWebcam
from utils.plots import colors
import subprocess as sp
# from scipy.stats import mode
from pycomm3 import LogixDriver
import traceback
def train_yolov5_ocr(q, message_log, moving_txt, q_stream, os_relative_path, q_train_websockets_txt, q_train_buckle,
q_frame): # 识别主程序
if not os.path.exists(os.path.join(os_relative_path, "train_archive_data")):
os.mkdir(os.path.join(os_relative_path, "train_archive_data"))
if os.path.exists("imgs"):
shutil.rmtree("imgs")
os.mkdir("imgs") # 创建图片暂存文件夹
label = ["train_id", "train_infor", "train_module", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "A", "C", "E",
"K", "T", "H", "B", "train_separate", "train_head", "M", "F"] # 模型类别
train_single = [] # 识别结果暂存列表
train_single_sta = 0 # 上一车间隔帧数,间隔状态
train_nums = 0 # 装车节数
font = cv2.FONT_HERSHEY_SIMPLEX
mp4_nums = 0 # 单列车视频帧数计数
train_couple_min = 0
train_separate_min = 0
train_couple_max = 0
train_single_mp4_name, train_single_mp4_name_next = None, None
train_one_start = 0 # 列车进站开始识别信号
train_date_name_single = "" # 当列车默认存档文件夹
train_out_station = 0 # 开始出站
train_out_sta = "waiting"
s_time, s_out_time = 0, 0
out_moving_x = 0
train_direction = [0] * 3
train_single_cp = [] # 识别列表暂存
train_in_time = 0 # 进站结束时间
train_in_sta = 0
error_in_time = 0
train_single_staing = False
while True:
try:
if q_frame.empty():
continue
img, im0s, xyxy, conf, name_id = q_frame.get(block=True, timeout=0.00001)
mp4_nums += 1
########################
if train_out_station == 1: # 出站
if train_out_sta == "starting" and moving_txt[0] == "kong" and s_time == 0:
s_time = time.time()
elif train_out_sta == "starting" and (moving_txt[0] == "zai" or moving_txt[0] == "chu"):
if s_out_time == 0:
s_out_time = time.time()
s_time = 0
if 1 in name_id: # 车辆信息
train_infor_name_index = name_id.index(1)
plot_one_box(xyxy[train_infor_name_index], im0s, label="train_infor", color=colors(
0, True), line_thickness=3)
if 0 in name_id: # 车号
train_infor_name_index = name_id.index(0)
plot_one_box(xyxy[train_infor_name_index], im0s, label="train_id", color=colors(
0, True), line_thickness=3)
if 2 in name_id: # 车型
train_infor_name_index = name_id.index(2)
plot_one_box(xyxy[train_infor_name_index], im0s, label="train_module", color=colors(
0, True), line_thickness=3)
if train_in_sta == 0 and train_out_sta == "waiting" and time.time() - train_in_time > 30:
message_sta = {"type": 2, "data": {"traveling_sta": 1, "state": 1}}
try:
shutil.rmtree("imgs")
os.mkdir("imgs") # 创建图片暂存文件夹
except:
pass
try:
shutil.rmtree("train_buckle")
os.mkdir("train_buckle") # 创建图片暂存文件夹
except:
pass
q_train_websockets_txt.put(message_sta) # 进站结束
message_log.put(f"{message_sta},sta:进站结束\n")
train_in_sta = 1
if (moving_txt[0] == "zai" or moving_txt[
0] == "chu") and train_out_sta == "waiting" and time.time() - train_in_time > 30:
train_out_sta = "starting"
message_sta = {"type": 2, "data": {"traveling_sta": 0, "state": 0}}
q_train_websockets_txt.put(message_sta) # 出站开始
message_log.put(f"{message_sta},sta:出站开始\n")
mp4_nums = 0
if (1 in name_id or 20 in name_id) and (
0 in name_id and 2 in name_id) and train_single_sta == 0: # 出站判断
train_couple_min += 1
if train_couple_min > 7 and train_single_sta == 0:
train_couple_min = 0
train_single_sta = 1
if 0 in name_id and 2 in name_id:
out_moving_index = name_id.index(2)
out_moving_x = int((xyxy[out_moving_index][0] + xyxy[out_moving_index][2]) / 2)
if 1 not in name_id and 20 not in name_id and 0 not in name_id and 2 not in name_id and train_single_sta == 1:
train_couple_min += 1
if train_single_sta == 1 and train_couple_min > 20:
if out_moving_x < 960:
train_nums -= 1
elif out_moving_x > 960:
train_nums += 1
message_log.put(f"out_station,{train_nums},{out_moving_x}\n")
train_couple_min = 0
train_single_sta = 0
if moving_txt[0] == "kong" and time.time() - s_time > 300 and s_time != 0:
train_nums = 0 # 2空轨 1在轨
train_out_station = 0
s_out_time, s_time = 0, 0
message_sta = {"type": 2, "data": {"traveling_sta": 0, "state": 1}}
q_train_websockets_txt.put(message_sta) # 出站结束
train_in_sta = 0
train_out_sta = "waiting"
message_log.put(f"{message_sta},sta:出站结束\n")
cv2.putText(im0s, f"Train out of station,sta:{train_out_sta},{train_nums},{moving_txt[0]}", (100, 100),
font, 2, (0, 255, 255), 2)
frame = cv2.resize(im0s, (640, 480))
q_stream.put(["train_in_stop", frame])
continue
#####################
if train_one_start == 1:
if error_in_time == 0 and moving_txt[0] == "kong":
error_in_time = time.time()
elif moving_txt[0] == "zai" or moving_txt[0] == "chu":
error_in_time = 0
if moving_txt[0] == "kong" and time.time() - error_in_time > 720 and error_in_time != 0:
error_in_time = 0
train_one_start = 0
train_couple_min = 0
train_separate_min = 0
message_sta = {"type": 2, "data": {"traveling_sta": 1, "state": 1}}
q_train_websockets_txt.put(message_sta) # 进站结束
if train_nums > 20 and train_single != []:
message_log.put(f"{message_sta},sta:进站意外结束,记录最后一次数据\n")
train_in_time = time.time()
train_nums = train_nums + 1
q.put({"train_nums": train_nums, "train_time_name": train_date_name_single,
"train_mp4_name": train_single_mp4_name, "train_single": train_single})
message_log.put(f"{message_sta},{train_direction},{train_nums}\n")
train_out_station = 1 # 列车开始出站
train_in_time = time.time()
train_one_start = 0 # 列车进站结束
else:
message_log.put(f"{message_sta},sta:进站意外结束\n")
train_nums = 0
time.sleep(1)
try:
shutil.rmtree("imgs")
os.mkdir("imgs") # 创建图片暂存文件夹
except:
pass
try:
shutil.rmtree("train_buckle")
os.mkdir("train_buckle") # 创建图片暂存文件夹
except:
pass
del train_single[:] # 退车
del train_single_cp[:]
if train_one_start == 0 and (1 in name_id or ((0 in name_id or 2 in name_id or 1) and (
3 in name_id or 4 in name_id or 5 in name_id or 6 in name_id or 7 in name_id or 8 in name_id or 9 in name_id or 10 in name_id or 11 in name_id))):
train_date_name_single = "train_archive_data/" + time.strftime("%Y%m%d%H%M%S", time.localtime())
error_in_time = 0
train_couple_min = 3
os.mkdir(f"{os_relative_path}/{train_date_name_single}") # 创建列车存档文件夹
os.mkdir(f"{os_relative_path}/{train_date_name_single}/{str(train_nums + 1).zfill(3)}") # 创建列车单节存档文件夹
train_single_mp4_name = f"{train_date_name_single}/{str(train_nums + 1).zfill(3)}/train_date.mp4"
train_one_start = 1 # 开始当前列视频采集和数据处理
message_sta = {"type": 2, "data": {"traveling_sta": 1, "state": 0}} # 当第一节车时,发送进车开始
message_log.put(f"{message_sta},sta:进站开始\n")
elif train_one_start != 1:
cv2.putText(im0s, f"sta:waiting", (100, 100), font, 3, (0, 255, 255), 2) # 开始等待
q_stream.put(["train_not_started", im0s])
continue
# if mp4_nums%25==0:
# q_train_buckle.put((img,im0s))
if train_single_sta == 0: # 方向判断
if train_couple_min != -1 and 1 in name_id and (
0 not in name_id and 2 not in name_id and 20 not in name_id):
train_couple_min += 1
if train_couple_min >= 3:
train_infor_name_index = name_id.index(1)
train_infor_name_xy = xyxy[train_infor_name_index]
img_w = im0s[int(train_infor_name_xy[1]):int(train_infor_name_xy[3]),
int(train_infor_name_xy[0]):int(train_infor_name_xy[2])]
cv2.putText(img_w, "hg", (0, 20), font, 1, (0, 0, 255), 1) # logo
cv2.imwrite("imgs/train_infor.jpg", img_w) # 截取车辆信息图片
train_couple_min = -1
if 20 in name_id and 21 not in name_id:
train_separate_min += 1
# print(f"train_separate_min:{train_separate_min},train_nums:{train_nums+1}\n")
# if ( train_separate_min >2 ) or (1 in name_id and (0 in name_id and 2 in name_id)): #当前列即将结束 空挡 车辆信息 车头
# if ( train_separate_min >1 and (1 in name_id or 0 in name_id or 2 in name_id)) or (1 in name_id and (0 in name_id and 2 in name_id)): #当前列即将结束 空挡 车辆信息 车头
if (train_separate_min >= 1):
# print('train_separate_min:', train_separate_min, ", curr node over.")
train_single_sta = 1 # 此节车准备结束信号
if train_separate_min > 2:
train_single_staing = True
else:
train_single_staing = False
# train_single_cp=copy.deepcopy(train_single)
# del train_single[:]#获取当节车号信息,并清空列表,等待下次
if 1 in name_id:
train_infor_name_index = name_id.index(1)
train_infor_name_xy = xyxy[train_infor_name_index]
img_w = im0s[int(train_infor_name_xy[1]):int(train_infor_name_xy[3]),
int(train_infor_name_xy[0]):int(train_infor_name_xy[2])]
cv2.putText(img_w, "hg", (0, 20), font, 1, (0, 0, 255), 1) # logo
cv2.imwrite("imgs/train_infor_next.jpg", img_w) # 截取车辆信息图片
try:
os.mkdir(
f"{os_relative_path}/{train_date_name_single}/{str(train_nums + 2).zfill(3)}") # 创建列车预先下节存档文件夹
except:
pass
train_single_mp4_name_next = train_single_mp4_name
train_single_mp4_name = f"{train_date_name_single}/{str(train_nums + 2).zfill(3)}/train_date.mp4" # 预先下一节视频
if train_single_sta == 1: # 此节车准备结束信号
if not train_single_staing:
if 1 not in name_id and 0 not in name_id and 2 not in name_id: # 当前节结束判断
train_couple_max += 1
else:
train_couple_max = 0
pass
else:
if 20 not in name_id: # 当前节结束判断
train_couple_max += 1
else:
train_couple_max = 0
if 1 in name_id:
train_separate_index = name_id.index(1)
train_separate_xy = xyxy[train_separate_index]
train_direction[1] = float((train_separate_xy[0] + train_separate_xy[2]) / 2) # 获取车辆信息最后出现的x坐标中点
if not os.path.exists("imgs/train_infor_next.jpg"):
train_infor_name_index = name_id.index(1)
train_infor_name_xy = xyxy[train_infor_name_index]
img_w = im0s[int(train_infor_name_xy[1]):int(train_infor_name_xy[3]),
int(train_infor_name_xy[0]):int(train_infor_name_xy[2])]
cv2.putText(img_w, "hg", (0, 20), font, 1, (0, 0, 255), 1) # logo
cv2.imwrite("imgs/train_infor_next.jpg", img_w) # 截取车辆信息图片
if 0 in name_id:
train_separate_index = name_id.index(0)
train_separate_xy = xyxy[train_separate_index]
train_direction[0] = float((train_separate_xy[0] + train_separate_xy[2]) / 2) # 获取车号最后出现的x坐标中点
if 2 in name_id:
train_separate_index = name_id.index(2)
train_separate_xy = xyxy[train_separate_index]
train_direction[2] = float((train_separate_xy[0] + train_separate_xy[2]) / 2) # 获取车型最后出现的x坐标中点
if train_couple_max >= 1 and train_direction != [0, 0, 0]: # 当前列结束,开始下一列
train_single_staing = False
f_x = sum(train_direction_i < 960 for train_direction_i in train_direction) - sum(
train_direction_i == 0 for train_direction_i in train_direction)
train_single_cp = copy.deepcopy(train_single)
del train_single[:] # 获取当节车号信息,并清空列表,等待下次
train_nums += 1 # 节数预加1
train_couple_min = 0 # 车辆信息计数归零
train_separate_min = 0 # 空挡计数归零
train_ocr_module = False
train_ocr_id = False
for train_ocr_i in train_single_cp:
if train_ocr_i[0] == "train_module":
train_ocr_module = True
if train_ocr_i[0] == "train_id":
train_ocr_id = True
if train_ocr_module == True and train_ocr_id == True:
break
if train_ocr_module != True or train_ocr_id != True:
train_nums -= 1
message_log.put(f"发现无效节数,已自行清理\n")
else:
message_log.put(f"车厢间隙,train_direction:{train_direction},计数+1\n")
# if f_x>1 :
# train_nums -= 2
# print('f_x:', f_x, 'train_numbs:', train_nums)
# message_log.put(f"{message_sta},{train_direction},{train_nums},发现倒车\n")
# train_single_mp4_name=f"{train_date_name_single}/{str(train_nums+1).zfill(3)}/train_date.mp4"#预先下一节视频
# q.put({"train_nums":-1,"train_time_name":-1,"train_mp4_name":-1,"train_single":train_single_cp})
# else:
q.put({"train_nums": train_nums, "train_time_name": train_date_name_single,
"train_mp4_name": train_single_mp4_name_next, "train_single": train_single_cp})
message_log.put(f"{message_sta},{train_direction},{train_nums}\n")
train_single_cp = []
train_direction = [0, 0, 0]
train_couple_max = 0
train_single_sta = 0
if moving_txt[0] == "chu" or moving_txt[0] == "kong":
train_in_time = time.time()
train_out_station = 1 # 列车开始出站
train_one_start = 0 # 列车进站结束
if train_nums > 20 and train_single != []:
message_log.put(f"{message_sta},sta:进站结束,补全最后一次数据\n")
train_nums = train_nums + 1
q.put({"train_nums": train_nums, "train_time_name": train_date_name_single,
"train_mp4_name": train_single_mp4_name, "train_single": train_single})
message_log.put(f"{message_sta},{train_direction},{train_nums}\n")
message_log.put(f"sta:检测到车头,准备进站结束,内\n")
if 1 in name_id: # 车辆信息识别框绘制
train_infor_name_index = name_id.index(1)
train_infor_name_xy = xyxy[train_infor_name_index]
plot_one_box(train_infor_name_xy, im0s, label="train_infor", color=colors(
0, True), line_thickness=3)
del name_id[train_infor_name_index]
xyxy = xyxy[torch.arange(xyxy.size(0)) != train_infor_name_index] # 删除车辆信息识别框
if 21 in name_id: # 车辆信息识别框绘制
train_infor_name_index = name_id.index(21)
train_infor_name_xy = xyxy[train_infor_name_index]
plot_one_box(train_infor_name_xy, im0s, label="train_head", color=colors(
0, True), line_thickness=3)
del name_id[train_infor_name_index]
xyxy = xyxy[torch.arange(xyxy.size(0)) != train_infor_name_index] # 删除车辆信息识别框
if 20 in name_id: # 车辆信息识别框绘制
train_infor_name_index = name_id.index(20)
train_infor_name_xy = xyxy[train_infor_name_index]
plot_one_box(train_infor_name_xy, im0s, label="train_separate", color=colors(
0, True), line_thickness=3)
del name_id[train_infor_name_index]
xyxy = xyxy[torch.arange(xyxy.size(0)) != train_infor_name_index] # 删除车辆信息识别框
if (0 in name_id or 2 in name_id):
ocr_list = {}
try:
id_index = name_id.index(0)
ocr_list["train_id"] = xyxy[id_index].int()
xyxy = xyxy[torch.arange(xyxy.size(0)) != id_index]
del name_id[id_index]
except:
pass
try:
id_index = name_id.index(2)
ocr_list["train_module"] = xyxy[id_index].int()
xyxy = xyxy[torch.arange(xyxy.size(0)) != id_index]
del name_id[id_index]
except:
pass
name_id = torch.tensor(name_id)
xyxy_index = xyxy[:, 0].sort()[1]
name_id = name_id[xyxy_index]
xyxy = xyxy[xyxy[:, 0].sort()[1]]
res_ocr = {}
for key in ocr_list.keys():
res_ocr[key] = ""
for i in range(len(xyxy)):
center_y = int((xyxy[i][1] + xyxy[i][3]) / 2)
for key, val in ocr_list.items():
if val[1] < center_y < val[3]:
res_ocr[key] = res_ocr[key] + label[int(name_id[i])]
for key in res_ocr.keys():
if len(res_ocr[key]) > 2:
train_single.append([key, res_ocr[key]]) # 采集汇总车号车型信息
img_w = im0s[ocr_list[key][1]:ocr_list[key][3], ocr_list[key][0]:ocr_list[key][2]]
cv2.putText(img_w, "hg", (0, 20), font, 1, (0, 255, 255), 1) # logo
cv2.imwrite("{}/imgs/{}_{}_{}.jpg".format(os.getcwd(), key, train_nums + 1, mp4_nums), img_w)
plot_one_box(ocr_list[key], im0s, label=key + "--" +
res_ocr[key], color=colors(0, True), line_thickness=3)
q_stream.put([f"{os_relative_path}/{train_single_mp4_name}", im0s, train_nums + 1])
except Exception as e:
message_log.put(f"{e.__traceback__.tb_lineno},{e},发现错误\n")