#!/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()