-
Notifications
You must be signed in to change notification settings - Fork 0
/
visualization.py
42 lines (31 loc) · 970 Bytes
/
visualization.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
import json
import threading
import websocket
import matplotlib.pyplot as plt
# WebSocket server URL
WS_URL = "ws://localhost:3000"
regions = []
num_of_vehicles = []
def on_message(ws, message):
data = json.loads(message)
region = data["partition"]
vehicles = data["numOfVehicles"]
update_visualization(region, vehicles)
def update_visualization(region, vehicles):
if region in regions:
idx = regions.index(region)
num_of_vehicles[idx] = vehicles
else:
regions.append(region)
num_of_vehicles.append(vehicles)
plt.bar(regions, num_of_vehicles, color='blue')
plt.xlabel('Region')
plt.ylabel('Number of Vehicles')
plt.title('Real-time Traffic Visualization')
plt.xticks(regions)
plt.show()
def start_websocket():
ws = websocket.WebSocketApp(WS_URL, on_message=on_message)
ws.run_forever()
websocket_thread = threading.Thread(target=start_websocket)
websocket_thread.start()