from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMenu, QAction
from ui.win import Ui_mainWindow
from PyQt5.QtCore import Qt, QPoint, QTimer, QThread, pyqtSignal
from PyQt5.QtGui import QImage, QPixmap, QPainter, QIcon
import sys
import os
import json
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import os
import time
import cv2
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)
from models.experimental import attempt_load
from utils.datasets import LoadImages, LoadWebcam
from utils.CustomMessageBox import MessageBox
from utils.general import check_img_size, check_requirements, check_imshow, colorstr, non_max_suppression, \
apply_classifier, scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
# from utils.plots import colors, plot_one_box, plot_one_box_PIL
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device
from utils.capnums import Camera
from deep_sort.deep_sort import DeepSort
def deepsort_update(Tracker, pred, xywh, np_img):
outputs = Tracker.update(xywh, pred[:, 4:5], pred[:, 5].tolist(), cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB))
return outputs
class DetThread(QThread):
send_img = pyqtSignal(np.ndarray)
send_raw = pyqtSignal(np.ndarray)
send_statistic = pyqtSignal(dict)
send_msg = pyqtSignal(str)
send_percent = pyqtSignal(int)
send_fps = pyqtSignal(str)
def __init__(self,tracker = 'ByteTrack',imgsz=(640,640)):
super(DetThread, self).__init__()
self.weights = './yolov5s.pt'
self.current_weight = './yolov5s.pt'
self.source = '0'
self.conf_thres = 0.25
self.iou_thres = 0.45
self.jump_out = False # jump out of the loop
self.is_continue = True # continue/pause
self.percent_length = 1000 # progress bar
self.rate_check = True # Whether to enable delay
self.rate = 100
self.imgsz = check_img_size(imgsz)
self.save_fold = './result'
if tracker == 'DeepSort':
self.tracker = DeepSort('deep_sort/deep_sort/deep/checkpoint/ckpt_car3.t7')
self._type = 'DeepSort'
@torch.no_grad()
def run(self,
imgsz=640,
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=True, # show results
save_txt=False, # save results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project='runs/detect', # save results to project/name
name='exp', # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
line_thickness=3, # bounding box thickness (pixels)
hide_labels=False, # hide labels
hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference
):
# 初始化
try:
# 选择设备(如CPU或GPU)
device = select_device(device)
# 如果设备不是CPU,设置half为True以使用半精度浮点数
half &= device.type != 'cpu'
# 尝试加载模型,并将模型加载到指定设备上
model = attempt_load(self.weights, map_location=device)
num_params = 0
# 计算模型参数的总数
for param in model.parameters():
num_params += param.numel()
# 获取模型的步幅(stride),并确保图像尺寸是步幅的倍数
stride = int(model.stride.max())
imgsz = check_img_size(imgsz, s=stride)
# 获取模型的类别名称
names = model.module.names if hasattr(model, 'module') else model.names
# 如果使用半精度浮点数,转换模型的权重为半精度
if half:
model.half()
# 如果使用CUDA设备,进行一次空推理以初始化模型
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))
deepsort_tracker = DeepSort('deep_sort/deep_sort/deep/checkpoint/ckpt_car3.t7')
# 数据加载
if self.source.isnumeric() or self.source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://')):
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadWebcam(self.source, img_size=imgsz, stride=stride)
# bs = len(dataset) # batch_size
else:
dataset = LoadImages(self.source, img_size=imgsz, stride=stride)
#根据输入源(视频流或本地图片/视频)加载数据。如果是视频流,启用 cudnn.benchmark 以加速推理
count = 0
jump_count = 0
start_time = time.time()
dataset = iter(dataset)
#主循环
while True:
if self.jump_out:
self.vid_cap.release()
self.send_percent.emit(0)
self.send_msg.emit('停止')
if hasattr(self, 'out'):
self.out.release()
break
#主循环处理每一帧数据。如果接收到跳出信号,则释放资源并退出循环
# 更换模型
# 如果当前权重与新权重不一致,更新模型
if self.current_weight != self.weights:
# 加载模型
model = attempt_load(self.weights, map_location=device) # 加载FP32模型(32位浮点数)
num_params = 0
# 计算模型参数总数
for param in model.parameters():
num_params += param.numel()
# 获取模型步幅(stride)
stride = int(model.stride.max())
# 检查图像尺寸是否为步幅的整数倍
imgsz = check_img_size(imgsz, s=stride)
# 获取模型的类别名称
names = model.module.names if hasattr(model, 'module') else model.names
# 如果使用半精度浮点数,转换模型权重为FP16(16位浮点数)
if half:
model.half()
# 如果设备不是CPU,运行一次模型以完成初始化
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))
# 更新当前权重
self.current_weight = self.weights
#如果检测到模型权重发生改变,重新加载模型。
#数据预处理
# 如果继续处理标志位为True,则执行以下代码
if self.is_continue:
# 从数据集中获取下一帧数据,包括路径、图像、原始图像和视频捕获对象
path, img, im0s, self.vid_cap = next(dataset)
count += 1 # 帧计数器加一
# 每处理30帧,计算一次FPS(帧率)
if count % 30 == 0 and count >= 30:
fps = int(30 / (time.time() - start_time)) # 计算30帧处理所需时间的平均FPS
self.send_fps.emit('fps:' + str(fps)) # 发送FPS信号
start_time = time.time() # 重新记录开始时间
# 如果有视频捕获对象,计算视频处理进度百分比
if self.vid_cap:
percent = int(count / self.vid_cap.get(cv2.CAP_PROP_FRAME_COUNT) * self.percent_length)
self.send_percent.emit(percent) # 发送进度百分比信号
else:
percent = self.percent_length # 如果没有视频捕获对象,设定进度为满
# 初始化统计字典,统计每个类别的数量
statistic_dic = {name: 0 for name in names}
# 将图像从numpy数组转换为torch张量,并加载到指定设备(如GPU)
img = torch.from_numpy(img).to(device)
# 如果使用半精度浮点数,转换图像数据类型,否则转换为浮点数
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 将像素值从0-255归一化到0.0-1.0
# 如果图像是三维的,增加一个维度以匹配模型输入要求
if img.ndimension() == 3:
img = img.unsqueeze(0)
# 推理
pred = model(img, augment=augment)[0]
#NMS
pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, classes, agnostic_nms, max_det=max_det)
# 处理检测结果 迭代每个检测结果,并绘制边界框和标签
# 遍历每张图像的检测结果
for i, det in enumerate(pred): # detections per image
# 复制原始图像,以便在上面绘制标注
im0 = im0s.copy()
# 创建一个Annotator对象,用于在图像上绘制标注
annotator = Annotator(im0, line_width=line_thickness, example=str(names))
# 如果有检测结果
if len(det):
# 将检测框的坐标从img_size尺度重新缩放到原始图像尺度,并四舍五入
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# 进行目标跟踪
track_outputs = deepsort_update(deepsort_tracker, det.cpu(), det[:, :4].cpu(), im0)
# Write results
# for *xyxy, conf, cls in reversed(det):
# c = int(cls) # integer class
# statistic_dic[names[c]] += 1
# label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
# annotator.box_label(xyxy, label, color=colors(c, True))
# 将跟踪结果绘制到图片上
if len(track_outputs)>0:
for track_output in track_outputs:
xyxy = track_output[:4]
c = int(track_output[4]) # integer class
track_id = 'ID_' + str(track_output[5])
statistic_dic[names[c]] += 1
label = (f'{track_id}') if hide_labels else (f'{track_id} {names[c]}')
annotator.box_label(xyxy, label, color=colors(c, True))
if self.rate_check:
time.sleep(1/self.rate)
#保存检测结果图像或视频,并通过信号发送处理后的图像和统计数据
im0 = annotator.result()
self.send_img.emit(im0)
self.send_raw.emit(im0s if isinstance(im0s, np.ndarray) else im0s[0])
self.send_statistic.emit(statistic_dic)
if self.save_fold:
os.makedirs(self.save_fold, exist_ok=True)
if self.vid_cap is None:
save_path = os.path.join(self.save_fold,
time.strftime('%Y_%m_%d_%H_%M_%S',
time.localtime()) + '.jpg')
cv2.imwrite(save_path, im0)
else:
if count == 1:
ori_fps = int(self.vid_cap.get(cv2.CAP_PROP_FPS))
if ori_fps == 0:
ori_fps = 25
# width = int(self.vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
# height = int(self.vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
width, height = im0.shape[1], im0.shape[0]
save_path = os.path.join(self.save_fold, time.strftime('%Y_%m_%d_%H_%M_%S', time.localtime()) + '.mp4')
self.out = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), ori_fps,
(width, height))
self.out.write(im0)
if percent == self.percent_length:
print(count)
self.send_percent.emit(0)
self.send_msg.emit('结束')
if hasattr(self, 'out'):
self.out.release()
break
except Exception as e:
self.send_msg.emit('%s' % e)
class MainWindow(QMainWindow, Ui_mainWindow):
def __init__(self, parent=None):
super(MainWindow, self).__init__(parent)
self.setupUi(self)
self.m_flag = False
# style 1: window can be stretched
# self.setWindowFlags(Qt.CustomizeWindowHint | Qt.WindowStaysOnTopHint)
# style 2: window can not be stretched
self.setWindowFlags(Qt.Window | Qt.FramelessWindowHint
| Qt.WindowSystemMenuHint | Qt.WindowMinimizeButtonHint | Qt.WindowMaximizeButtonHint)
# self.setWindowOpacity(0.85) # Transparency of window
self.minButton.clicked.connect(self.showMinimized)
self.maxButton.clicked.connect(self.max_or_restore)
# show Maximized window
self.maxButton.animateClick(10)
self.closeButton.clicked.connect(self.close)
self.qtimer = QTimer(self)
self.qtimer.setSingleShot(True)
self.qtimer.timeout.connect(lambda: self.statistic_label.clear())
# search models automatically
self.comboBox.clear()
self.pt_list = os.listdir('./pt')
self.pt_list = [file for file in self.pt_list if file.endswith('.pt')]
self.pt_list.sort(key=lambda x: os.path.getsize('./pt/'+x))
self.comboBox.clear()
self.comboBox.addItems(self.pt_list)
self.qtimer_search = QTimer(self)
self.qtimer_search.timeout.connect(lambda: self.search_pt())
self.qtimer_search.start(2000)
# yolov5 thread
self.det_thread = DetThread()
self.model_type = self.comboBox.currentText()
self.det_thread.weights = "./pt/%s" % self.model_type
self.det_thread.source = '0'
self.det_thread.percent_length = self.progressBar.maximum()
self.det_thread.send_raw.connect(lambda x: self.show_image(x, self.raw_video))
self.det_thread.send_img.connect(lambda x: self.show_image(x, self.out_video))
self.det_thread.send_statistic.connect(self.show_statistic)
self.det_thread.send_msg.connect(lambda x: self.show_msg(x))
self.det_thread.send_percent.connect(lambda x: self.progressBar.setValue(x))
self.det_thread.send_fps.connect(lambda x: self.fps_label.setText(x))
self.fileButton.clicked.connect(self.open_file)
self.cameraButton.clicked.connect(self.chose_cam)
#self.rtspButton.clicked.connect(self.chose_rtsp)
self.runButton.clicked.connect(self.run_or_continue)
self.stopButton.clicked.connect(self.stop)
self.comboBox.currentTextChanged.connect(self.change_model)
self.confSpinBox.valueChanged.connect(lambda x: self.change_val(x, 'confSpinBox'))
self.confSlider.valueChanged.connect(lambda x: self.change_val(x, 'confSlider'))
self.iouSpinBox.valueChanged.connect(lambda x: self.change_val(x, 'iouSpinBox'))
self.iouSlider.valueChanged.connect(lambda x: self.change_val(x, 'iouSlider'))
self.rateSpinBox.valueChanged.connect(lambda x: self.change_val(x, 'rateSpinBox'))
self.rateSlider.valueChanged.connect(lambda x: self.change_val(x, 'rateSlider'))
self.checkBox.clicked.connect(self.checkrate)
self.saveCheckBox.clicked.connect(self.is_save)
self.load_setting()
def search_pt(self):
pt_list = os.listdir('./pt')
pt_list = [file for file in pt_list if file.endswith('.pt')]
pt_list.sort(key=lambda x: os.path.getsize('./pt/' + x))
if pt_list != self.pt_list:
self.pt_list = pt_list
self.comboBox.clear()
self.comboBox.addItems(self.pt_list)
def is_save(self):
if self.saveCheckBox.isChecked():
self.det_thread.save_fold = './result'
else:
self.det_thread.save_fold = None
def checkrate(self):
if self.checkBox.isChecked():
self.det_thread.rate_check = True
else:
self.det_thread.rate_check = False
def chose_cam(self):
try:
self.stop()
MessageBox(
self.closeButton, title='提示', text='加载摄像头', time=2000, auto=True).exec_()
# get the number of local cameras
_, cams = Camera().get_cam_num()
popMenu = QMenu()
popMenu.setFixedWidth(self.cameraButton.width())
popMenu.setStyleSheet('''
QMenu {
font-size: 16px;
font-family: "Microsoft YaHei UI";
font-weight: light;
color:white;
padding-left: 5px;
padding-right: 5px;
padding-top: 4px;
padding-bottom: 4px;
border-style: solid;
border-width: 0px;
border-color: rgba(255, 255, 255, 255);
border-radius: 3px;
background-color: rgba(200, 200, 200,50);}
''')
for cam in cams:
exec("action_%s = QAction('%s')" % (cam, cam))
exec("popMenu.addAction(action_%s)" % cam)
x = self.groupBox_5.mapToGlobal(self.cameraButton.pos()).x()
y = self.groupBox_5.mapToGlobal(self.cameraButton.pos()).y()
y = y + self.cameraButton.frameGeometry().height()
pos = QPoint(x, y)
action = popMenu.exec_(pos)
if action:
self.det_thread.source = action.text()
self.statistic_msg('加载视频:{}'.format(action.text()))
except Exception as e:
self.statistic_msg('%s' % e)
def load_setting(self):
config_file = 'config/setting.json'
if not os.path.exists(config_file):
iou = 0.26
conf = 0.33
rate = 10
check = 0
savecheck = 0
new_config = {"iou": iou,
"conf": conf,
"rate": rate,
"check": check,
"savecheck": savecheck
}
new_json = json.dumps(new_config, ensure_ascii=False, indent=2)
with open(config_file, 'w', encoding='utf-8') as f:
f.write(new_json)
else:
config = json.load(open(config_file, 'r', encoding='utf-8'))
if len(config) != 5:
iou = 0.26
conf = 0.33
rate = 10
check = 0
savecheck = 0
else:
iou = config['iou']
conf = config['conf']
rate = config['rate']
check = config['check']
savecheck = config['savecheck']
self.confSpinBox.setValue(conf)
self.iouSpinBox.setValue(iou)
self.rateSpinBox.setValue(rate)
self.checkBox.setCheckState(check)
self.det_thread.rate_check = check
self.saveCheckBox.setCheckState(savecheck)
self.is_save()
def change_val(self, x, flag):
if flag == 'confSpinBox':
self.confSlider.setValue(int(x*100))
elif flag == 'confSlider':
self.confSpinBox.setValue(x/100)
self.det_thread.conf_thres = x/100
elif flag == 'iouSpinBox':
self.iouSlider.setValue(int(x*100))
elif flag == 'iouSlider':
self.iouSpinBox.setValue(x/100)
self.det_thread.iou_thres = x/100
elif flag == 'rateSpinBox':
self.rateSlider.setValue(x)
elif flag == 'rateSlider':
self.rateSpinBox.setValue(x)
self.det_thread.rate = x * 10
else:
pass
def statistic_msg(self, msg):
self.statistic_label.setText(msg)
# self.qtimer.start(3000)
def show_msg(self, msg):
self.runButton.setChecked(Qt.Unchecked)
self.statistic_msg(msg)
if msg == "Finished":
self.saveCheckBox.setEnabled(True)
def change_model(self, x):
self.model_type = self.comboBox.currentText()
self.det_thread.weights = "./pt/%s" % self.model_type
self.statistic_msg('Change model to %s' % x)
def open_file(self):
config_file = 'config/fold.json'
# config = json.load(open(config_file, 'r', encoding='utf-8'))
config = json.load(open(config_file, 'r', encoding='utf-8'))
open_fold = config['open_fold']
if not os.path.exists(open_fold):
open_fold = os.getcwd()
name, _ = QFileDialog.getOpenFileName(self, 'Video/image', open_fold, "Pic File(*.mp4 *.mkv *.avi *.flv "
"*.jpg *.png)")
if name:
self.det_thread.source = name
self.statistic_msg('Loaded file:{}'.format(os.path.basename(name)))
config['open_fold'] = os.path.dirname(name)
config_json = json.dumps(config, ensure_ascii=False, indent=2)
with open(config_file, 'w', encoding='utf-8') as f:
f.write(config_json)
self.stop()
def max_or_restore(self):
if self.maxButton.isChecked():
self.showMaximized()
else:
self.showNormal()
def run_or_continue(self):
self.det_thread.jump_out = False
if self.runButton.isChecked():
self.saveCheckBox.setEnabled(False)
self.det_thread.is_continue = True
if not self.det_thread.isRunning():
self.det_thread.start()
source = os.path.basename(self.det_thread.source)
source = 'camera' if source.isnumeric() else source
self.statistic_msg('正在检测 >> 模型:{},文件:{}'.
format(os.path.basename(self.det_thread.weights),
source))
else:
self.det_thread.is_continue = False
self.statistic_msg('暂停')
def stop(self):
self.det_thread.jump_out = True
self.saveCheckBox.setEnabled(True)
def mousePressEvent(self, event):
self.m_Position = event.pos()
if event.button() == Qt.LeftButton:
if 0 < self.m_Position.x() < self.groupBox.pos().x() + self.groupBox.width() and \
0 < self.m_Position.y() < self.groupBox.pos().y() + self.groupBox.height():
self.m_flag = True
def mouseMoveEvent(self, QMouseEvent):
if Qt.LeftButton and self.m_flag:
self.move(QMouseEvent.globalPos() - self.m_Position)
def mouseReleaseEvent(self, QMouseEvent):
self.m_flag = False
@staticmethod
def show_image(img_src, label):
try:
ih, iw, _ = img_src.shape
w = label.geometry().width()
h = label.geometry().height()
# keep original aspect ratio
if iw/w > ih/h:
scal = w / iw
nw = w
nh = int(scal * ih)
img_src_ = cv2.resize(img_src, (nw, nh))
else:
scal = h / ih
nw = int(scal * iw)
nh = h
img_src_ = cv2.resize(img_src, (nw, nh))
frame = cv2.cvtColor(img_src_, cv2.COLOR_BGR2RGB)
img = QImage(frame.data, frame.shape[1], frame.shape[0], frame.shape[2] * frame.shape[1],
QImage.Format_RGB888)
label.setPixmap(QPixmap.fromImage(img))
except Exception as e:
print(repr(e))
def show_statistic(self, statistic_dic):
try:
self.resultWidget.clear()
statistic_dic = sorted(statistic_dic.items(), key=lambda x: x[1], reverse=True)
statistic_dic = [i for i in statistic_dic if i[1] > 0]
results = [' '+str(i[0]) + ':' + str(i[1]) for i in statistic_dic]
self.resultWidget.addItems(results)
except Exception as e:
print(repr(e))
def closeEvent(self, event):
self.det_thread.jump_out = True
config_file = 'config/setting.json'
config = dict()
config['iou'] = self.confSpinBox.value()
config['conf'] = self.iouSpinBox.value()
config['rate'] = self.rateSpinBox.value()
config['check'] = self.checkBox.checkState()
config['savecheck'] = self.saveCheckBox.checkState()
config_json = json.dumps(config, ensure_ascii=False, indent=2)
with open(config_file, 'w', encoding='utf-8') as f:
f.write(config_json)
MessageBox(
self.closeButton, title='提示', text='正在关闭程序', time=2000, auto=True).exec_()
sys.exit(0)
if __name__ == "__main__":
app = QApplication(sys.argv)
myWin = MainWindow()
myWin.show()
# myWin.showMaximized()
sys.exit(app.exec_())
最新发布