如何将FastAPI POST请求响应传入WebSocket的get_stream函数

你的核心问题是错误地将POST接口的处理函数作为内部函数直接调用，且没有正确传递参数。正确的做法是抽离地址验证逻辑，再通过合理的方式让WebSocket获取到合法的摄像头地址，下面提供两种可行的实现方案：

方案1：先通过POST提交并验证地址，再通过WebSocket传递标识获取地址 #

适合需要先验证地址有效性，再建立视频流连接的场景，需通过全局存储（或会话）暂存验证后的地址。

步骤1：抽离地址验证逻辑 #

把原POST接口里的地址验证逻辑单独写成一个函数，供POST和WebSocket复用：

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from pydantic import BaseModel
from typing import Union, Optional
import cv2
import uvicorn

app = FastAPI()

# 全局存储验证后的摄像头地址（多客户端场景建议用会话或数据库，这里仅作示例）
valid_camera_addresses: dict[str, Union[str, int]] = {}

class Address(BaseModel):
    camera_id: Union[str, int] = 0
    client_id: str  # 客户端唯一标识，用于后续WebSocket关联地址

def validate_camera_address(camera_id: Union[str, int]) -> bool:
    """验证摄像头地址是否合法"""
    if isinstance(camera_id, int):
        return camera_id in (0, 1)
    elif isinstance(camera_id, str):
        return camera_id.isnumeric() or camera_id.endswith('.txt') or camera_id.lower().startswith(
            ('rtsp://', 'rtmp://', 'http://', 'https://'))
    return False

步骤2：修改POST接口存储合法地址 #

@app.post("/camera_id")
async def submit_camera_address(address: Address):
    is_valid = validate_camera_address(address.camera_id)
    if is_valid:
        valid_camera_addresses[address.client_id] = address.camera_id
        return {"status": "success", "message": "Camera address validated"}
    else:
        return {
           "message": "Incorrect Camera Address",
           "status": str(address.camera_id),
           "expect_input": "The camera address should be (0, 1) or start with ('rtsp://', 'rtmp://', 'http://', 'https://')"
        }

步骤3：修改WebSocket函数获取存储的地址 #

@app.websocket("/ws/{client_id}")
async def get_stream(websocket: WebSocket, client_id: str):
    await websocket.accept()
    
    # 获取之前验证过的地址
    camera_id = valid_camera_addresses.get(client_id)
    if not camera_id:
        await websocket.send_text("Error: No valid camera address found for this client")
        await websocket.close()
        return
    
    # 处理视频流
    camera = cv2.VideoCapture(str(camera_id) if isinstance(camera_id, str) else camera_id)
    if not camera.isOpened():
        await websocket.send_text("Error: Failed to open camera")
        await websocket.close()
        return
    
    try:
        while True:
            ret, frame = camera.read()  # 原代码camera.frame()是错误的，应该用read()
            if ret:
                ret, buffer = cv2.imencode('.jpg', frame)
                await websocket.send_bytes(buffer.tobytes())
            else:
                await websocket.send_text("No frame captured")
                break
    except WebSocketDisconnect:
        print("Client disconnected")
    finally:
        camera.release()

方案2：WebSocket直接接收地址并验证 #

适合不需要单独POST验证，直接在WebSocket连接后传递地址的场景，更简洁：

修正后的完整代码 #

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from pydantic import BaseModel
from typing import Union
import cv2
import uvicorn

app = FastAPI()

class Address(BaseModel):
    camera_id: Union[str, int] = 0

def validate_camera_address(camera_id: Union[str, int]) -> bool:
    """验证摄像头地址是否合法"""
    if isinstance(camera_id, int):
        return camera_id in (0, 1)
    elif isinstance(camera_id, str):
        return camera_id.isnumeric() or camera_id.endswith('.txt') or camera_id.lower().startswith(
            ('rtsp://', 'rtmp://', 'http://', 'https://'))
    return False

@app.websocket("/ws")
async def get_stream(websocket: WebSocket):
    await websocket.accept()
    
    # 第一步：接收客户端发送的摄像头地址
    address_data = await websocket.receive_json()
    try:
        address = Address(**address_data)
    except Exception:
        await websocket.send_text("Error: Invalid address format")
        await websocket.close()
        return
    
    # 验证地址合法性
    if not validate_camera_address(address.camera_id):
        await websocket.send_text("Error: Incorrect Camera Address. Should be (0,1) or start with rtsp:///rtmp:///http:///https://")
        await websocket.close()
        return
    
    # 处理视频流
    camera = cv2.VideoCapture(str(address.camera_id) if isinstance(address.camera_id, str) else address.camera_id)
    if not camera.isOpened():
        await websocket.send_text("Error: Failed to open camera")
        await websocket.close()
        return
    
    try:
        while True:
            ret, frame = camera.read()
            if ret:
                ret, buffer = cv2.imencode('.jpg', frame)
                await websocket.send_bytes(buffer.tobytes())
            else:
                await websocket.send_text("No frame captured")
                break
    except WebSocketDisconnect:
        print("Client disconnected")
    finally:
        camera.release()

原代码的关键错误修正 #

1. 调用方式错误：不能直接调用POST接口的处理函数address，需抽离复用逻辑，或通过客户端传递参数。
2. OpenCV调用错误：原代码camera.frame()不存在，应使用ret, frame = camera.read()获取帧。
3. 变量未定义：原代码del frame, result中的result未声明，需删除或修正。
4. 类型导入缺失：Union需从typing模块导入。

内容的提问来源于stack exchange，提问作者benjamin olise