from flask import Flask, request, jsonify, send_file, stream_with_context, Response from werkzeug.utils import secure_filename import os import uuid from flask_cors import CORS import time from workflow import WorkflowResult, WorkflowResultStream from utils import insert_into_mysql, setup_logging, json_format, json_to_excel import pandas as pd import json logger = setup_logging() app = Flask(__name__) CORS(app) # 允许跨域请求 # 配置上传参数 UPLOAD_FOLDER = 'uploads' ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'bmp'} MAX_FILE_SIZE = 10 * 1024 * 1024 # 10MB app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER app.config['MAX_CONTENT_LENGTH'] = MAX_FILE_SIZE def allowed_file(filename): return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS @app.route('/') def index(): return app.send_static_file('imageOCR.html') @app.route('/upload', methods=['POST']) def upload_file(): if 'file' not in request.files: return jsonify({'status': 'error', 'message': 'No file uploaded'}), 400 file = request.files['file'] # 检查文件名是否合法 if file.filename == '': return jsonify({'status': 'error', 'message': 'Empty filename'}), 400 # 验证文件类型和大小 if not allowed_file(file.filename): return jsonify({'status': 'error', 'message': 'Invalid file type'}), 400 try: # 安全保存文件 filename = secure_filename(file.filename) if not os.path.exists(app.config['UPLOAD_FOLDER']): os.makedirs(app.config['UPLOAD_FOLDER']) filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) file.save(filepath) # 文件处理逻辑 file_id = str(uuid.uuid4()) ocr_workflow = WorkflowResult(local_file_path=filepath, cos_path_prefix=f"images/ruiwen", file_id=file_id, filename=filename, json_output_path=os.path.join(app.config['UPLOAD_FOLDER'], f"{file_id}_result.json"), excel_output_path=os.path.join(app.config['UPLOAD_FOLDER'], f"{file_id}_result.xlsx")) json_res = ocr_workflow.run_single_workflow() # 插入数据库 logger.info("开始插入数据库记录") db_params = ocr_workflow.db_insert_params() insert_into_mysql(**db_params) logger.info("数据库记录插入完成") return jsonify({ 'status': 'success', 'message': 'File uploaded successfully', 'file_id': file_id, 'json': json_res }), 200 except Exception as e: return jsonify({'status': 'error', 'message': str(e)}), 500 @app.route('/upload_stream', methods=['POST']) def upload_file_stream(): if 'file' not in request.files: return jsonify({'status': 'error', 'message': 'No file uploaded'}), 400 file = request.files['file'] if file.filename == '': return jsonify({'status': 'error', 'message': 'Empty filename'}), 400 if not allowed_file(file.filename): return jsonify({'status': 'error', 'message': 'Invalid file type'}), 400 try: # 安全保存文件 filename = secure_filename(file.filename) upload_folder = app.config['UPLOAD_FOLDER'] if not os.path.exists(upload_folder): os.makedirs(upload_folder) filepath = os.path.join(upload_folder, filename) file.save(filepath) file_id = str(uuid.uuid4()) start_time = time.time() json_output_path = os.path.join(upload_folder, f"{file_id}_result.json") excel_output_path = os.path.join(upload_folder, f"{file_id}_result.xlsx") @stream_with_context def generate(): try: # 执行OCR工作流，获取流式结果 ocr_workflow_stream = WorkflowResultStream( local_file_path=filepath, cos_path_prefix="images/ruiwen", file_id=file_id, filename=filename, json_output_path=json_output_path, excel_output_path=excel_output_path ) stream_ocr_result = ocr_workflow_stream.run_single_workflow() input_tokens = 0 output_tokens = 0 content_parts = [] table_data = None # 用于存储表格数据 if hasattr(stream_ocr_result, '__iter__'): logger.info("开始流式返回OCR结果") for chunk in stream_ocr_result: if isinstance(chunk, tuple): # Token统计信息 in_tokens, out_tokens = chunk input_tokens += in_tokens output_tokens += out_tokens else: content_str = str(chunk) content_parts.append(content_str) yield f"event: message\ndata: {json.dumps({'content': content_str}, ensure_ascii=False)}\n\n" else: logger.info("返回非流式OCR结果") content_str = str(stream_ocr_result) content_parts.append(content_str) yield f"event: message\ndata: {json.dumps({'content': content_str}, ensure_ascii=False)}\n\n" try: logger.info("开始插入数据库记录") ocr_result = "".join(content_parts) logger.info("开始保存Excel文件, 路径: %s", excel_output_path) json_res = json_format(ocr_result) data_detail = json_to_excel(json_res, excel_output_path) yield f"event: table\ndata: {json.dumps(data_detail, ensure_ascii=False)}\n\n" db_params = ocr_workflow_stream.db_insert_params() db_params.update({ 'ocr_cost_time': int(time.time() - start_time), 'input_tokens': input_tokens, 'output_tokens': output_tokens, 'ocr_json_result': ocr_result }) insert_into_mysql(**db_params) logger.info("数据库记录插入完成") # 🔥 关键修复：添加完成事件通知 yield f"event: completed\ndata: {json.dumps({'file_id': file_id, 'status': 'success'})}\n\n" logger.info(f"SSE流处理完成，file_id: {file_id}") except Exception as db_error: logger.error(f"数据库插入失败: {str(db_error)}") # 发送错误事件 yield f"event: error\ndata: {json.dumps({'error': f'数据库操作失败: {str(db_error)}'})}\n\n" except Exception as e: logger.error(f"流式处理发生错误: {str(e)}", exc_info=True) yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n" return Response( generate(), mimetype='text/event-stream', headers={ 'Cache-Control': 'no-cache', 'Connection': 'keep-alive', 'Access-Control-Allow-Origin': '*', 'Access-Control-Allow-Headers': 'Content-Type' } ) except Exception as e: logger.error(f"文件处理失败: {str(e)}") return jsonify({'status': 'error', 'message': str(e)}), 500 @app.route('/export_file/') def export_file(file_id): """处理Excel文件下载""" try: # 验证文件合法性 if not valid_file_id(file_id): return jsonify({'error': 'Invalid file ID'}), 400 # 构建完整的Excel文件路径 excel_file = os.path.join(app.config['UPLOAD_FOLDER'], f"{file_id}_result.xlsx") if not os.path.exists(excel_file): return jsonify({'error': 'File expired or not found'}), 404 # 发送文件：使用正确的MIME类型和下载文件名 response = send_file( excel_file, mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', # 修正MIME类型 as_attachment=True, download_name='data_result.xlsx' # 确保Flask版本支持此参数 ) return response except Exception as e: # 记录错误日志到服务器控制台，便于调试 app.logger.error(f"Error exporting file {file_id}: {str(e)}") return jsonify({'error': 'An internal server error occurred during file export.'}), 500 def valid_file_id(file_id): """验证文件ID合法性""" try: uuid.UUID(file_id) return True except ValueError: return False def cleanup_file(path): """安全清理文件""" try: if os.path.exists(path): os.remove(path) except Exception as e: print(f"文件清理失败: {str(e)}") if __name__ == '__main__': app.run(host='0.0.0.0', port=5000, debug=True)