diff --git a/src/visualise/__init__.py b/src/visualise/__init__.py
index 750597c..9407b6e 100644
--- a/src/visualise/__init__.py
+++ b/src/visualise/__init__.py
@@ -1,20 +1,20 @@
-from flask import Flask, render_template, Response, request
-from flask_sock import Sock
-import numpy as np
-import numpy.typing as npt
-from simple_websocket import Server
-import time
-from datetime import datetime
-import multiprocessing as mp
-
-import matplotlib.pyplot as plt
 import io
 import logging
-
-from ..processing.aoa import AoA
-from .. import config
+import multiprocessing as mp
+import time
+from datetime import datetime
+from typing import Generator
 
 import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import numpy.typing as npt
+from flask import Flask, Response, render_template, request
+from flask_sock import Sock
+from simple_websocket import Server
+
+from .. import config
+from ..processing.aoa import AoA
 
 matplotlib.use("agg")
 
@@ -30,7 +30,7 @@ aoa: AoA = AoA()
 
 
 @app.route("/preprocessed")
-def preprocessed():
+def preprocessed() -> str:
     return render_template("preprocessed.html")
 
 
@@ -39,7 +39,7 @@ subscriber_settings: dict[Subscriber, tuple[int, int, int]] = {}
 
 
 @sock.route("/data")
-def get_data(sock: Subscriber):
+def get_data(sock: Subscriber) -> None:
     while True:
         msg = sock.receive()
         if len(msg.split()) != 3:
@@ -50,7 +50,7 @@ def get_data(sock: Subscriber):
 
 def add_data(
     raw_data: npt.NDArray[np.complex128], new_data: npt.NDArray[np.complex128]
-):
+) -> None:
     if config.VISUALISE_RAW:
         magn = np.abs(raw_data)
         phase = np.angle(raw_data)
@@ -92,7 +92,7 @@ def add_data(
         del subscriber_settings[subscriber]
 
 
-def make_heatmap(aoa: AoA, max_tof: float):
+def make_heatmap(aoa: AoA, max_tof: float) -> io.BytesIO:
     logger.info(f"Making heatmap with aoa of {aoa.timestamp}")
     fig = plt.figure()
     ax = fig.add_axes([0, 0, 1, 1], polar=True)
@@ -112,7 +112,7 @@ def make_heatmap(aoa: AoA, max_tof: float):
     return buf
 
 
-def gather_aoa(max_tof: float):
+def gather_aoa(max_tof: float) -> Generator[bytes, None, None]:
     assert aoa_queue is not None
 
     prev_frame = datetime.now()
@@ -130,7 +130,7 @@ def gather_aoa(max_tof: float):
 
 
 @app.route("/aoa_tof")
-def aoa_tof():
+def aoa_tof() -> Response:
     max_tof_str = request.args.get("max_tof")
     try:
         max_tof = float(max_tof_str)
@@ -141,8 +141,9 @@ def aoa_tof():
     )
 
 
-def start(conn: "mp.Queue[AoA]"):
-    global aoa_queue
+def start(conn: "mp.Queue[AoA]") -> None:
+    global app, aoa_queue
 
     aoa_queue = conn
     app.run(debug=True, use_reloader=False, host="0.0.0.0")
+    logger.info("Visualisation server shut down")