317 lines
10 KiB
Python
317 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Screen2Feishu - AI飞书多维表格自动录入助手
|
||
监控文件夹，自动处理图片并写入飞书多维表格
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import time
|
||
import argparse
|
||
from pathlib import Path
|
||
from watchdog.observers import Observer
|
||
from watchdog.events import FileSystemEventHandler
|
||
|
||
# 添加项目根目录到Python路径
|
||
project_root = Path(__file__).parent
|
||
sys.path.insert(0, str(project_root))
|
||
|
||
from services.ai_service import AIService
|
||
from services.feishu_service import FeishuService
|
||
from utils.logger import setup_logger
|
||
from utils.notifier import send_notification
|
||
from utils.config_loader import load_config, validate_config
|
||
|
||
|
||
class ImageFileHandler(FileSystemEventHandler):
|
||
"""文件系统事件处理器"""
|
||
|
||
def __init__(self, config, logger):
|
||
self.config = config
|
||
self.logger = logger
|
||
self.ai_service = AIService(config['ai'])
|
||
self.feishu_service = FeishuService(config['feishu'])
|
||
self.processed_files = set() # 已处理文件记录
|
||
self.supported_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.gif'}
|
||
|
||
# 创建必要的文件夹
|
||
self._create_directories()
|
||
|
||
def _create_directories(self):
|
||
"""创建必要的文件夹"""
|
||
watch_folder = Path(self.config['system']['watch_folder'])
|
||
watch_folder.mkdir(exist_ok=True)
|
||
|
||
if self.config['system'].get('post_process') == 'move':
|
||
processed_folder = Path(self.config['system']['processed_folder'])
|
||
processed_folder.mkdir(exist_ok=True)
|
||
|
||
def _is_supported_image(self, file_path):
|
||
"""检查文件是否为支持的图片格式"""
|
||
return Path(file_path).suffix.lower() in self.supported_extensions
|
||
|
||
def _is_file_ready(self, file_path):
|
||
"""检查文件是否已完全写入（通过文件大小稳定判断）"""
|
||
try:
|
||
size1 = os.path.getsize(file_path)
|
||
time.sleep(0.5) # 等待0.5秒
|
||
size2 = os.path.getsize(file_path)
|
||
return size1 == size2 and size1 > 0
|
||
except (OSError, IOError):
|
||
return False
|
||
|
||
def _process_image(self, file_path):
|
||
"""处理单个图片文件"""
|
||
file_path = Path(file_path)
|
||
|
||
# 检查是否已处理
|
||
if str(file_path) in self.processed_files:
|
||
return
|
||
|
||
self.logger.info(f"开始处理图片: {file_path.name}")
|
||
|
||
ai_result = None # 初始化变量
|
||
|
||
try:
|
||
# 检查是否启用内存处理
|
||
memory_processing = self.config['system'].get('memory_processing', False)
|
||
|
||
if memory_processing:
|
||
# 内存处理模式：直接读取图片到内存，不保存到processed_images
|
||
try:
|
||
with open(file_path, 'rb') as f:
|
||
image_bytes = f.read()
|
||
|
||
# 1. AI分析图片（内存模式）
|
||
ai_result = self.ai_service.analyze_image_from_bytes(image_bytes, file_path.name)
|
||
if not ai_result:
|
||
self.logger.error(f"AI分析失败: {file_path.name}")
|
||
return
|
||
except Exception as e:
|
||
self.logger.error(f"内存处理图片失败: {str(e)}")
|
||
# 回退到文件处理模式
|
||
memory_processing = False
|
||
|
||
if not memory_processing:
|
||
# 文件处理模式：传统方式
|
||
# 1. AI分析图片
|
||
ai_result = self.ai_service.analyze_image(str(file_path))
|
||
if not ai_result:
|
||
self.logger.error(f"AI分析失败: {file_path.name}")
|
||
return
|
||
|
||
# 2. 写入飞书
|
||
success = self.feishu_service.add_task(ai_result)
|
||
if success:
|
||
self.logger.info(f"成功写入飞书: {file_path.name}")
|
||
self.processed_files.add(str(file_path))
|
||
|
||
# 3. 后处理
|
||
self._post_process_file(file_path)
|
||
|
||
# 4. 发送通知
|
||
send_notification(
|
||
"任务处理成功",
|
||
f"已成功处理 {file_path.name} 并写入飞书表格"
|
||
)
|
||
else:
|
||
self.logger.error(f"写入飞书失败: {file_path.name}")
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"处理图片 {file_path.name} 时出错: {str(e)}")
|
||
|
||
def _post_process_file(self, file_path):
|
||
"""文件后处理"""
|
||
post_process = self.config['system'].get('post_process', 'keep')
|
||
|
||
if post_process == 'delete':
|
||
try:
|
||
file_path.unlink()
|
||
self.logger.info(f"已删除文件: {file_path.name}")
|
||
except Exception as e:
|
||
self.logger.error(f"删除文件失败: {str(e)}")
|
||
|
||
elif post_process == 'move':
|
||
try:
|
||
processed_folder = Path(self.config['system']['processed_folder'])
|
||
target_path = processed_folder / file_path.name
|
||
|
||
# 如果目标文件已存在，添加时间戳
|
||
if target_path.exists():
|
||
timestamp = int(time.time())
|
||
target_path = processed_folder / f"{file_path.stem}_{timestamp}{file_path.suffix}"
|
||
|
||
file_path.rename(target_path)
|
||
self.logger.info(f"已移动文件到: {target_path}")
|
||
except Exception as e:
|
||
self.logger.error(f"移动文件失败: {str(e)}")
|
||
|
||
def on_created(self, event):
|
||
"""文件创建事件处理"""
|
||
if event.is_directory:
|
||
return
|
||
|
||
file_path = event.src_path
|
||
if not self._is_supported_image(file_path):
|
||
return
|
||
|
||
self.logger.info(f"检测到新文件: {file_path}")
|
||
|
||
# 等待文件完全写入
|
||
if self._is_file_ready(file_path):
|
||
self._process_image(file_path)
|
||
else:
|
||
self.logger.warning(f"文件未完全写入，跳过处理: {file_path}")
|
||
|
||
def on_modified(self, event):
|
||
"""文件修改事件处理"""
|
||
if event.is_directory:
|
||
return
|
||
|
||
file_path = event.src_path
|
||
if not self._is_supported_image(file_path):
|
||
return
|
||
|
||
# 对于修改事件，也检查文件是否已完全写入
|
||
if self._is_file_ready(file_path):
|
||
self._process_image(file_path)
|
||
|
||
|
||
def parse_arguments():
|
||
"""解析命令行参数"""
|
||
parser = argparse.ArgumentParser(
|
||
description="AI飞书多维表格自动录入助手",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
示例用法:
|
||
python main.py # 使用默认配置启动
|
||
python main.py --config custom.yaml # 使用自定义配置文件
|
||
python main.py --test # 测试模式（不监控文件夹）
|
||
"""
|
||
)
|
||
|
||
parser.add_argument(
|
||
'--config', '-c',
|
||
default='config.yaml',
|
||
help='配置文件路径 (默认: config.yaml)'
|
||
)
|
||
|
||
parser.add_argument(
|
||
'--test', '-t',
|
||
action='store_true',
|
||
help='测试模式：处理现有文件后退出，不监控文件夹'
|
||
)
|
||
|
||
parser.add_argument(
|
||
'--verbose', '-v',
|
||
action='store_true',
|
||
help='显示详细日志'
|
||
)
|
||
|
||
return parser.parse_args()
|
||
|
||
|
||
def main():
|
||
"""主函数"""
|
||
args = parse_arguments()
|
||
|
||
# 加载配置
|
||
try:
|
||
config = load_config(args.config)
|
||
except Exception as e:
|
||
print(f" 配置加载失败: {e}")
|
||
sys.exit(1)
|
||
|
||
# 验证配置
|
||
try:
|
||
validate_config(config)
|
||
except Exception as e:
|
||
print(f" 配置验证失败: {e}")
|
||
sys.exit(1)
|
||
|
||
# 设置日志
|
||
log_level = 'DEBUG' if args.verbose else 'INFO'
|
||
logger = setup_logger(
|
||
name='screen2feishu',
|
||
log_file='app.log',
|
||
level=log_level
|
||
)
|
||
|
||
logger.info("=" * 50)
|
||
logger.info("Screen2Feishu 启动")
|
||
logger.info(f"配置文件: {args.config}")
|
||
logger.info(f"监控文件夹: {config['system']['watch_folder']}")
|
||
logger.info("=" * 50)
|
||
|
||
# 创建文件处理器
|
||
handler = ImageFileHandler(config, logger)
|
||
|
||
if args.test:
|
||
# 测试模式：处理现有文件
|
||
watch_folder = Path(config['system']['watch_folder'])
|
||
image_files = [
|
||
f for f in watch_folder.iterdir()
|
||
if f.is_file() and handler._is_supported_image(f)
|
||
]
|
||
|
||
if not image_files:
|
||
logger.info("没有找到待处理的图片文件")
|
||
return
|
||
|
||
logger.info(f"找到 {len(image_files)} 个图片文件，开始处理...")
|
||
for image_file in image_files:
|
||
handler._process_image(image_file)
|
||
|
||
logger.info("测试模式处理完成")
|
||
return
|
||
|
||
# 监控模式
|
||
observer = Observer()
|
||
observer.schedule(
|
||
handler,
|
||
path=config['system']['watch_folder'],
|
||
recursive=False
|
||
)
|
||
|
||
try:
|
||
observer.start()
|
||
logger.info("文件监控已启动，按 Ctrl+C 停止程序")
|
||
|
||
# 发送启动通知
|
||
send_notification(
|
||
"Screen2Feishu 已启动",
|
||
f"正在监控文件夹: {config['system']['watch_folder']}"
|
||
)
|
||
|
||
while True:
|
||
time.sleep(1)
|
||
|
||
except KeyboardInterrupt:
|
||
logger.info("收到停止信号，正在关闭...")
|
||
observer.stop()
|
||
|
||
# 发送停止通知
|
||
send_notification(
|
||
"Screen2Feishu 已停止",
|
||
"程序已正常关闭"
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"程序运行出错: {str(e)}")
|
||
observer.stop()
|
||
|
||
# 发送错误通知
|
||
send_notification(
|
||
"Screen2Feishu 错误",
|
||
f"程序运行出错: {str(e)}"
|
||
)
|
||
sys.exit(1)
|
||
|
||
finally:
|
||
observer.join()
|
||
logger.info("程序已完全停止")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |