diff --git a/src/config.py b/src/config.py
index 699829b..05614ca 100644
--- a/src/config.py
+++ b/src/config.py
@@ -34,3 +34,6 @@ FRAME_FORMAT = "HT"
 
 CENTRAL_FREQUENCY_HZ = CENTRAL_FREQUENCY_MHZ * 1_000_000
 C = 299_792_458  # m/s
+
+VISUALISE_RAW = False
+HEATMAP_FPS = 10
diff --git a/src/visualise/__init__.py b/src/visualise/__init__.py
index 9e97267..4ce0db5 100644
--- a/src/visualise/__init__.py
+++ b/src/visualise/__init__.py
@@ -4,13 +4,14 @@ import numpy as np
 import numpy.typing as npt
 from simple_websocket import Server
 import time
+from datetime import datetime
 
 import matplotlib.pyplot as plt
 import io
-from PIL import Image
 import logging
 
 from ..aoa import AoA
+from .. import config
 
 import matplotlib
 
@@ -48,23 +49,24 @@ def get_data(sock: Subscriber):
 def add_data(
     raw_data: npt.NDArray[np.complex128], new_data: npt.NDArray[np.complex128]
 ):
-    # magn = np.abs(raw_data)
-    # phase = np.angle(raw_data)
+    if config.VISUALISE_RAW:
+        magn = np.abs(raw_data)
+        phase = np.angle(raw_data)
 
-    # new_mag = np.abs(new_data)
-    # new_phase = np.angle(new_data)
+        new_mag = np.abs(new_data)
+        new_phase = np.angle(new_data)
 
-    # fig, axs = plt.subplots(2, 2)
-    # axs[0, 0].plot(magn[:, 0, 0], c="b")
-    # axs[0, 0].plot(magn[:, 1, 0], c="orange")
-    # axs[1, 0].plot(new_mag[:, 0, 0], c="b")
-    # axs[1, 0].plot(new_mag[:, 1, 0], c="orange")
-    # axs[0, 1].plot(phase[:, 0, 0], c="b")
-    # axs[0, 1].plot(phase[:, 1, 0], c="orange")
-    # axs[1, 1].plot(new_phase[:, 0, 0], c="b")
-    # axs[1, 1].plot(new_phase[:, 1, 0], c="orange")
-    # fig.savefig("/tmp/plot.png")
-    # plt.close(fig)
+        fig, axs = plt.subplots(2, 2)
+        axs[0, 0].plot(magn[:, 0, 0], c="b")
+        axs[0, 0].plot(magn[:, 1, 0], c="orange")
+        axs[1, 0].plot(new_mag[:, 0, 0], c="b")
+        axs[1, 0].plot(new_mag[:, 1, 0], c="orange")
+        axs[0, 1].plot(phase[:, 0, 0], c="b")
+        axs[0, 1].plot(phase[:, 1, 0], c="orange")
+        axs[1, 1].plot(new_phase[:, 0, 0], c="b")
+        axs[1, 1].plot(new_phase[:, 1, 0], c="orange")
+        fig.savefig("/tmp/plot.png")
+        plt.close(fig)
 
     global data
     if data.size == 0:
@@ -108,9 +110,12 @@ def make_heatmap(max_tof: float):
 
 
 def gather_aoa(max_tof: float):
+    prev_frame = datetime.now()
     while True:
-        # time.sleep(0.05)
-        logger.info("Got AoA heatmap")
+        while (datetime.now() - prev_frame).total_seconds() < 1 / config.HEATMAP_FPS:
+            time.sleep(0.01)
+            pass
+        logger.debug("Generating heatmap")
         buf = make_heatmap(max_tof)
         yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + buf.read() + b"\r\n")
         buf.close()
@@ -121,7 +126,7 @@ def aoa_tof():
     max_tof_str = request.args.get("max_tof")
     try:
         max_tof = float(max_tof_str)
-    except Exception as e:
+    except Exception:
         max_tof = 5e-8
     return Response(
         gather_aoa(max_tof), mimetype="multipart/x-mixed-replace; boundary=frame"