class WidgetConfig(QGroupBox):
def __init__(self):
super(WidgetConfig, self).__init__()
HEIGHT = 40
grid = QGridLayout()
# 使用默认摄像头复选框
self.check_camera = QCheckBox('Use default camera')
self.check_camera.setChecked(False)
self.check_camera.stateChanged.connect(self.slot_check_camera)
grid.addWidget(self.check_camera, 0, 0, 1, 3) # 一行三列
# 选择视频文件
label_video = QLabel('Detect File')
self.line_video = QLineEdit()
if 'video' in GLOBAL.config:
self.line_video.setText(GLOBAL.config['video'])
self.line_video.setFixedHeight(HEIGHT)
self.line_video.setEnabled(False)
self.line_video.editingFinished.connect(
lambda: GLOBAL.record_config({'video': self.line_video.text()}))
self.btn_video = QPushButton('Choose')
self.btn_video.setFixedHeight(HEIGHT)
self.btn_video.setEnabled(False)
self.btn_video.clicked.connect(self.choose_video_file)
self.slot_check_camera()
grid.addWidget(label_video, 1, 0)
grid.addWidget(self.line_video, 1, 1)
grid.addWidget(self.btn_video, 1, 2)
# 选择权重文件
label_weights = QLabel('Weights File')
self.line_weights = QLineEdit()
if 'weights' in GLOBAL.config:
self.line_weights.setText(GLOBAL.config['weights'])
self.line_weights.setFixedHeight(HEIGHT)
self.line_weights.editingFinished.connect(lambda: GLOBAL.record_config(
{'weights': self.line_weights.text()}))
self.btn_weights = QPushButton('Choose')
self.btn_weights.setFixedHeight(HEIGHT)
self.btn_weights.clicked.connect(self.choose_weights_file)
grid.addWidget(label_weights, 2, 0)
grid.addWidget(self.line_weights, 2, 1)
grid.addWidget(self.btn_weights, 2, 2)
# 是否使用GPU
label_device = QLabel('CUDA device')
self.line_device = QLineEdit('gpu')
if 'device' in GLOBAL.config:
self.line_device.setText(GLOBAL.config['device'])
else:
self.line_device.setText('cpu')
self.line_device.setPlaceholderText('cpu or 0 or 0,1,2,3')
self.line_device.setFixedHeight(HEIGHT)
self.line_device.editingFinished.connect(
lambda: GLOBAL.record_config({'device': self.line_device.text()}))
grid.addWidget(label_device, 3, 0)
grid.addWidget(self.line_device, 3, 1, 1, 2)
# 设置图像大小
label_size = QLabel('Img Size')
self.combo_size = QComboBox()
self.combo_size.setFixedHeight(HEIGHT)
self.combo_size.setStyleSheet(
'QAbstractItemView::item {height: 40px;}')
self.combo_size.setView(QListView())
self.combo_size.addItem('320', 320)
self.combo_size.addItem('416', 416)
self.combo_size.addItem('480', 480)
self.combo_size.addItem('544', 544)
self.combo_size.addItem('640', 640)
self.combo_size.setCurrentIndex(2)
self.combo_size.currentIndexChanged.connect(
lambda: GLOBAL.record_config(
{'img_size': self.combo_size.currentData()}))
grid.addWidget(label_size, 4, 0)
grid.addWidget(self.combo_size, 4, 1, 1, 2)
#choose net camera
label_stream = QLabel('NetVedioStream')
self.combo_stream = QComboBox()
self.combo_stream.setFixedHeight(HEIGHT)
self.combo_stream.setStyleSheet(
'QAbstractItemView::item {height: 40px;}')
self.combo_stream.setView(QListView())
self.combo_stream.addItem('rtsp://*****:*****@192.168.0.65/',
'rtsp://*****:*****@192.168.0.65/')
self.combo_stream.addItem('rtsp://*****:*****@192.168.0.66/',
'rtsp://*****:*****@192.168.0.66/')
self.combo_stream.addItem('rtsp://*****:*****@192.168.0.67/',
'rtsp://*****:*****@192.168.0.67/')
self.combo_stream.addItem('rtsp://*****:*****@192.168.0.68/',
'rtsp://*****:*****@192.168.0.68/')
self.combo_stream.addItem('rtsp://*****:*****@192.168.0.65/',
'rtsp://*****:*****@192.168.0.65/')
self.combo_stream.setCurrentIndex(0)
self.combo_stream.currentIndexChanged.connect(
lambda: GLOBAL.record_config(
{'netstreamvedio': self.combo_stream.currentData()}))
grid.addWidget(label_stream, 5, 0)
grid.addWidget(self.combo_stream, 5, 1, 1, 2)
# 设置置信度阈值
label_conf = QLabel('Confidence')
self.spin_conf = QDoubleSpinBox()
self.spin_conf.setFixedHeight(HEIGHT)
self.spin_conf.setDecimals(1)
self.spin_conf.setRange(0.1, 0.9)
self.spin_conf.setSingleStep(0.1)
if 'conf_thresh' in GLOBAL.config:
self.spin_conf.setValue(GLOBAL.config['conf_thresh'])
else:
self.spin_conf.setValue(0.4) # 默认值
GLOBAL.record_config({'conf_thresh': 0.4})
self.spin_conf.valueChanged.connect(lambda: GLOBAL.record_config(
{'conf_thresh': round(self.spin_conf.value(), 1)}))
grid.addWidget(label_conf, 6, 0)
grid.addWidget(self.spin_conf, 6, 1, 1, 2)
# 设置IOU阈值
label_iou = QLabel('IOU')
self.spin_iou = QDoubleSpinBox()
self.spin_iou.setFixedHeight(HEIGHT)
self.spin_iou.setDecimals(1)
self.spin_iou.setRange(0.1, 0.9)
self.spin_iou.setSingleStep(0.1)
if 'iou_thresh' in GLOBAL.config:
self.spin_iou.setValue(GLOBAL.config['iou_thresh'])
else:
self.spin_iou.setValue(0.5) # 默认值
GLOBAL.record_config({'iou_thresh': 0.5})
self.spin_iou.valueChanged.connect(lambda: GLOBAL.record_config(
{'iou_thresh': round(self.spin_iou.value(), 1)}))
grid.addWidget(label_iou, 7, 0)
grid.addWidget(self.spin_iou, 7, 1, 1, 2)
# class-agnostic NMS
self.check_agnostic = QCheckBox('Agnostic')
if 'agnostic' in GLOBAL.config:
self.check_agnostic.setChecked(GLOBAL.config['agnostic'])
else:
self.check_agnostic.setChecked(True)
self.check_agnostic.stateChanged.connect(lambda: GLOBAL.record_config(
{'agnostic': self.check_agnostic.isChecked()}))
grid.addWidget(self.check_agnostic, 8, 0, 1, 3) # 一行三列
# augmented inference
self.check_augment = QCheckBox('Augment')
if 'augment' in GLOBAL.config:
self.check_augment.setChecked(GLOBAL.config['augment'])
else:
self.check_augment.setChecked(True)
self.check_augment.stateChanged.connect(lambda: GLOBAL.record_config(
{'augment': self.check_augment.isChecked()}))
grid.addWidget(self.check_augment, 9, 0, 1, 3) # 一行三列
self.setLayout(grid) # 设置布局
def slot_check_camera(self):
check = self.check_camera.isChecked()
GLOBAL.record_config({'use_camera': check}) # 保存配置
if check:
self.line_video.setEnabled(False)
self.btn_video.setEnabled(False)
else:
self.line_video.setEnabled(True)
self.btn_video.setEnabled(True)
def choose_weights_file(self):
"""从系统中选择权重文件"""
file = QFileDialog.getOpenFileName(
self, "Pre-trained YOLOv5 Weights", "./",
"Weights Files (*.pt);;All Files (*)")
if file[0] != '':
self.line_weights.setText(file[0])
GLOBAL.record_config({'weights': file[0]})
def choose_video_file(self):
"""从系统中选择视频文件"""
file = QFileDialog.getOpenFileName(self, "Video Files", "./",
"Video Files (*)")
if file[0] != '':
self.line_video.setText(file[0])
GLOBAL.record_config({'video': file[0]})
def save_config(self):
"""保存当前的配置到配置文件"""
config = {
'use_camera': self.check_camera.isChecked(),
'video': self.line_video.text(),
'weights': self.line_weights.text(),
'device': self.line_device.text(),
'img_size': self.combo_size.currentData(),
'conf_thresh': round(self.spin_conf.value(), 1),
'iou_thresh': round(self.spin_iou.value(), 1),
'agnostic': self.check_agnostic.isChecked(),
'augment': self.check_augment.isChecked(),
'netstreamvedio': self.combo_stream.currentData()
}
GLOBAL.record_config(config)
class SettingsDialog(QDialog):
def __init__(self):
super(SettingsDialog, self).__init__()
self.setWindowTitle('Settings')
self.setWindowIcon(QIcon('img/yologo.png'))
self.setWindowFlags(Qt.WindowCloseButtonHint)
self.setMinimumWidth(600)
HEIGHT = 30
grid = QGridLayout()
# 选择权重文件
label_weights = QLabel('Weights')
self.line_weights = QLineEdit()
self.line_weights.setFixedHeight(HEIGHT)
self.btn_weights = QPushButton('...')
self.btn_weights.setFixedWidth(40)
self.btn_weights.setFixedHeight(HEIGHT)
self.btn_weights.clicked.connect(self.choose_weights_file)
grid.addWidget(label_weights, 2, 0)
grid.addWidget(self.line_weights, 2, 1, 1, 2)
grid.addWidget(self.btn_weights, 2, 3)
# 是否使用GPU
label_device = QLabel('CUDA device')
self.line_device = QLineEdit('cpu')
self.line_device.setToolTip('cuda device, i.e. 0 or 0,1,2,3 or cpu')
self.line_device.setPlaceholderText('cpu or 0 or 0,1,2,3')
self.line_device.setFixedHeight(HEIGHT)
grid.addWidget(label_device, 3, 0)
grid.addWidget(self.line_device, 3, 1, 1, 3)
# 设置图像大小
label_size = QLabel('Img Size')
self.combo_size = QComboBox()
self.combo_size.setToolTip('inference size (pixels)')
self.combo_size.setFixedHeight(HEIGHT)
self.combo_size.setStyleSheet(
'QAbstractItemView::item {height: 40px;}')
self.combo_size.setView(QListView())
self.combo_size.addItem('320', 320)
self.combo_size.addItem('384', 384)
self.combo_size.addItem('448', 448)
self.combo_size.addItem('512', 512)
self.combo_size.addItem('576', 576)
self.combo_size.addItem('640', 640)
grid.addWidget(label_size, 4, 0)
grid.addWidget(self.combo_size, 4, 1, 1, 3)
# 设置置信度阈值
label_conf = QLabel('Confidence')
self.spin_conf = QDoubleSpinBox()
self.spin_conf.setToolTip('confidence threshold')
self.spin_conf.setFixedHeight(HEIGHT)
self.spin_conf.setDecimals(1)
self.spin_conf.setRange(0.1, 0.9)
self.spin_conf.setSingleStep(0.1)
grid.addWidget(label_conf, 5, 0)
grid.addWidget(self.spin_conf, 5, 1, 1, 3)
# 设置IOU阈值
label_iou = QLabel('IOU')
self.spin_iou = QDoubleSpinBox()
self.spin_iou.setToolTip('NMS IoU threshold')
self.spin_iou.setFixedHeight(HEIGHT)
self.spin_iou.setDecimals(1)
self.spin_iou.setRange(0.1, 0.9)
self.spin_iou.setSingleStep(0.1)
grid.addWidget(label_iou, 6, 0)
grid.addWidget(self.spin_iou, 6, 1, 1, 3)
# maximum detections per image
label_max_det = QLabel('maximum detections')
self.spin_max_det = QDoubleSpinBox()
self.spin_max_det.setToolTip('Maximum detections per image')
self.spin_max_det.setFixedHeight(HEIGHT)
self.spin_max_det.setDecimals(0)
self.spin_max_det.setRange(10, 1000)
self.spin_max_det.setSingleStep(10)
grid.addWidget(label_max_det, 7, 0)
grid.addWidget(self.spin_max_det, 7, 1, 1, 3)
# class-agnostic NMS
self.check_agnostic = QCheckBox('Agnostic')
self.check_agnostic.setToolTip('class-agnostic NMS')
# augmented inference
self.check_augment = QCheckBox('Augment')
self.check_augment.setToolTip('augmented inference')
# half
self.check_half = QCheckBox('Half')
self.check_half.setToolTip('use FP16 half-precision inference')
grid.addWidget(self.check_agnostic, 8, 0)
grid.addWidget(self.check_augment, 8, 1)
grid.addWidget(self.check_half, 8, 2)
# use OpenCV DNN for ONNX inference
self.check_dnn = QCheckBox('DNN')
self.check_dnn.setToolTip('Use OpenCV DNN for ONNX inference')
grid.addWidget(self.check_dnn, 9, 0)
box = QGroupBox()
box.setLayout(grid)
hbox = QHBoxLayout()
self.btn_cancel = QPushButton('Cancel')
self.btn_cancel.clicked.connect(self.restore)
self.btn_ok = QPushButton('Ok')
self.btn_ok.clicked.connect(self.save_settings)
hbox.addStretch()
hbox.addWidget(self.btn_cancel)
hbox.addWidget(self.btn_ok)
vbox = QVBoxLayout()
vbox.addWidget(box)
vbox.addLayout(hbox)
self.setLayout(vbox)
self.load_settings()
def choose_weights_file(self):
"""从系统中选择权重文件"""
weights_path = '.'
weights = gb.get_config('weights', '')
if os.path.exists(weights):
weights_path = os.path.dirname(weights)
file = QFileDialog.getOpenFileName(
self, "Pre-trained YOLOv5 Weights", weights_path,
"Weights Files (*.pt);;All Files (*)")
if file[0] != '':
self.line_weights.setText(file[0])
def load_settings(self):
self.line_weights.setText(gb.get_config('weights', ''))
self.line_device.setText(gb.get_config('device', 'cpu'))
self.combo_size.setCurrentText(gb.get_config('img_size', '640'))
self.spin_conf.setValue(gb.get_config('conf_thresh', 0.5))
self.spin_iou.setValue(gb.get_config('iou_thresh', 0.5))
self.spin_max_det.setValue(gb.get_config('max_det', 50))
self.check_agnostic.setChecked(gb.get_config('agnostic', True))
self.check_augment.setChecked(gb.get_config('augment', True))
self.check_half.setChecked(gb.get_config('half', True))
self.check_dnn.setChecked(gb.get_config('dnn', False))
def save_settings(self):
"""更新配置"""
config = {
'weights': self.line_weights.text(),
'device': self.line_device.text(),
'img_size': self.combo_size.currentText(),
'conf_thresh': round(self.spin_conf.value(), 1),
'iou_thresh': round(self.spin_iou.value(), 1),
'max_det': int(self.spin_max_det.value()),
'agnostic': self.check_agnostic.isChecked(),
'augment': self.check_augment.isChecked(),
'half': self.check_half.isChecked(),
'dnn': self.check_dnn.isChecked()
}
gb.record_config(config)
self.accept()
def restore(self):
"""恢复原配置"""
self.load_settings()
self.reject()
def closeEvent(self, event):
self.restore()
