401 lines
14 KiB
Python
401 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
实时手部检测Web服务器
|
||
支持WebSocket通信，实时视频流处理和机械臂控制
|
||
"""
|
||
|
||
import asyncio
|
||
import base64
|
||
import json
|
||
import logging
|
||
import os
|
||
import time
|
||
from threading import Thread
|
||
from typing import Dict, Optional, Any
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from flask import Flask, render_template, request
|
||
from flask_socketio import SocketIO, emit
|
||
from PIL import Image
|
||
import io
|
||
|
||
from hand_detection_3d import load_mediapipe_model, process_frame_3d
|
||
|
||
# 配置日志
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class HandDetectionWebServer:
|
||
"""实时手部检测Web服务器"""
|
||
|
||
def __init__(self, host='0.0.0.0', port=5000):
|
||
self.host = host
|
||
self.port = port
|
||
|
||
# Flask应用和SocketIO
|
||
# 设置模板和静态文件路径
|
||
import os
|
||
template_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'templates')
|
||
self.app = Flask(__name__, template_folder=template_dir)
|
||
self.app.config['SECRET_KEY'] = 'hand_detection_secret_key'
|
||
self.socketio = SocketIO(self.app, cors_allowed_origins="*", async_mode='threading')
|
||
|
||
# MediaPipe模型
|
||
self.hands_model = load_mediapipe_model()
|
||
|
||
# 状态管理
|
||
self.clients = {} # 连接的客户端
|
||
self.current_frame = None
|
||
self.previous_hand_data = None
|
||
self.detection_results = {
|
||
'x_angle': 90,
|
||
'y_angle': 90,
|
||
'z_angle': 90,
|
||
'grip': 0,
|
||
'action': 'none',
|
||
'speed': 5,
|
||
'timestamp': time.time()
|
||
}
|
||
|
||
# 性能监控
|
||
self.fps_counter = 0
|
||
self.last_fps_time = time.time()
|
||
self.current_fps = 0
|
||
|
||
# 配置路由和事件处理
|
||
self._setup_routes()
|
||
self._setup_socket_events()
|
||
|
||
def _setup_routes(self):
|
||
"""设置HTTP路由"""
|
||
|
||
@self.app.route('/')
|
||
def index():
|
||
"""主页面"""
|
||
return render_template('index.html')
|
||
|
||
@self.app.route('/api/status')
|
||
def api_status():
|
||
"""获取系统状态"""
|
||
return {
|
||
'status': 'running',
|
||
'fps': self.current_fps,
|
||
'clients': len(self.clients),
|
||
'detection_results': self.detection_results
|
||
}
|
||
|
||
def _setup_socket_events(self):
|
||
"""设置WebSocket事件处理"""
|
||
|
||
@self.socketio.on('connect')
|
||
def handle_connect():
|
||
"""客户端连接"""
|
||
client_id = request.sid
|
||
self.clients[client_id] = {
|
||
'connected_at': time.time(),
|
||
'type': 'unknown',
|
||
'last_ping': time.time()
|
||
}
|
||
logger.info(f"客户端 {client_id} 已连接")
|
||
|
||
# 发送欢迎消息和当前状态
|
||
emit('status', {
|
||
'message': '连接成功',
|
||
'client_id': client_id,
|
||
'current_results': self.detection_results
|
||
})
|
||
|
||
@self.socketio.on('disconnect')
|
||
def handle_disconnect():
|
||
"""客户端断开连接"""
|
||
client_id = request.sid
|
||
if client_id in self.clients:
|
||
del self.clients[client_id]
|
||
logger.info(f"客户端 {client_id} 已断开连接")
|
||
|
||
@self.socketio.on('register_client')
|
||
def handle_register_client(data):
|
||
"""注册客户端类型"""
|
||
client_id = request.sid
|
||
client_type = data.get('type', 'unknown')
|
||
|
||
if client_id in self.clients:
|
||
self.clients[client_id]['type'] = client_type
|
||
logger.info(f"客户端 {client_id} 注册为: {client_type}")
|
||
|
||
emit('registration_success', {
|
||
'client_id': client_id,
|
||
'type': client_type
|
||
})
|
||
|
||
@self.socketio.on('video_frame')
|
||
def handle_video_frame(data):
|
||
"""处理视频帧"""
|
||
try:
|
||
# 解码base64图像
|
||
frame_data = base64.b64decode(data['frame'])
|
||
frame = self._decode_frame(frame_data)
|
||
|
||
if frame is not None:
|
||
# 处理帧并检测手部
|
||
control_signal, hand_data = process_frame_3d(
|
||
frame, self.hands_model, self.previous_hand_data
|
||
)
|
||
|
||
# 更新检测结果
|
||
self.detection_results = control_signal
|
||
self.detection_results['timestamp'] = time.time()
|
||
self.previous_hand_data = hand_data
|
||
|
||
# 编码处理后的帧
|
||
processed_frame_data = self._encode_frame(frame)
|
||
|
||
# 发送结果给web预览客户端
|
||
self.socketio.emit('detection_results', {
|
||
'control_signal': control_signal,
|
||
'processed_frame': processed_frame_data,
|
||
'fps': self.current_fps
|
||
}, room=None)
|
||
|
||
# 发送控制信号给机械臂客户端
|
||
self._send_to_robot_clients(control_signal)
|
||
|
||
# 更新FPS
|
||
self._update_fps()
|
||
|
||
except Exception as e:
|
||
logger.error(f"处理视频帧时出错: {e}")
|
||
emit('error', {'message': str(e)})
|
||
|
||
@self.socketio.on('ping')
|
||
def handle_ping():
|
||
"""处理ping请求"""
|
||
client_id = request.sid
|
||
if client_id in self.clients:
|
||
self.clients[client_id]['last_ping'] = time.time()
|
||
emit('pong', {'timestamp': time.time()})
|
||
|
||
@self.socketio.on('get_detection_results')
|
||
def handle_get_detection_results():
|
||
"""获取最新的检测结果"""
|
||
emit('detection_results', {
|
||
'control_signal': self.detection_results,
|
||
'fps': self.current_fps
|
||
})
|
||
|
||
@self.socketio.on('start_local_test')
|
||
def handle_start_local_test(data=None):
|
||
"""处理开始本地测试请求"""
|
||
try:
|
||
# 如果提供了视频路径，使用指定的视频
|
||
if data and 'video_path' in data:
|
||
test_video = data['video_path']
|
||
if not os.path.exists(test_video):
|
||
emit('test_error', {
|
||
'message': f'视频文件不存在: {test_video}'
|
||
})
|
||
return
|
||
else:
|
||
# 检查是否有默认测试视频
|
||
test_videos = [
|
||
'data/videos/test_basic.mp4',
|
||
'data/videos/test_gesture.mp4'
|
||
]
|
||
|
||
# 找到第一个存在的测试视频
|
||
test_video = None
|
||
for video_path in test_videos:
|
||
if os.path.exists(video_path):
|
||
test_video = video_path
|
||
break
|
||
|
||
if not test_video:
|
||
# 没有找到测试视频，提供帮助信息
|
||
emit('test_error', {
|
||
'message': '未找到测试视频文件',
|
||
'help': '请先生成测试视频：python create_test_video.py'
|
||
})
|
||
return
|
||
|
||
logger.info(f"开始本地测试，使用视频: {test_video}")
|
||
self.start_local_video_test(test_video)
|
||
|
||
emit('test_started', {
|
||
'message': f'本地测试已开始，使用视频: {os.path.basename(test_video)}',
|
||
'video_path': test_video
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"启动本地测试时出错: {e}")
|
||
emit('test_error', {
|
||
'message': f'启动本地测试失败: {str(e)}'
|
||
})
|
||
|
||
@self.socketio.on('get_video_list')
|
||
def handle_get_video_list():
|
||
"""获取可用的视频文件列表"""
|
||
try:
|
||
video_dirs = ['data/videos', 'videos', '.']
|
||
video_extensions = ['.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv']
|
||
videos = []
|
||
|
||
for video_dir in video_dirs:
|
||
if os.path.exists(video_dir):
|
||
for file in os.listdir(video_dir):
|
||
if any(file.lower().endswith(ext) for ext in video_extensions):
|
||
file_path = os.path.join(video_dir, file)
|
||
try:
|
||
file_size = os.path.getsize(file_path)
|
||
size_mb = file_size / (1024 * 1024)
|
||
videos.append({
|
||
'path': file_path,
|
||
'name': file,
|
||
'size': f'{size_mb:.1f}MB'
|
||
})
|
||
except OSError:
|
||
continue
|
||
|
||
# 按文件名排序
|
||
videos.sort(key=lambda x: x['name'])
|
||
|
||
emit('video_list', {
|
||
'videos': videos
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取视频列表时出错: {e}")
|
||
emit('video_list', {
|
||
'videos': []
|
||
})
|
||
|
||
def _decode_frame(self, frame_data: bytes) -> Optional[np.ndarray]:
|
||
"""解码图像帧"""
|
||
try:
|
||
# 使用PIL解码
|
||
image = Image.open(io.BytesIO(frame_data))
|
||
frame = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
|
||
return frame
|
||
except Exception as e:
|
||
logger.error(f"解码帧时出错: {e}")
|
||
return None
|
||
|
||
def _encode_frame(self, frame: np.ndarray) -> str:
|
||
"""编码图像帧为base64"""
|
||
try:
|
||
# 转换为RGB格式
|
||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
image = Image.fromarray(frame_rgb)
|
||
|
||
# 编码为JPEG
|
||
buffer = io.BytesIO()
|
||
image.save(buffer, format='JPEG', quality=80)
|
||
frame_data = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
||
|
||
return f"data:image/jpeg;base64,{frame_data}"
|
||
except Exception as e:
|
||
logger.error(f"编码帧时出错: {e}")
|
||
return ""
|
||
|
||
def _send_to_robot_clients(self, control_signal: Dict[str, Any]):
|
||
"""发送控制信号给机械臂客户端"""
|
||
robot_clients = [
|
||
client_id for client_id, info in self.clients.items()
|
||
if info.get('type') == 'robot'
|
||
]
|
||
|
||
if robot_clients:
|
||
for client_id in robot_clients:
|
||
self.socketio.emit('robot_control', control_signal, room=client_id)
|
||
|
||
def _update_fps(self):
|
||
"""更新FPS计数"""
|
||
self.fps_counter += 1
|
||
current_time = time.time()
|
||
|
||
if current_time - self.last_fps_time >= 1.0: # 每秒更新一次
|
||
self.current_fps = self.fps_counter
|
||
self.fps_counter = 0
|
||
self.last_fps_time = current_time
|
||
|
||
def start_local_video_test(self, video_path: str):
|
||
"""启动本地视频测试"""
|
||
def video_test_thread():
|
||
cap = cv2.VideoCapture(video_path)
|
||
|
||
while cap.isOpened():
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
break
|
||
|
||
# 处理帧
|
||
control_signal, hand_data = process_frame_3d(
|
||
frame, self.hands_model, self.previous_hand_data
|
||
)
|
||
|
||
# 更新状态
|
||
self.detection_results = control_signal
|
||
self.detection_results['timestamp'] = time.time()
|
||
self.previous_hand_data = hand_data
|
||
|
||
# 编码帧
|
||
processed_frame_data = self._encode_frame(frame)
|
||
|
||
# 广播结果
|
||
self.socketio.emit('detection_results', {
|
||
'control_signal': control_signal,
|
||
'processed_frame': processed_frame_data,
|
||
'fps': self.current_fps
|
||
})
|
||
|
||
# 发送给机械臂
|
||
self._send_to_robot_clients(control_signal)
|
||
|
||
# 更新FPS
|
||
self._update_fps()
|
||
|
||
# 控制帧率
|
||
time.sleep(1/30) # 30 FPS
|
||
|
||
cap.release()
|
||
|
||
thread = Thread(target=video_test_thread)
|
||
thread.daemon = True
|
||
thread.start()
|
||
logger.info(f"本地视频测试已启动: {video_path}")
|
||
|
||
def run(self, debug=False):
|
||
"""启动Web服务器"""
|
||
logger.info(f"启动手部检测Web服务器 http://{self.host}:{self.port}")
|
||
self.socketio.run(
|
||
self.app,
|
||
host=self.host,
|
||
port=self.port,
|
||
debug=debug,
|
||
allow_unsafe_werkzeug=True
|
||
)
|
||
|
||
def main():
|
||
"""主函数"""
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description='实时手部检测Web服务器')
|
||
parser.add_argument('--host', default='0.0.0.0', help='服务器地址')
|
||
parser.add_argument('--port', type=int, default=5000, help='端口号')
|
||
parser.add_argument('--debug', action='store_true', help='调试模式')
|
||
parser.add_argument('--test-video', help='本地测试视频路径')
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 创建服务器实例
|
||
server = HandDetectionWebServer(host=args.host, port=args.port)
|
||
|
||
# 如果指定了测试视频，启动本地视频测试
|
||
if args.test_video:
|
||
server.start_local_video_test(args.test_video)
|
||
|
||
# 启动服务器
|
||
server.run(debug=args.debug)
|
||
|
||
if __name__ == '__main__':
|
||
main() |