339 lines
11 KiB
Python
339 lines
11 KiB
Python
|
import sys
|
|||
|
|
import os
|
|||
|
|
import random
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from PySide6.QtCore import QTimer, Qt
|
|||
|
|
from PySide6.QtGui import QImage, QPixmap
|
|||
|
|
from PySide6.QtWidgets import QApplication, QWidget, QLabel, QPushButton, QSlider, QVBoxLayout, QProgressBar
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
# 加载动画界面
|
|||
|
|
# ======================================================
|
|||
|
|
class LoadingScreen(QWidget):
|
|||
|
|
def __init__(self):
|
|||
|
|
super().__init__()
|
|||
|
|
self.setWindowTitle("Face Random Selector")
|
|||
|
|
self.setFixedSize(300, 150)
|
|||
|
|
self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
|
|||
|
|
self.setStyleSheet("""
|
|||
|
|
background-color: #1E90FF;
|
|||
|
|
border-radius: 10px;
|
|||
|
|
""")
|
|||
|
|
|
|||
|
|
# 创建布局
|
|||
|
|
layout = QVBoxLayout()
|
|||
|
|
layout.setAlignment(Qt.AlignCenter)
|
|||
|
|
layout.setSpacing(15)
|
|||
|
|
|
|||
|
|
# 加载提示文字
|
|||
|
|
self.loading_label = QLabel("正在加载模型...")
|
|||
|
|
self.loading_label.setAlignment(Qt.AlignCenter)
|
|||
|
|
self.loading_label.setStyleSheet("""
|
|||
|
|
color: white;
|
|||
|
|
font-size: 18px;
|
|||
|
|
font-weight: bold;
|
|||
|
|
""")
|
|||
|
|
layout.addWidget(self.loading_label)
|
|||
|
|
|
|||
|
|
# 进度条 - 设置为不确定模式，显示加载动画
|
|||
|
|
self.progress_bar = QProgressBar()
|
|||
|
|
self.progress_bar.setRange(0, 0) # 设置为0,0表示不确定模式
|
|||
|
|
self.progress_bar.setFixedHeight(10)
|
|||
|
|
self.progress_bar.setStyleSheet("""
|
|||
|
|
QProgressBar {
|
|||
|
|
border: 1px solid rgba(255, 255, 255, 100);
|
|||
|
|
border-radius: 5px;
|
|||
|
|
background-color: rgba(255, 255, 255, 50);
|
|||
|
|
}
|
|||
|
|
QProgressBar::chunk {
|
|||
|
|
background-color: white;
|
|||
|
|
border-radius: 4px;
|
|||
|
|
}
|
|||
|
|
""")
|
|||
|
|
layout.addWidget(self.progress_bar)
|
|||
|
|
|
|||
|
|
self.setLayout(layout)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
# YuNet 模型路径
|
|||
|
|
# ======================================================
|
|||
|
|
def get_yunet_model_path():
|
|||
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|||
|
|
model_path = os.path.join(base_dir, "model", "face_detection_yunet_2023mar.onnx")
|
|||
|
|
|
|||
|
|
if not os.path.exists(model_path):
|
|||
|
|
raise FileNotFoundError(f"YuNet 模型不存在: {model_path}")
|
|||
|
|
return model_path
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
# 主窗口
|
|||
|
|
# ======================================================
|
|||
|
|
class FaceRandomApp(QWidget):
|
|||
|
|
def __init__(self):
|
|||
|
|
super().__init__()
|
|||
|
|
|
|||
|
|
# 先显示加载页面
|
|||
|
|
self.loading_screen = LoadingScreen()
|
|||
|
|
self.loading_screen.show()
|
|||
|
|
|
|||
|
|
# 强制立即处理GUI事件，确保加载页面能立即显示
|
|||
|
|
QApplication.processEvents()
|
|||
|
|
|
|||
|
|
# 延迟初始化主界面，让加载页面先显示
|
|||
|
|
QTimer.singleShot(100, self.initialize_app)
|
|||
|
|
|
|||
|
|
def initialize_app(self):
|
|||
|
|
"""初始化应用程序"""
|
|||
|
|
# 状态
|
|||
|
|
self.state = "normal"
|
|||
|
|
self.selected_face_index = -1
|
|||
|
|
self.static_frame = None
|
|||
|
|
self.faces_snapshot = []
|
|||
|
|
|
|||
|
|
# 参数
|
|||
|
|
self.detection_confidence = 0.6
|
|||
|
|
self.faces = []
|
|||
|
|
|
|||
|
|
# 摄像头
|
|||
|
|
self.cap = cv2.VideoCapture(0)
|
|||
|
|
self.cap.set(3, 1280)
|
|||
|
|
self.cap.set(4, 720)
|
|||
|
|
|
|||
|
|
# YuNet 加载
|
|||
|
|
model_path = get_yunet_model_path()
|
|||
|
|
self.detector = cv2.FaceDetectorYN.create(
|
|||
|
|
model_path,
|
|||
|
|
"",
|
|||
|
|
(640, 480),
|
|||
|
|
score_threshold=self.detection_confidence,
|
|||
|
|
nms_threshold=0.3,
|
|||
|
|
top_k=5000
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# UI
|
|||
|
|
self.setup_ui()
|
|||
|
|
|
|||
|
|
# 关闭加载界面并显示主窗口
|
|||
|
|
self.loading_screen.close()
|
|||
|
|
self.show()
|
|||
|
|
|
|||
|
|
# Timer
|
|||
|
|
self.timer = QTimer()
|
|||
|
|
self.timer.timeout.connect(self.update_frame)
|
|||
|
|
self.timer.start(30)
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
# UI
|
|||
|
|
# ======================================================
|
|||
|
|
def setup_ui(self):
|
|||
|
|
self.setWindowTitle("Face Random Selector")
|
|||
|
|
self.resize(1280, 720)
|
|||
|
|
|
|||
|
|
self.video_label = QLabel(self)
|
|||
|
|
self.video_label.setGeometry(0, 0, self.width(), self.height())
|
|||
|
|
self.video_label.setStyleSheet("background: black;")
|
|||
|
|
|
|||
|
|
# 按钮 - 简约黑色半透明样式，无边框
|
|||
|
|
self.btn = QPushButton("随机", self)
|
|||
|
|
self.btn.setFixedSize(140, 55)
|
|||
|
|
self.btn.move(20, self.height() - 75)
|
|||
|
|
self.btn.clicked.connect(self.on_random_clicked)
|
|||
|
|
self.btn.setStyleSheet("""
|
|||
|
|
QPushButton {
|
|||
|
|
background-color: rgba(0, 0, 0, 180);
|
|||
|
|
color: white;
|
|||
|
|
font-size: 20px;
|
|||
|
|
font-weight: bold;
|
|||
|
|
border: none;
|
|||
|
|
border-radius: 8px;
|
|||
|
|
}
|
|||
|
|
QPushButton:hover {
|
|||
|
|
background-color: rgba(30, 30, 30, 200);
|
|||
|
|
}
|
|||
|
|
QPushButton:pressed {
|
|||
|
|
background-color: rgba(50, 50, 50, 220);
|
|||
|
|
}
|
|||
|
|
""")
|
|||
|
|
|
|||
|
|
# 置信度滑条 - 简约半透明样式
|
|||
|
|
self.slider = QSlider(Qt.Horizontal, self)
|
|||
|
|
self.slider.setRange(30, 90)
|
|||
|
|
self.slider.setValue(int(self.detection_confidence * 100))
|
|||
|
|
self.slider.setFixedWidth(200)
|
|||
|
|
self.slider.move(120, 20)
|
|||
|
|
self.slider.valueChanged.connect(self.on_confidence_change)
|
|||
|
|
self.slider.setStyleSheet("""
|
|||
|
|
QSlider::groove:horizontal {
|
|||
|
|
background: rgba(0, 0, 0, 120);
|
|||
|
|
height: 4px;
|
|||
|
|
border-radius: 2px;
|
|||
|
|
}
|
|||
|
|
QSlider::handle:horizontal {
|
|||
|
|
background: white;
|
|||
|
|
width: 14px;
|
|||
|
|
height: 14px;
|
|||
|
|
border-radius: 7px;
|
|||
|
|
margin: -5px 0;
|
|||
|
|
}
|
|||
|
|
""")
|
|||
|
|
|
|||
|
|
# 置信度标签 - 简约白色文字
|
|||
|
|
self.conf_label = QLabel("置信度", self)
|
|||
|
|
self.conf_label.move(20, 20)
|
|||
|
|
self.conf_label.setStyleSheet("""
|
|||
|
|
color: white;
|
|||
|
|
font-size: 16px;
|
|||
|
|
font-weight: bold;
|
|||
|
|
background: transparent;
|
|||
|
|
""")
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
def resizeEvent(self, event):
|
|||
|
|
self.video_label.setGeometry(0, 0, self.width(), self.height())
|
|||
|
|
self.btn.move(20, self.height() - 75)
|
|||
|
|
super().resizeEvent(event)
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
# 置信度
|
|||
|
|
# ======================================================
|
|||
|
|
def on_confidence_change(self, value):
|
|||
|
|
self.detection_confidence = value / 100.0
|
|||
|
|
self.detector.setScoreThreshold(self.detection_confidence)
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
# 随机按钮
|
|||
|
|
# ======================================================
|
|||
|
|
def on_random_clicked(self):
|
|||
|
|
if self.state == "normal":
|
|||
|
|
self.state = "random"
|
|||
|
|
self.btn.setText("重置")
|
|||
|
|
|
|||
|
|
if len(self.faces) > 0:
|
|||
|
|
self.selected_face_index = random.randint(0, len(self.faces) - 1)
|
|||
|
|
|
|||
|
|
# 捕获静态帧
|
|||
|
|
ret, img = self.cap.read()
|
|||
|
|
img = cv2.flip(img, 1)
|
|||
|
|
img = self.resize_cover(img, self.video_label.width(), self.video_label.height())
|
|||
|
|
|
|||
|
|
self.static_frame = img.copy()
|
|||
|
|
self.faces_snapshot = self.faces.copy()
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
self.state = "normal"
|
|||
|
|
self.btn.setText("随机")
|
|||
|
|
self.selected_face_index = -1
|
|||
|
|
self.static_frame = None
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
# 等比例覆盖（无黑边）
|
|||
|
|
# ======================================================
|
|||
|
|
def resize_cover(self, img, target_w, target_h):
|
|||
|
|
"""等比例缩放 + 居中裁切 = 无黑边"""
|
|||
|
|
h, w = img.shape[:2]
|
|||
|
|
|
|||
|
|
scale = max(target_w / w, target_h / h)
|
|||
|
|
new_w = int(w * scale)
|
|||
|
|
new_h = int(h * scale)
|
|||
|
|
|
|||
|
|
resized = cv2.resize(img, (new_w, new_h))
|
|||
|
|
|
|||
|
|
# 居中裁切
|
|||
|
|
x_start = (new_w - target_w) // 2
|
|||
|
|
y_start = (new_h - target_h) // 2
|
|||
|
|
|
|||
|
|
cropped = resized[y_start:y_start + target_h, x_start:x_start + target_w]
|
|||
|
|
return cropped
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
# 绘制人脸框和置信度 - 优化样式
|
|||
|
|
# ======================================================
|
|||
|
|
def draw_face_with_confidence(self, img, face, color, thickness=2):
|
|||
|
|
"""绘制人脸框和置信度"""
|
|||
|
|
x, y, w, h = map(int, face[:4])
|
|||
|
|
score = face[4] # 置信度分数，范围0-1
|
|||
|
|
|
|||
|
|
# 绘制人脸框
|
|||
|
|
cv2.rectangle(img, (x, y), (x + w, y + h), color, thickness)
|
|||
|
|
|
|||
|
|
# 绘制置信度文本
|
|||
|
|
confidence_text = f"{score:.2f}"
|
|||
|
|
text_y = max(y - 10, 20) # 确保文本不超出图像顶部
|
|||
|
|
|
|||
|
|
# 绘制半透明文本背景
|
|||
|
|
text_size = cv2.getTextSize(confidence_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
|
|||
|
|
|
|||
|
|
# 创建半透明覆盖层
|
|||
|
|
overlay = img.copy()
|
|||
|
|
cv2.rectangle(overlay, (x, text_y - text_size[1] - 5),
|
|||
|
|
(x + text_size[0] + 10, text_y + 5), color, -1)
|
|||
|
|
|
|||
|
|
# 应用半透明效果
|
|||
|
|
alpha = 0.7 # 透明度
|
|||
|
|
cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0, img)
|
|||
|
|
|
|||
|
|
# 绘制文本
|
|||
|
|
cv2.putText(img, confidence_text, (x + 5, text_y),
|
|||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
# 主循环
|
|||
|
|
# ======================================================
|
|||
|
|
def update_frame(self):
|
|||
|
|
ret, img = self.cap.read()
|
|||
|
|
if not ret:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
img = cv2.flip(img, 1)
|
|||
|
|
img = self.resize_cover(img, self.video_label.width(), self.video_label.height())
|
|||
|
|
|
|||
|
|
if self.state == "normal":
|
|||
|
|
# 设置 YuNet 输入大小
|
|||
|
|
self.detector.setInputSize((img.shape[1], img.shape[0]))
|
|||
|
|
|
|||
|
|
# 检测
|
|||
|
|
_, detected = self.detector.detect(img)
|
|||
|
|
self.faces = detected if detected is not None else []
|
|||
|
|
|
|||
|
|
# 画框和置信度
|
|||
|
|
for face in self.faces:
|
|||
|
|
self.draw_face_with_confidence(img, face, (0, 255, 0), 2)
|
|||
|
|
|
|||
|
|
else: # random 模式
|
|||
|
|
img = self.static_frame.copy()
|
|||
|
|
|
|||
|
|
for i, face in enumerate(self.faces_snapshot):
|
|||
|
|
color = (0, 0, 255) if i == self.selected_face_index else (0, 255, 0)
|
|||
|
|
thickness = 3 if i == self.selected_face_index else 2
|
|||
|
|
self.draw_face_with_confidence(img, face, color, thickness)
|
|||
|
|
|
|||
|
|
self.display(img)
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
def display(self, img):
|
|||
|
|
rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|||
|
|
h, w, c = rgb.shape
|
|||
|
|
qimg = QImage(rgb.data, w, h, c * w, QImage.Format_RGB888)
|
|||
|
|
self.video_label.setPixmap(QPixmap.fromImage(qimg))
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
def closeEvent(self, event):
|
|||
|
|
if self.cap.isOpened():
|
|||
|
|
self.cap.release()
|
|||
|
|
event.accept()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ======================================================
|
|||
|
|
def main():
|
|||
|
|
app = QApplication(sys.argv)
|
|||
|
|
window = FaceRandomApp()
|
|||
|
|
# 注意：这里不再调用 window.show()，因为 FaceRandomApp 内部会处理显示逻辑
|
|||
|
|
sys.exit(app.exec())
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|