在 Python 中使用 Websocket 协议传输数据

Webscoket 是客户端和服务器之间的一种全双工通信协议。一旦客户端与服务器建立起连接,之后的全部数据通信都通过这个连接进行。通信过程中,可互相发送 JSON、XML、HTML 或 Image 等任意格式的数据。

本文介绍了如何在 Python 中搭建一个 Websocket 服务器和对应的客户端,并实现一个简单的功能:客户端发送一串图片给服务器,服务器在处理完成之后发送结果给客户端。

简介

Websocket 协议的优点:

  • 建立在 TCP 协议之上,可以建立持久性的 TCP 连接
  • 服务器可以主动向客户端推送信息,客户端也可以主动向服务器发送信息,是真正的双向平等对话
  • 数据格式比较轻量,性能开销小,通信高效
  • 同一服务器可以和多台客户端建立连接,连接之间互不干扰

安装

使用 pip

pip install websockets

代码实现

服务器

下面是 WebSocket 服务端代码。在和客户端建立连接之后,服务器会一直接收客户端发送的图片,直到客户端发送 end 字符串表示传输完成为止。服务器会对接收到的图片进行处理(比如 OCR),这个过程可能比较漫长(用 time.sleep(1) 模拟)。全部处理完成后,服务器会把结果(示例中,结果就是 Hello from server! 这个字符串)发送给客户端。

import cv2
import time
import base64
import asyncio
import websockets
import numpy as np


async def hello(websocket, path):
while True:
data = await websocket.recv()

if data == 'end':
break

image = cv2.imdecode(np.frombuffer(base64.b64decode(data), dtype=np.uint8), flags=1)

print('Size of received image:', image.shape, 'Processing...')
time.sleep(1)

print('Sending final results to client...')
await websocket.send('Hello from server!')

print('All Done!')


start_server = websockets.serve(hello, "localhost", 8888, max_size=2**30)
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()

客户端

下面是 WebSocket 客户端代码。在和服务器建立连接之后,客户端会一直发送图片给服务器,最后发送 end 字符串表示传输完成,并等待服务器返回结果。

import cv2
import time
import glob
import base64
import asyncio
import websockets


async def hello():
uri = "ws://localhost:8888"
async with websockets.connect(uri, max_size=2**30) as websocket:
for image_path in glob.glob('images/*.jpg'):
image = cv2.imread(image_path)
image_str = base64.b64encode(cv2.imencode('.jpg', image)[1]).decode()

# send data
print('Sending image:', image_path)
await websocket.send(image_str)

print('Sending `end`...')
await websocket.send('end')

print('Waiting for result...')
res = await websocket.recv()
print('Result:', res)

asyncio.get_event_loop().run_until_complete(hello())

执行

上述场景中,Websocket 相比于传统 HTTP 请求的优势在于:

  • 客户端可以非阻塞发送图片
  • 服务器在处理完成之后,可以主动将结果发送给客户端,无需客户端轮询

服务器的执行结果如下(示例):

Size of received image: (3024, 4032, 3) Processing...
Size of received image: (3024, 4032, 3) Processing...
Size of received image: (3024, 4032, 3) Processing...
Size of received image: (2974, 3947, 3) Processing...
Size of received image: (3024, 4032, 3) Processing...
Size of received image: (3024, 4032, 3) Processing...
Size of received image: (3024, 4032, 3) Processing...
Sending final results to client...
All Done!

客户端的执行结果如下(示例):

Sending image: images/IMG_6667.jpg
Sending image: images/IMG_6673.jpg
Sending image: images/IMG_6666.jpg
Sending image: images/IMG_6664.jpg
Sending image: images/IMG_6658.jpg
Sending image: images/IMG_6659.jpg
Sending image: images/IMG_6660.jpg
Sending `end`...
Waiting for result...
Result: Hello from server!