import cv2 import numpy as np from flask import Flask, render_template, request, jsonify, send_file from flask_socketio import SocketIO, emit import io import base64 import time import threading import logging import os import json import sqlite3 from datetime import datetime import zipfile import tempfile # 导入视觉处理相关的模块 from llm_req import VisionAPIClient, DetectionResult from cap_trigger import ImageClient # --- 配置 --- SAVE_PATH_LEFT = "./static/received/left" SAVE_PATH_RIGHT = "./static/received/right" SAVE_PATH_LEFT_MARKED = "./static/received/left_marked" # 标注图片保存路径 SAVE_PATH_RIGHT_MARKED = "./static/received/right_marked" # 标注图片保存路径 FLASK_HOST = "0.0.0.0" FLASK_PORT = 5000 MAX_LIVE_FRAMES = 2 # 保留最新的几帧用于实时显示 DATABASE_PATH = "received_images.db" # SQLite 数据库文件路径 # --- 配置 --- # 设置日志 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # --- 数据库初始化 --- def init_db(): """初始化 SQLite 数据库和表""" conn = sqlite3.connect(DATABASE_PATH) cursor = conn.cursor() # 创建表，存储图片信息，添加标注图片字段 cursor.execute(''' CREATE TABLE IF NOT EXISTS images ( id INTEGER PRIMARY KEY AUTOINCREMENT, left_filename TEXT NOT NULL, right_filename TEXT NOT NULL, left_marked_filename TEXT, right_marked_filename TEXT, timestamp REAL NOT NULL, metadata TEXT, comment TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) ''') conn.commit() conn.close() logger.info(f"Database {DATABASE_PATH} initialized.") # --- 全局变量 --- # 用于存储最新的左右帧，用于实时显示 latest_left_frame = None latest_right_frame = None latest_timestamp = None frame_lock = threading.Lock() # 保护全局帧变量 # --- Flask & SocketIO 应用 --- app = Flask(__name__) # 为生产环境配置 SECRET_KEY app.config['SECRET_KEY'] = 'your-secret-key-change-this' # 配置异步模式，如果需要异步处理可以调整 socketio = SocketIO(app, cors_allowed_origins="*") # 允许所有来源，生产环境请具体配置 # 初始化图像客户端 image_client = ImageClient("tcp://127.0.0.1:54321", client_id="local") # 初始化数据库 init_db() # --- 辅助函数 --- def draw_detections_on_image(image: np.ndarray, detections: list) -> np.ndarray: """在图像上绘制检测框""" # 复制原图以避免修改原始图像 marked_image = image.copy() # 定义颜色映射 color_map = { 1: (0, 255, 0), # 绿色 - 弹药箱 2: (255, 0, 0), # 蓝色 - 士兵 3: (0, 0, 255), # 红色 - 枪支 4: (255, 255, 0) # 青色 - 数字牌 } # 获取图像尺寸 h, w = image.shape[:2] # 绘制每个检测框 for detection in detections: # 获取检测信息 obj_id = detection.get("id", 0) label = detection.get("label", "") bbox = detection.get("bbox", []) if len(bbox) == 4: # 将归一化的坐标转换为实际像素坐标 x_min = int(bbox[0] * w / 999) y_min = int(bbox[1] * h / 999) x_max = int(bbox[2] * w / 999) y_max = int(bbox[3] * h / 999) # 获取颜色 color = color_map.get(obj_id, (255, 255, 255)) # 默认白色 # 绘制边界框 cv2.rectangle(marked_image, (x_min, y_min), (x_max, y_max), color, 2) # 添加标签 label_text = f"{label} ({obj_id})" cv2.putText(marked_image, label_text, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1) return marked_image # --- Flask 路由 --- @app.route('/') def index(): """主页，加载实时图像页面""" return render_template('index.html') @app.route('/list') # 新增路由用于显示图片列表 def list_images(): """加载图片列表页面""" return render_template('list_images.html') @app.route('/api/images', methods=['GET']) def get_images_api(): """API: 获取图片列表 (JSON 格式)""" conn = sqlite3.connect(DATABASE_PATH) cursor = conn.cursor() # 按时间倒序排列，包含标注图片字段 cursor.execute("SELECT id, left_filename, right_filename, left_marked_filename, right_marked_filename, timestamp, metadata, comment, created_at FROM images ORDER BY timestamp DESC") rows = cursor.fetchall() conn.close() images = [] for row in rows: images.append({ "id": row[0], "left_filename": row[1], "right_filename": row[2], "left_marked_filename": row[3] or "", # 如果没有标注图片则显示空字符串 "right_marked_filename": row[4] or "", # 如果没有标注图片则显示空字符串 "timestamp": row[5], "metadata": row[6], "comment": row[7] or "", # 如果没有 comment 则显示空字符串 "created_at": row[8] }) return jsonify(images) @app.route('/api/images', methods=['DELETE']) def delete_image_api(): """API: 删除单张图片记录及其文件""" image_id = request.json.get('id') if not image_id: return jsonify({"error": "Image ID is required"}), 400 conn = sqlite3.connect(DATABASE_PATH) cursor = conn.cursor() # 查询文件名，包含标注图片 cursor.execute("SELECT left_filename, right_filename, left_marked_filename, right_marked_filename FROM images WHERE id = ?", (image_id,)) row = cursor.fetchone() if not row: conn.close() return jsonify({"error": "Image not found"}), 404 left_filename, right_filename, left_marked_filename, right_marked_filename = row # 删除数据库记录 cursor.execute("DELETE FROM images WHERE id = ?", (image_id,)) conn.commit() conn.close() # 删除对应的文件，包括标注图片 left_path = os.path.join(SAVE_PATH_LEFT, left_filename) right_path = os.path.join(SAVE_PATH_RIGHT, right_filename) left_marked_path = os.path.join(SAVE_PATH_LEFT_MARKED, left_marked_filename) if left_marked_filename else None right_marked_path = os.path.join(SAVE_PATH_RIGHT_MARKED, right_marked_filename) if right_marked_filename else None try: if os.path.exists(left_path): os.remove(left_path) logger.info(f"Deleted file: {left_path}") if os.path.exists(right_path): os.remove(right_path) logger.info(f"Deleted file: {right_path}") if left_marked_path and os.path.exists(left_marked_path): os.remove(left_marked_path) logger.info(f"Deleted file: {left_marked_path}") if right_marked_path and os.path.exists(right_marked_path): os.remove(right_marked_path) logger.info(f"Deleted file: {right_marked_path}") except OSError as e: logger.error(f"Error deleting files: {e}") # 即使删除文件失败，数据库记录也已删除，返回成功 pass return jsonify({"message": f"Image {image_id} deleted successfully"}) @app.route('/api/images/export', methods=['POST']) def export_images_api(): """API: 打包导出选中的图片，优先导出标注图片""" selected_ids = request.json.get('ids', []) if not selected_ids: return jsonify({"error": "No image IDs selected"}), 400 conn = sqlite3.connect(DATABASE_PATH) cursor = conn.cursor() placeholders = ','.join('?' * len(selected_ids)) # 查询包含标注图片的文件名 cursor.execute(f"SELECT left_filename, right_filename, left_marked_filename, right_marked_filename FROM images WHERE id IN ({placeholders})", selected_ids) rows = cursor.fetchall() conn.close() if not rows: return jsonify({"error": "No matching images found"}), 404 # 创建临时 ZIP 文件 temp_zip_fd, temp_zip_path = tempfile.mkstemp(suffix='.zip') os.close(temp_zip_fd) # 关闭文件描述符，让 zipfile 模块管理 try: with zipfile.ZipFile(temp_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: for left_fn, right_fn, left_marked_fn, right_marked_fn in rows: # 优先使用标注图片，如果没有则使用原图 left_export_fn = left_marked_fn if left_marked_fn else left_fn right_export_fn = right_marked_fn if right_marked_fn else right_fn # 确定文件路径 left_export_path = os.path.join(SAVE_PATH_LEFT_MARKED if left_marked_fn else SAVE_PATH_LEFT, left_export_fn) right_export_path = os.path.join(SAVE_PATH_RIGHT_MARKED if right_marked_fn else SAVE_PATH_RIGHT, right_export_fn) # 添加到 ZIP 文件 if os.path.exists(left_export_path): zipf.write(left_export_path, os.path.join('left', left_export_fn)) if os.path.exists(right_export_path): zipf.write(right_export_path, os.path.join('right', right_export_fn)) logger.info(f"Exported {len(rows)} image pairs to {temp_zip_path}") # 返回 ZIP 文件给客户端 return send_file(temp_zip_path, as_attachment=True, download_name='exported_images.zip') except Exception as e: logger.error(f"Error creating export ZIP: {e}") if os.path.exists(temp_zip_path): os.remove(temp_zip_path) return jsonify({"error": str(e)}), 500 @app.route('/upload', methods=['POST']) def upload_images(): """接收左右摄像头图片，保存并推送更新，同时生成标注图片""" try: # 从 multipart/form-data 中获取文件 left_file = request.files.get('left_image') right_file = request.files.get('right_image') metadata_str = request.form.get('metadata') # 如果需要处理元数据 comment = request.form.get('comment', '') # 获取 comment 字段 if not left_file or not right_file: logger.warning("Received request without required image files.") return jsonify({"error": "Missing left_image or right_image"}), 400 # 读取图片数据 left_img_bytes = left_file.read() right_img_bytes = right_file.read() # 解码图片用于后续处理 (如显示、保存) nparr_left = np.frombuffer(left_img_bytes, np.uint8) nparr_right = np.frombuffer(right_img_bytes, np.uint8) img_left = cv2.imdecode(nparr_left, cv2.IMREAD_COLOR) img_right = cv2.imdecode(nparr_right, cv2.IMREAD_COLOR) if img_left is None or img_right is None: logger.error("Failed to decode received images.") return jsonify({"error": "Could not decode images"}), 400 # 解析元数据 (如果提供) metadata = {} if metadata_str: try: metadata = json.loads(metadata_str) except json.JSONDecodeError as e: logger.warning(f"Could not parse metadata: {e}") timestamp_str = str(metadata.get("timestamp", str(int(time.time())))) timestamp_str_safe = timestamp_str.replace(".", "_") # 避免文件名中的点号问题 # 生成文件名 left_filename = f"left_{timestamp_str_safe}.jpg" right_filename = f"right_{timestamp_str_safe}.jpg" left_marked_filename = f"left_marked_{timestamp_str_safe}.jpg" # 标注图片文件名 right_marked_filename = f"right_marked_{timestamp_str_safe}.jpg" # 标注图片文件名 # 保存原图到本地 left_path = os.path.join(SAVE_PATH_LEFT, left_filename) right_path = os.path.join(SAVE_PATH_RIGHT, right_filename) # 确保目录存在 os.makedirs(SAVE_PATH_LEFT, exist_ok=True) os.makedirs(SAVE_PATH_RIGHT, exist_ok=True) os.makedirs(SAVE_PATH_LEFT_MARKED, exist_ok=True) # 创建标注图片目录 os.makedirs(SAVE_PATH_RIGHT_MARKED, exist_ok=True) # 创建标注图片目录 cv2.imwrite(left_path, img_left) cv2.imwrite(right_path, img_right) logger.info(f"Saved original images: {left_path}, {right_path}") # 使用 VisionAPIClient 处理图片并生成标注图片 left_marked_path = None right_marked_path = None try: with VisionAPIClient() as client: # 处理左图 left_task_id = client.submit_task(image_id=1, image=img_left) # 处理右图 right_task_id = client.submit_task(image_id=2, image=img_right) # 等待任务完成 client.task_queue.join() # 获取处理结果 left_result = client.get_result(left_task_id) right_result = client.get_result(right_task_id) # 生成标注图片 if left_result and left_result.success: marked_left_img = draw_detections_on_image(img_left, left_result.detections) left_marked_path = os.path.join(SAVE_PATH_LEFT_MARKED, left_marked_filename) cv2.imwrite(left_marked_path, marked_left_img) logger.info(f"Saved marked left image: {left_marked_path}") if right_result and right_result.success: marked_right_img = draw_detections_on_image(img_right, right_result.detections) right_marked_path = os.path.join(SAVE_PATH_RIGHT_MARKED, right_marked_filename) cv2.imwrite(right_marked_path, marked_right_img) logger.info(f"Saved marked right image: {right_marked_path}") except Exception as e: logger.error(f"Error processing images with VisionAPIClient: {e}") # 即使处理失败，也继续保存原图 # 将图片信息写入数据库，包含标注图片字段 conn = sqlite3.connect(DATABASE_PATH) cursor = conn.cursor() cursor.execute(''' INSERT INTO images (left_filename, right_filename, left_marked_filename, right_marked_filename, timestamp, metadata, comment) VALUES (?, ?, ?, ?, ?, ?, ?) ''', (left_filename, right_filename, left_marked_filename if left_marked_path else None, right_marked_filename if right_marked_path else None, float(timestamp_str), json.dumps(metadata), comment)) conn.commit() image_id = cursor.lastrowid # 获取新插入记录的 ID conn.close() logger.info(f"Recorded image pair (ID: {image_id}) in database.") # 将 OpenCV 图像编码为 base64 字符串，用于 WebSocket 传输 _, left_encoded = cv2.imencode('.jpg', img_left) _, right_encoded = cv2.imencode('.jpg', img_right) left_b64 = base64.b64encode(left_encoded).decode('utf-8') right_b64 = base64.b64encode(right_encoded).decode('utf-8') # 更新用于实时显示的全局变量 (如果需要) with frame_lock: global latest_left_frame, latest_right_frame, latest_timestamp latest_left_frame = img_left latest_right_frame = img_right latest_timestamp = timestamp_str socketio.emit('update_image', { 'left_image': left_b64, 'right_image': right_b64, 'timestamp': timestamp_str }) return jsonify({"message": "Images received, saved, and pushed via WebSocket", "timestamp": timestamp_str, "id": image_id}) except Exception as e: logger.error(f"Error processing upload: {e}") return jsonify({"error": str(e)}), 500 @app.route('/api/images/comment', methods=['PUT']) def update_image_comment(): """API: 更新图片的 comment""" data = request.json image_id = data.get('id') comment = data.get('comment', '') if not image_id: return jsonify({"error": "Image ID is required"}), 400 conn = sqlite3.connect(DATABASE_PATH) cursor = conn.cursor() # 更新 comment 字段 cursor.execute("UPDATE images SET comment = ? WHERE id = ?", (comment, image_id)) conn.commit() conn.close() return jsonify({"message": f"Comment for image {image_id} updated successfully"}) @app.route('/status') def status(): with frame_lock: has_frames = latest_left_frame is not None and latest_right_frame is not None timestamp = latest_timestamp if has_frames else "N/A" return jsonify({"has_frames": has_frames, "latest_timestamp": timestamp}) # --- SocketIO 事件处理程序 --- @socketio.event def connect(): """客户端连接事件""" logger.info("Client connected") @socketio.event def disconnect(): """客户端断开连接事件""" logger.info("Client disconnected") @socketio.event def capture_button(data): """ SocketIO 事件处理器，当客户端发送 'button_pressed' 事件时触发 """ try: image_client.request_sync() logger.info("Request sent to server.") except Exception as e: logger.error(f"Error sending request: {e}") if __name__ == '__main__': logger.info(f"Starting Flask-SocketIO server on {FLASK_HOST}:{FLASK_PORT}") socketio.run(app, host=FLASK_HOST, port=FLASK_PORT, debug=False, allow_unsafe_werkzeug=True)