318 lines
13 KiB
Python
318 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
屏幕录制管理器
|
||
仅保留FFmpeg录制功能
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import logging
|
||
import subprocess
|
||
import signal
|
||
import base64
|
||
from pathlib import Path
|
||
from datetime import datetime
|
||
from typing import Dict, Any, List, Optional
|
||
|
||
try:
|
||
import pyautogui
|
||
except ImportError:
|
||
pyautogui = None
|
||
|
||
try:
|
||
from .utils.config_manager import ConfigManager
|
||
except ImportError:
|
||
from utils.config_manager import ConfigManager
|
||
|
||
class RecordingManager:
|
||
def __init__(self, config_manager: Optional[ConfigManager] = None):
|
||
"""
|
||
初始化录制管理器
|
||
|
||
Args:
|
||
config_manager: 配置管理器实例
|
||
"""
|
||
self.logger = logging.getLogger(__name__)
|
||
self.config_manager = config_manager or ConfigManager()
|
||
|
||
# FFmpeg进程管理
|
||
self._ffmpeg_processes = {}
|
||
self._ffmpeg_meta = {}
|
||
|
||
# 默认参数
|
||
self.screen_fps = 25
|
||
self.screen_size = self._get_screen_size()
|
||
|
||
def _get_screen_size(self):
|
||
try:
|
||
import pyautogui
|
||
return pyautogui.size()
|
||
except ImportError:
|
||
return (1920, 1080)
|
||
|
||
def _get_primary_screen_bounds(self) -> Dict[str, int]:
|
||
try:
|
||
import ctypes
|
||
user32 = ctypes.windll.user32
|
||
w = user32.GetSystemMetrics(0)
|
||
h = user32.GetSystemMetrics(1)
|
||
return {'x': 0, 'y': 0, 'width': int(w), 'height': int(h)}
|
||
except Exception:
|
||
sw, sh = self.screen_size
|
||
return {'x': 0, 'y': 0, 'width': int(sw), 'height': int(sh)}
|
||
|
||
def start_recording_ffmpeg(self, session_id: str, patient_id: str, screen_location: List[int], fps: int = None) -> Dict[str, Any]:
|
||
result = {'success': False, 'message': ''}
|
||
try:
|
||
x, y, w, h = screen_location
|
||
bounds = self._get_primary_screen_bounds()
|
||
x_clamped = max(bounds['x'], min(int(x), bounds['x'] + bounds['width'] - 1))
|
||
y_clamped = max(bounds['y'], min(int(y), bounds['y'] + bounds['height'] - 1))
|
||
max_w = (bounds['x'] + bounds['width']) - x_clamped
|
||
max_h = (bounds['y'] + bounds['height']) - y_clamped
|
||
w_clamped = max(1, min(int(w), int(max_w)))
|
||
h_clamped = max(1, min(int(h), int(max_h)))
|
||
off_x = x_clamped - bounds['x']
|
||
off_y = y_clamped - bounds['y']
|
||
|
||
file_dir = self.config_manager.get_config_value('FILEPATH', 'path')
|
||
timestamp = datetime.now().strftime('%H%M%S%f')[:-3]
|
||
base_path = os.path.join(file_dir, patient_id, session_id, f'video_{timestamp}')
|
||
os.makedirs(base_path, exist_ok=True)
|
||
screen_video_path = os.path.join(base_path, 'screen.mp4')
|
||
target_fps = fps or self.screen_fps
|
||
ffmpeg_path = None
|
||
if self.config_manager:
|
||
ffmpeg_path = (
|
||
self.config_manager.get_config_value('SCREEN_RECORDING', 'ffmpeg_path', fallback=None) or
|
||
self.config_manager.get_config_value('RECORDING', 'ffmpeg_path', fallback=None)
|
||
)
|
||
if not ffmpeg_path or not os.path.isfile(str(ffmpeg_path)):
|
||
base_dir = os.path.dirname(sys.executable) if getattr(sys, 'frozen', False) else os.path.dirname(os.path.abspath(__file__))
|
||
alt_path = os.path.join(base_dir, 'ffmpeg', 'bin', 'ffmpeg.exe')
|
||
if os.path.isfile(alt_path):
|
||
ffmpeg_path = alt_path
|
||
else:
|
||
result['message'] = '未配置有效的ffmpeg_path，请在配置中设置 SCREEN_RECORDING.ffmpeg_path 或 RECORDING.ffmpeg_path'
|
||
return result
|
||
|
||
cmd = [
|
||
str(ffmpeg_path),
|
||
'-y',
|
||
'-f', 'gdigrab',
|
||
'-framerate', str(target_fps),
|
||
'-draw_mouse', str(int((self.config_manager.get_config_value('SCREEN_RECORDING', 'ffmpeg_draw_mouse', fallback='0') or '0'))),
|
||
'-offset_x', str(off_x),
|
||
'-offset_y', str(off_y),
|
||
'-video_size', f'{w_clamped}x{h_clamped}',
|
||
'-i', 'desktop',
|
||
]
|
||
|
||
codec = (self.config_manager.get_config_value('SCREEN_RECORDING', 'ffmpeg_codec', fallback=None) or
|
||
self.config_manager.get_config_value('RECORDING', 'ffmpeg_codec', fallback=None) or
|
||
'libx264')
|
||
preset = (self.config_manager.get_config_value('SCREEN_RECORDING', 'ffmpeg_preset', fallback=None) or
|
||
self.config_manager.get_config_value('RECORDING', 'ffmpeg_preset', fallback=None) or
|
||
('p1' if codec == 'h264_nvenc' else 'ultrafast'))
|
||
threads = int((self.config_manager.get_config_value('SCREEN_RECORDING', 'ffmpeg_threads', fallback='2') or '2'))
|
||
bframes = int((self.config_manager.get_config_value('SCREEN_RECORDING', 'ffmpeg_bframes', fallback='0') or '0'))
|
||
gop = int((self.config_manager.get_config_value('SCREEN_RECORDING', 'ffmpeg_gop', fallback=str(max(1, int(target_fps*2)))) or str(max(1, int(target_fps*2)))))
|
||
|
||
cmd += ['-c:v', codec]
|
||
# 统一的低CPU选项
|
||
cmd += ['-preset', str(preset)]
|
||
cmd += ['-bf', str(bframes)]
|
||
cmd += ['-g', str(gop)]
|
||
cmd += ['-pix_fmt', 'yuv420p']
|
||
cmd += ['-threads', str(threads)]
|
||
cmd += ['-r', str(target_fps)]
|
||
cmd += [screen_video_path]
|
||
proc = subprocess.Popen(
|
||
cmd,
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.DEVNULL,
|
||
stderr=subprocess.DEVNULL,
|
||
creationflags=getattr(subprocess, 'CREATE_NEW_PROCESS_GROUP', 0)
|
||
)
|
||
self._ffmpeg_processes['screen'] = proc
|
||
self._ffmpeg_meta['screen'] = {'base_path': base_path, 'patient_id': patient_id, 'session_id': session_id, 'video_path': screen_video_path}
|
||
result['success'] = True
|
||
result['message'] = 'ffmpeg录制已启动'
|
||
result['database_updates'] = {
|
||
'session_id': session_id,
|
||
'status': 'recording',
|
||
'video_paths': {
|
||
'screen_video_path': os.path.relpath(screen_video_path, file_dir)
|
||
}
|
||
}
|
||
return result
|
||
except Exception as e:
|
||
result['message'] = f'ffmpeg启动失败: {e}'
|
||
return result
|
||
|
||
def stop_recording_ffmpeg(self, session_id: str = None) -> Dict[str, Any]:
|
||
result = {'success': False, 'message': ''}
|
||
try:
|
||
proc = self._ffmpeg_processes.get('screen')
|
||
meta = self._ffmpeg_meta.get('screen')
|
||
if proc:
|
||
try:
|
||
if proc.stdin and proc.poll() is None:
|
||
try:
|
||
proc.communicate(input=b'q', timeout=2.0)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
proc.send_signal(getattr(signal, 'CTRL_BREAK_EVENT', signal.SIGTERM))
|
||
except Exception:
|
||
pass
|
||
try:
|
||
proc.terminate()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
proc.wait(timeout=3.0)
|
||
except Exception:
|
||
try:
|
||
proc.kill()
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
self._ffmpeg_processes.pop('screen', None)
|
||
result['success'] = True
|
||
result['message'] = 'ffmpeg录制已停止'
|
||
if meta:
|
||
result['database_updates'] = {
|
||
'session_id': meta.get('session_id'),
|
||
'status': 'recorded'
|
||
}
|
||
return result
|
||
except Exception as e:
|
||
result['message'] = f'ffmpeg停止失败: {e}'
|
||
return result
|
||
|
||
def save_detection_images(self, session_id: str, patient_id: str, detection_data: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""
|
||
保存前端传入的检测图片到指定目录
|
||
|
||
Args:
|
||
session_id: 检测会话ID
|
||
patient_id: 患者ID
|
||
detection_data: 前端传入的检测数据，包含base64格式的图片数据
|
||
|
||
Returns:
|
||
Dict: 包含所有采集数据的字典，符合detection_data表结构
|
||
"""
|
||
# 生成采集时间戳
|
||
timestamp = datetime.now().strftime('%H%M%S%f')[:-3] # 精确到毫秒
|
||
file_path = self.config_manager.get_config_value('FILEPATH', 'path')
|
||
data_dir = Path(os.path.join(file_path,patient_id, session_id, f"image_{timestamp}"))
|
||
# 创建数据存储目录
|
||
data_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 设置目录权限为777（完全权限）
|
||
try:
|
||
import stat
|
||
os.chmod(str(data_dir), stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO) # 777权限
|
||
self.logger.debug(f"已设置目录权限为777: {data_dir}")
|
||
except Exception as perm_error:
|
||
self.logger.warning(f"设置目录权限失败: {perm_error}，但目录创建成功")
|
||
|
||
# 初始化数据字典
|
||
data = {
|
||
'session_id': session_id,
|
||
'head_pose': detection_data.get('head_pose'),
|
||
'screen_location': detection_data.get('screen_location'),
|
||
'body_pose': None,
|
||
'body_image': None,
|
||
'foot_data': detection_data.get('foot_data'),
|
||
'foot_data_image': detection_data.get('foot_data_image'),
|
||
'foot1_image': None,
|
||
'foot2_image': None,
|
||
'screen_image': None,
|
||
'timestamp': timestamp
|
||
}
|
||
|
||
try:
|
||
# 保存图片数据
|
||
image_fields = [
|
||
('body_image', 'body'),
|
||
('foot1_image', 'foot1'),
|
||
('foot2_image', 'foot2')
|
||
]
|
||
|
||
for field, prefix in image_fields:
|
||
base64_data = detection_data.get(field)
|
||
if base64_data:
|
||
try:
|
||
# 移除base64头部信息
|
||
if ';base64,' in base64_data:
|
||
base64_data = base64_data.split(';base64,')[1]
|
||
|
||
# 解码base64数据
|
||
image_data = base64.b64decode(base64_data)
|
||
|
||
# 生成图片文件名
|
||
filename = f'{prefix}_{timestamp}.jpg'
|
||
file_path = data_dir / filename
|
||
|
||
# 保存图片
|
||
with open(file_path, 'wb') as f:
|
||
f.write(image_data)
|
||
|
||
# 更新数据字典中的图片路径
|
||
|
||
data[field] = str(os.path.join(patient_id, session_id, f"image_{timestamp}", filename))
|
||
self.logger.debug(f'{field}保存成功: {filename}')
|
||
|
||
except Exception as e:
|
||
self.logger.error(f'保存{field}失败: {e}')
|
||
# 完整屏幕截图--根据screen_location 进行截图
|
||
screen_image = self._capture_screen_image(data_dir, data.get('screen_location'),'screen', timestamp=timestamp)
|
||
if screen_image:
|
||
data['screen_image'] = str(os.path.join( patient_id, session_id, f"image_{timestamp}", screen_image))
|
||
# 足部压力屏幕截图——根据foot_data_image 进行截图
|
||
foot_data_image = self._capture_screen_image(data_dir, data.get('foot_data_image'),'foot_data', timestamp=timestamp)
|
||
if foot_data_image:
|
||
data['foot_data_image'] = str(os.path.join( patient_id, session_id, f"image_{timestamp}", foot_data_image))
|
||
|
||
self.logger.debug(f'数据保存完成: {session_id}, 时间戳: {timestamp}')
|
||
|
||
except Exception as e:
|
||
self.logger.error(f'数据保存失败: {e}')
|
||
|
||
return data
|
||
|
||
|
||
def _capture_screen_image(self, data_dir, screen_location,type, timestamp) -> Optional[str]:
|
||
"""
|
||
采集屏幕截图，根据screen_region 进行截图
|
||
|
||
Args:
|
||
data_dir: 数据存储目录路径
|
||
|
||
Returns:
|
||
str: 截图文件的相对路径，失败返回None
|
||
"""
|
||
try:
|
||
# 截取屏幕
|
||
if screen_location:
|
||
# 使用指定区域截图
|
||
x, y, width, height = screen_location
|
||
screenshot = pyautogui.screenshot(region=(x, y, width, height))
|
||
else:
|
||
# 全屏截图
|
||
screenshot = pyautogui.screenshot()
|
||
|
||
# 保存截图
|
||
from pathlib import Path
|
||
screen_filename = f'{type}_{timestamp}.jpg'
|
||
image_path = Path(data_dir) / screen_filename
|
||
screenshot.save(str(image_path), quality=95, optimize=True)
|
||
|
||
return screen_filename
|
||
|
||
except Exception as e:
|
||
self.logger.error(f'屏幕截图失败: {e}')
|
||
return None |