Add image back

demo_visualization_api.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python3
+"""
+演示可视化API功能的示例脚本
+展示如何获取绘制了关键点和骨架的图像
+"""
+
+import requests
+import base64
+import json
+import cv2
+import numpy as np
+import time
+
+def demo_api_usage():
+    """演示API使用方法"""
+    
+    print("🎯 单鼠姿态检测可视化API演示")
+    print("=" * 50)
+    
+    # 1. 准备图像数据
+    # 这里使用您提供的示例格式
+    image_path = "test_mouse.jpg"
+    
+    if not cv2.imread(image_path) is not None:
+        print("❌ 找不到测试图像，请确保test_mouse.jpg存在")
+        return
+    
+    # 将图像转换为Base64
+    with open(image_path, "rb") as image_file:
+        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
+    
+    # 2. 准备API请求数据（按照您的示例格式）
+    request_data = {
+        "image": image_base64,
+        "conf_threshold": 0.3,
+        "frame_id": 1,
+        "timestamp": 1641234567.123,
+        "return_image": True,    # 🔥 关键参数：返回可视化图像
+        "draw_bbox": False       # 只绘制关键点和骨架，不绘制边界框
+    }
+    
+    print("📤 发送API请求...")
+    print(f"请求参数: conf_threshold={request_data['conf_threshold']}")
+    print(f"         return_image={request_data['return_image']}")
+    print(f"         draw_bbox={request_data['draw_bbox']}")
+    
+    # 3. 发送请求
+    api_url = "http://localhost:8765/api/process_frame"
+    
+    try:
+        start_time = time.time()
+        response = requests.post(api_url, json=request_data)
+        request_time = time.time() - start_time
+        
+        if response.status_code == 200:
+            result = response.json()
+            
+            print(f"✅ 请求成功！(耗时: {request_time:.3f}s)")
+            print("\n📊 检测结果:")
+            print(f"  - 检测到小鼠: {result['mouse_detected']}")
+            
+            if result['mouse_detected']:
+                print(f"  - 检测置信度: {result['confidence']:.3f}")
+                print(f"  - 关键点数量: {len(result['keypoints'])}")
+                print(f"  - 处理FPS: {result['fps']:.1f}")
+                
+                # 显示关键点信息
+                print("\n🔍 检测到的关键点:")
+                for kpt in result['keypoints']:
+                    print(f"    {kpt['id']}: {kpt['name']} "
+                          f"({kpt['x']:.1f}, {kpt['y']:.1f}) "
+                          f"conf={kpt['confidence']:.3f}")
+            
+            # 4. 🎨 处理可视化图像
+            if "visualization_image" in result and result["visualization_image"]:
+                print("\n🎨 可视化图像处理:")
+                print("  ✅ 接收到绘制了关键点和骨架的图像")
+                
+                # 解码Base64图像
+                img_bytes = base64.b64decode(result["visualization_image"])
+                nparr = np.frombuffer(img_bytes, np.uint8)
+                vis_image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+                
+                # 保存可视化结果
+                output_path = "demo_visualization_result.jpg"
+                cv2.imwrite(output_path, vis_image)
+                
+                print(f"  💾 可视化图像已保存: {output_path}")
+                print(f"  📐 图像尺寸: {vis_image.shape}")
+                
+                # 显示图像信息
+                print(f"  📊 Base64数据大小: {len(result['visualization_image'])} 字符")
+                print(f"  📁 文件大小: {len(img_bytes)} 字节")
+                
+            else:
+                print("❌ 没有接收到可视化图像")
+            
+            print("\n" + "=" * 50)
+            print("🎉 演示完成！")
+            print("\n💡 使用说明:")
+            print("  1. 设置 return_image=true 来获取可视化图像")
+            print("  2. 设置 draw_bbox=false 只绘制关键点和骨架（如您所需）")
+            print("  3. 设置 draw_bbox=true 同时绘制边界框")
+            print("  4. 返回的 visualization_image 是Base64编码的JPEG图像")
+            
+        else:
+            print(f"❌ API请求失败: {response.status_code}")
+            print(f"错误信息: {response.text}")
+    
+    except Exception as e:
+        print(f"❌ 请求出错: {str(e)}")
+
+def compare_with_without_visualization():
+    """对比有无可视化的API响应"""
+    
+    print("\n" + "=" * 50)
+    print("🔄 对比测试：有无可视化图像的API响应")
+    print("=" * 50)
+    
+    # 准备图像
+    with open("test_mouse.jpg", "rb") as image_file:
+        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
+    
+    # 测试配置
+    configs = [
+        {
+            "name": "不返回可视化图像",
+            "data": {
+                "image": image_base64,
+                "conf_threshold": 0.3,
+                "return_image": False
+            }
+        },
+        {
+            "name": "返回可视化图像",
+            "data": {
+                "image": image_base64,
+                "conf_threshold": 0.3,
+                "return_image": True,
+                "draw_bbox": False
+            }
+        }
+    ]
+    
+    for config in configs:
+        print(f"\n📋 测试: {config['name']}")
+        
+        try:
+            response = requests.post("http://localhost:8765/api/process_frame", 
+                                   json=config["data"])
+            
+            if response.status_code == 200:
+                result = response.json()
+                
+                # 计算响应大小
+                response_size = len(response.content)
+                has_vis_image = "visualization_image" in result and result["visualization_image"]
+                
+                print(f"  📊 响应大小: {response_size:,} 字节")
+                print(f"  🖼️  包含可视化图像: {'是' if has_vis_image else '否'}")
+                
+                if has_vis_image:
+                    vis_size = len(result["visualization_image"])
+                    print(f"  📸 图像数据大小: {vis_size:,} 字符")
+                
+                print(f"  ⚡ 处理时间: {result['processing_time']:.3f}s")
+                print(f"  🎯 检测结果: {'检测到小鼠' if result['mouse_detected'] else '未检测到小鼠'}")
+                
+        except Exception as e:
+            print(f"  ❌ 测试失败: {str(e)}")
+
+if __name__ == "__main__":
+    # 检查服务是否可用
+    try:
+        response = requests.get("http://localhost:8765/api/status", timeout=5)
+        if response.status_code != 200:
+            print("❌ WebRTC服务不可用，请先启动服务:")
+            print("   python gradio_webrtc_api.py")
+            exit(1)
+    except Exception:
+        print("❌ 无法连接到WebRTC服务，请先启动服务:")
+        print("   python gradio_webrtc_api.py")
+        exit(1)
+    
+    # 运行演示
+    demo_api_usage()
+    compare_with_without_visualization() 

gradio_webrtc_api.py

@@ -75,6 +75,8 @@ class FrameData(BaseModel):
     conf_threshold: float = 0.3
     timestamp: float = None
     frame_id: int = None
+    return_image: bool = False  # 是否返回可视化图像
+    draw_bbox: bool = False     # 是否绘制边界框
 
 class ProcessingResult(BaseModel):
     success: bool
@@ -87,6 +89,7 @@ class ProcessingResult(BaseModel):
     processing_time: float = 0.0
     fps: float = 0.0
     error: Optional[str] = None
+    visualization_image: Optional[str] = None  # Base64编码的可视化图像
 
 # API端点
 @app.get("/")
@@ -120,8 +123,13 @@ async def process_frame_api(frame_data: FrameData):
         if frame is None:
             raise HTTPException(status_code=400, detail="无法解码图像")
         
-        # 处理帧
-        result = processor.process_frame(frame, frame_data.conf_threshold)
+        # 处理帧，支持返回可视化图像
+        result = processor.process_frame(
+            frame, 
+            frame_data.conf_threshold,
+            return_image=frame_data.return_image,
+            draw_bbox=frame_data.draw_bbox
+        )
         
         # 如果有frame_id，使用提供的值
         if frame_data.frame_id is not None:
@@ -229,9 +237,16 @@ async def process_frame_websocket(frame_data: Dict[str, Any]) -> Dict[str, Any]:
         
         # 获取配置参数
         conf_threshold = frame_data.get("conf_threshold", 0.3)
+        return_image = frame_data.get("return_image", False)
+        draw_bbox = frame_data.get("draw_bbox", False)
         
-        # 处理帧
-        result = processor.process_frame(frame, conf_threshold)
+        # 处理帧，支持返回可视化图像
+        result = processor.process_frame(
+            frame, 
+            conf_threshold,
+            return_image=return_image,
+            draw_bbox=draw_bbox
+        )
         
         # 保留原始的frame_id和timestamp（如果提供）
         if "frame_id" in frame_data:
@@ -277,8 +292,13 @@ async def process_batch_frames(frames: List[FrameData]):
                 })
                 continue
             
-            # 处理帧
-            result = processor.process_frame(frame, frame_data.conf_threshold)
+            # 处理帧，支持返回可视化图像
+            result = processor.process_frame(
+                frame, 
+                frame_data.conf_threshold,
+                return_image=frame_data.return_image,
+                draw_bbox=frame_data.draw_bbox
+            )
             
             # 保留原始信息
             if frame_data.frame_id is not None:

gradio_webrtc_server.py

@@ -121,13 +121,57 @@ class SingleMouseProcessor:
             logger.error(traceback.format_exc())
             return False
 
-    def process_frame(self, frame: np.ndarray, conf_threshold: float = 0.3) -> Dict[str, Any]:
+    def _draw_keypoints_and_skeleton(self, frame: np.ndarray, keypoints: List[Dict], draw_bbox: bool = False, bbox: Dict = None) -> np.ndarray:
+        """
+        在图像上绘制关键点和骨架
+        
+        Args:
+            frame: 输入图像帧
+            keypoints: 关键点列表
+            draw_bbox: 是否绘制边界框
+            bbox: 边界框信息
+            
+        Returns:
+            绘制后的图像
+        """
+        vis_frame = frame.copy()
+        
+        # 绘制边界框（如果需要）
+        if draw_bbox and bbox is not None:
+            cv2.rectangle(vis_frame, 
+                         (int(bbox["x1"]), int(bbox["y1"])), 
+                         (int(bbox["x2"]), int(bbox["y2"])), 
+                         (0, 255, 0), 2)
+        
+        # 绘制关键点
+        for kpt in keypoints:
+            cv2.circle(vis_frame, 
+                      (int(kpt["x"]), int(kpt["y"])), 
+                      5, (0, 0, 255), -1)
+            cv2.putText(vis_frame, 
+                       kpt["name"], 
+                       (int(kpt["x"]) + 5, int(kpt["y"]) - 5),
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
+        
+        # 绘制骨架连接
+        kpt_dict = {kpt["id"]: (kpt["x"], kpt["y"]) for kpt in keypoints}
+        for connection in self.keypoint_connections:
+            if connection[0] in kpt_dict and connection[1] in kpt_dict:
+                pt1 = (int(kpt_dict[connection[0]][0]), int(kpt_dict[connection[0]][1]))
+                pt2 = (int(kpt_dict[connection[1]][0]), int(kpt_dict[connection[1]][1]))
+                cv2.line(vis_frame, pt1, pt2, (255, 0, 0), 2)
+        
+        return vis_frame
+
+    def process_frame(self, frame: np.ndarray, conf_threshold: float = 0.3, return_image: bool = False, draw_bbox: bool = False) -> Dict[str, Any]:
         """
         处理单帧图像并返回检测结果
         
         Args:
             frame: 输入图像帧
             conf_threshold: 置信度阈值
+            return_image: 是否返回绘制了关键点的图像
+            draw_bbox: 是否在返回图像中绘制边界框
             
         Returns:
             包含检测结果的字典
@@ -220,6 +264,27 @@ class SingleMouseProcessor:
                                                 "confidence": float(kpt[2])
                                             })
                 
+                # 生成可视化图像（如果需要）
+                if return_image and detection_data["mouse_detected"]:
+                    vis_frame = self._draw_keypoints_and_skeleton(
+                        frame, 
+                        detection_data["keypoints"], 
+                        draw_bbox=draw_bbox,
+                        bbox=detection_data["bbox"]
+                    )
+                    
+                    # 将图像编码为Base64
+                    _, buffer = cv2.imencode('.jpg', vis_frame)
+                    image_base64 = base64.b64encode(buffer).decode('utf-8')
+                    detection_data["visualization_image"] = image_base64
+                elif return_image:
+                    # 即使没有检测到小鼠也返回原图
+                    _, buffer = cv2.imencode('.jpg', frame)
+                    image_base64 = base64.b64encode(buffer).decode('utf-8')
+                    detection_data["visualization_image"] = image_base64
+                else:
+                    logger.info("不需要生成可视化图像")
+                
                 # 更新统计信息
                 process_time = time.time() - start_time
                 self.frame_count += 1
@@ -384,9 +449,50 @@ def create_gradio_app():
                 "image": "base64_encoded_image",
                 "conf_threshold": 0.3,
                 "frame_id": 1,
-                "timestamp": 1641234567.123
+                "timestamp": 1641234567.123,
+                "return_image": true,
+                "draw_bbox": false
             }
             ```
+            
+            ### 新增参数说明
+            - **return_image** (bool, 可选): 是否返回绘制了关键点和骨架的可视化图像
+            - **draw_bbox** (bool, 可选): 是否在可视化图像中绘制边界框（需要return_image=true）
+            
+            ### 响应格式
+            当 `return_image=true` 时，响应会包含额外的字段：
+            ```json
+            {
+                "success": true,
+                "mouse_detected": true,
+                "keypoints": [...],
+                "visualization_image": "base64_encoded_image_with_keypoints_and_skeleton"
+            }
+            ```
+            
+            ### 使用示例
+            ```javascript
+            // 只获取检测数据
+            fetch('/api/process_frame', {
+                method: 'POST',
+                headers: {'Content-Type': 'application/json'},
+                body: JSON.stringify({
+                    image: base64Image,
+                    return_image: false
+                })
+            });
+            
+            // 获取检测数据和可视化图像
+            fetch('/api/process_frame', {
+                method: 'POST', 
+                headers: {'Content-Type': 'application/json'},
+                body: JSON.stringify({
+                    image: base64Image,
+                    return_image: true,
+                    draw_bbox: false  // 只绘制关键点和骨架，不绘制边界框
+                })
+            });
+            ```
             """)
     
     return app
@@ -432,9 +538,16 @@ def process_webrtc_data(data: Dict[str, Any]) -> Dict[str, Any]:
         
         # 获取配置参数
         conf_threshold = data.get("conf_threshold", 0.3)
+        return_image = data.get("return_image", False)
+        draw_bbox = data.get("draw_bbox", False)
         
-        # 处理帧
-        result = processor.process_frame(frame, conf_threshold)
+        # 处理帧，支持返回可视化图像
+        result = processor.process_frame(
+            frame, 
+            conf_threshold,
+            return_image=return_image,
+            draw_bbox=draw_bbox
+        )
         
         return result
         

test_visualization_api.py

@@ -0,0 +1,235 @@
+#!/usr/bin/env python3
+"""
+测试可视化图像API功能
+"""
+
+import requests
+import base64
+import json
+import cv2
+import numpy as np
+import time
+import os
+
+def load_test_image():
+    """加载测试图像"""
+    # 尝试查找测试图像
+    test_images = ["test_mouse.jpg", "test_mouse.png", "sample.jpg"]
+    
+    for img_path in test_images:
+        if os.path.exists(img_path):
+            print(f"✅ 找到测试图像: {img_path}")
+            return cv2.imread(img_path)
+    
+    # 如果没有找到测试图像，创建一个简单的
+    print("📸 创建合成测试图像...")
+    test_img = np.random.randint(50, 200, (640, 640, 3), dtype=np.uint8)
+    
+    # 添加简单的小鼠形状
+    cv2.circle(test_img, (320, 300), 50, (100, 100, 100), -1)  # 头部
+    cv2.ellipse(test_img, (320, 380), (80, 120), 0, 0, 360, (120, 120, 120), -1)  # 身体
+    cv2.ellipse(test_img, (320, 500), (20, 80), 0, 0, 360, (90, 90, 90), -1)  # 尾巴
+    
+    return test_img
+
+def image_to_base64(image):
+    """将OpenCV图像转换为Base64字符串"""
+    _, buffer = cv2.imencode('.jpg', image)
+    return base64.b64encode(buffer).decode('utf-8')
+
+def base64_to_image(base64_str):
+    """将Base64字符串转换为OpenCV图像"""
+    img_bytes = base64.b64decode(base64_str)
+    nparr = np.frombuffer(img_bytes, np.uint8)
+    return cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+def test_api_with_visualization():
+    """测试带可视化功能的API"""
+    print("=== 测试可视化图像API功能 ===")
+    
+    # 加载测试图像
+    test_image = load_test_image()
+    if test_image is None:
+        print("❌ 无法加载测试图像")
+        return False
+    
+    print(f"测试图像尺寸: {test_image.shape}")
+    
+    # 转换为Base64
+    image_base64 = image_to_base64(test_image)
+    
+    # 测试不同的配置
+    test_configs = [
+        {
+            "name": "不返回图像",
+            "data": {
+                "image": image_base64,
+                "conf_threshold": 0.1,
+                "return_image": False,
+                "frame_id": 1,
+                "timestamp": time.time()
+            }
+        },
+        {
+            "name": "返回图像（不包括边界框）",
+            "data": {
+                "image": image_base64,
+                "conf_threshold": 0.1,
+                "return_image": True,
+                "draw_bbox": False,
+                "frame_id": 2,
+                "timestamp": time.time()
+            }
+        },
+        {
+            "name": "返回图像（包括边界框）",
+            "data": {
+                "image": image_base64,
+                "conf_threshold": 0.1,
+                "return_image": True,
+                "draw_bbox": True,
+                "frame_id": 3,
+                "timestamp": time.time()
+            }
+        }
+    ]
+    
+    api_url = "http://localhost:8765/api/process_frame"
+    
+    for i, config in enumerate(test_configs):
+        print(f"\n--- 测试 {i+1}: {config['name']} ---")
+        
+        try:
+            # 发送请求
+            start_time = time.time()
+            response = requests.post(api_url, json=config["data"])
+            request_time = time.time() - start_time
+            
+            if response.status_code == 200:
+                result = response.json()
+                print(f"✅ 请求成功 (耗时: {request_time:.3f}s)")
+                print(f"检测到小鼠: {result['mouse_detected']}")
+                
+                if result['mouse_detected']:
+                    print(f"置信度: {result['confidence']:.3f}")
+                    print(f"关键点数量: {len(result['keypoints'])}")
+                    print(f"处理FPS: {result['fps']:.1f}")
+                
+                # 检查是否有可视化图像
+                if "visualization_image" in result and result["visualization_image"]:
+                    print("✅ 包含可视化图像")
+                    
+                    # 解码并保存可视化图像
+                    vis_image = base64_to_image(result["visualization_image"])
+                    output_filename = f"test_result_{i+1}_{config['name'].replace(' ', '_').replace('（', '').replace('）', '')}.jpg"
+                    cv2.imwrite(output_filename, vis_image)
+                    print(f"可视化图像已保存: {output_filename}")
+                    
+                else:
+                    if "visualization_image" in result:
+                        print(f"❌ visualization_image存在但为空: {result['visualization_image'] is None}")
+                    else:
+                        print("❌ 响应中没有visualization_image字段")
+                    print(f"响应中的所有字段: {list(result.keys())}")
+                    # 打印部分响应内容（不包括过长的字段）
+                    debug_result = {k: v for k, v in result.items() if k != 'visualization_image'}
+                    print(f"响应内容（除图像外）: {json.dumps(debug_result, indent=2, ensure_ascii=False)}")
+                
+            else:
+                print(f"❌ 请求失败: {response.status_code}")
+                print(f"错误信息: {response.text}")
+                
+        except Exception as e:
+            print(f"❌ 请求出错: {str(e)}")
+    
+    return True
+
+def test_websocket_with_visualization():
+    """测试WebSocket的可视化功能"""
+    print("\n=== 测试WebSocket可视化功能 ===")
+    
+    try:
+        import websockets
+        import asyncio
+        
+        async def test_ws():
+            # 加载测试图像
+            test_image = load_test_image()
+            image_base64 = image_to_base64(test_image)
+            
+            uri = "ws://localhost:8765/ws/stream"
+            
+            async with websockets.connect(uri) as websocket:
+                print("✅ WebSocket连接成功")
+                
+                # 发送带可视化的数据
+                frame_data = {
+                    "image": image_base64,
+                    "conf_threshold": 0.1,  # 降低置信度阈值
+                    "return_image": True,
+                    "draw_bbox": False,
+                    "frame_id": 100,
+                    "timestamp": time.time()
+                }
+                
+                await websocket.send(json.dumps(frame_data))
+                
+                # 接收结果
+                response = await websocket.recv()
+                result = json.loads(response)
+                
+                if result['success']:
+                    print(f"✅ WebSocket处理成功")
+                    print(f"检测到小鼠: {result['mouse_detected']}")
+                    
+                    if "visualization_image" in result and result["visualization_image"]:
+                        print("✅ 包含可视化图像")
+                        
+                        # 保存WebSocket结果
+                        vis_image = base64_to_image(result["visualization_image"])
+                        cv2.imwrite("websocket_result.jpg", vis_image)
+                        print("WebSocket可视化图像已保存: websocket_result.jpg")
+                    else:
+                        print("❌ WebSocket结果不包含可视化图像")
+                else:
+                    print(f"❌ WebSocket处理失败: {result.get('error', 'Unknown error')}")
+        
+        # 运行异步测试
+        asyncio.run(test_ws())
+        return True
+        
+    except ImportError:
+        print("❌ 未安装websockets库，跳过WebSocket测试")
+        return False
+    except Exception as e:
+        print(f"❌ WebSocket测试出错: {str(e)}")
+        return False
+
+def main():
+    """主函数"""
+    print("🚀 开始测试可视化图像API功能")
+    print("确保WebRTC服务已启动在 http://localhost:8765")
+    
+    # 检查服务是否可用
+    try:
+        response = requests.get("http://localhost:8765/api/status", timeout=5)
+        if response.status_code == 200:
+            print("✅ WebRTC服务可用")
+        else:
+            print("❌ WebRTC服务不可用")
+            return
+    except Exception as e:
+        print(f"❌ 无法连接到WebRTC服务: {str(e)}")
+        print("请先启动服务: python gradio_webrtc_api.py")
+        return
+    
+    # 测试REST API
+    test_api_with_visualization()
+    
+    # 测试WebSocket
+    test_websocket_with_visualization()
+    
+    print("\n🎉 测试完成！检查生成的图像文件查看结果。")
+
+if __name__ == "__main__":
+    main()