DevelopmentT/background-remover

1"""2WebSocket test client for manual testing.3 4Usage:5    python test_ws_client.py [IMAGE_PATH]6 7Connects to ws://localhost:7860/ws/ocr, sends an image, and prints8each streamed OCR chunk as it arrives.9"""10 11import asyncio12import json13import sys14 15import websockets16 17 18async def main(image_path: str) -> None:19    uri = "ws://localhost:7860/ws/ocr?user_id=test_user_1"20 21    print(f"🔌 Connecting to {uri}")22    async with websockets.connect(uri) as ws:23        # Read the connected message24        msg = json.loads(await ws.recv())25        print(f"✅ {msg['event']}: {msg.get('data', {}).get('message', '')}")26 27        # Send image as binary28        with open(image_path, "rb") as f:29            image_data = f.read()30        print(f"📤 Sending image ({len(image_data)} bytes)")31        await ws.send(image_data)32 33        # Listen for responses34        while True:35            raw = await ws.recv()36            msg = json.loads(raw)37            event = msg.get("event")38 39            if event == "status":40                print(f"  ℹ️  {msg['data']['message']}")41            elif event == "ocr_chunk":42                print(f"  📝 [{msg['index']}] {msg['data']}  (conf={msg.get('confidence', '?')})")43            elif event == "ocr_complete":44                print(f"  ✅ Complete – {msg['data']['total_lines']} lines extracted")45                break46            elif event == "error":47                print(f"  ❌ Error: {msg.get('message', msg)}")48                break49            else:50                print(f"  ❓ {event}: {msg}")51 52 53if __name__ == "__main__":54    if len(sys.argv) < 2:55        print("Usage: python test_ws_client.py <image_path>")56        sys.exit(1)57    asyncio.run(main(sys.argv[1]))58