diff --git a/where_fi/cli/__init__.py b/where_fi/cli/__init__.py
index 52d993c..e03d808 100644
--- a/where_fi/cli/__init__.py
+++ b/where_fi/cli/__init__.py
@@ -52,7 +52,7 @@ def heatmap() -> None:
     def callback(antenna_data: npt.NDArray[np.complex64]) -> None:
         logger.info(f"Got final CSI data with shape {antenna_data.shape}")
         visualise_data(antenna_data, visualise.DataType.RAW_CSI)
-        processed = preprocessor.preprocess(antenna_data)
+        processed = preprocessor.preprocess(antenna_data, visualiser=visualise_data)
         visualise_data(processed, visualise.DataType.PROCESSED_CSI)
         logger.info(f"Processed CSI data with shape {processed.shape}")
         processed_tensor = torch.tensor(processed, device=device)
diff --git a/where_fi/visualise/Makefile b/where_fi/visualise/Makefile
index 978f494..a359e89 100644
--- a/where_fi/visualise/Makefile
+++ b/where_fi/visualise/Makefile
@@ -6,4 +6,5 @@ server/generated: protos/*.proto
 	find protos/ -type f -name "*.proto" | xargs uv run protol --create-package --in-place --python-out server/generated/  protoc --proto-path=protos/
 
 frontend/src/grpc: protos/*.proto
-	find protos/ -name "*.proto" | xargs npm exec --prefix frontend/ protoc --ts_out frontend/src/grpc -I protos/
+	mkdir -p frontend/src/grpc
+	cd frontend && find ../protos/ -name "*.proto" | xargs npx protoc --ts_out=src/grpc -I../protos/
diff --git a/where_fi/visualise/frontend/src/components/Figure.vue b/where_fi/visualise/frontend/src/components/Figure.vue
index 3ec605f..aa42ae7 100644
--- a/where_fi/visualise/frontend/src/components/Figure.vue
+++ b/where_fi/visualise/frontend/src/components/Figure.vue
@@ -11,6 +11,18 @@ import { computed } from 'vue'
 const { figure, paused = false } = defineProps<{ figure: Figure; paused?: boolean }>()
 const cancel = ref<boolean>(false)
 
+function reshape(data: number[], width: number) {
+  if (data.length % width != 0) {
+    throw new Error('Data length is not divisible by width')
+  }
+  const height = data.length / width
+  const result: number[][] = new Array(height)
+  for (let i = 0; i < height; i++) {
+    result[i] = data.slice(i * width, (i + 1) * width)
+  }
+  return result
+}
+
 const data = ref<FigureData | null>(null)
 const plotData = computed(() => {
   if (data.value) {
@@ -24,16 +36,12 @@ const plotData = computed(() => {
           color: line.color ? line.color : undefined,
         }))
       case 'heatmap':
-        return []
-      /* data.value.figure.heatmap.lines.map((line) => {
-          return {
-            x: line.x,
-            y: line.y,
-            type: 'scatter',
-            name: line.label,
-            color: line.color,
-          }
-        })*/
+        return [
+          {
+            z: reshape(data.value.figure.heatmap.data, data.value.figure.heatmap.width),
+            type: 'heatmap' as const,
+          },
+        ]
     }
   }
   return []
@@ -43,6 +51,7 @@ const layout = {
   title: { text: figure.title },
   xaxis: { title: { text: figure.xLabel } },
   yaxis: { title: { text: figure.yLabel } },
+  height: 700,
 }
 
 function updateGraph() {
diff --git a/where_fi/visualise/protos/figure_type/heatmap.proto b/where_fi/visualise/protos/figure_type/heatmap.proto
index cfd8a38..9a8be5a 100644
--- a/where_fi/visualise/protos/figure_type/heatmap.proto
+++ b/where_fi/visualise/protos/figure_type/heatmap.proto
@@ -2,9 +2,8 @@ syntax = "proto3";
 
 message HeatmapData {
     string uuid = 1;
-    repeated float x = 2;
-    repeated float y = 3;
-    uint32 width = 4;
-    uint32 height = 5;
-    string cmap = 6;
+    repeated float data = 2;
+    uint32 width = 3;
+    uint32 height = 4;
+    string cmap = 5;
 }
diff --git a/where_fi/visualise/server/__init__.py b/where_fi/visualise/server/__init__.py
index d428bab..75031dc 100644
--- a/where_fi/visualise/server/__init__.py
+++ b/where_fi/visualise/server/__init__.py
@@ -23,8 +23,9 @@ clients_lock = threading.Lock()
 
 class DataType(Enum):
     RAW_CSI = 1
-    PROCESSED_CSI = 2
-    HEATMAP = 3
+    UNWRAPPED_PHASE = 2
+    PROCESSED_CSI = 3
+    HEATMAP = 4
 
 
 @dataclass
@@ -40,6 +41,12 @@ figures = {
         x_label="Subcarrier",
         y_label="Phase",
     ),
+    "unwrapped_phase": figure_pb2.Figure(
+        uuid=str(uuid.uuid4()),
+        title="Unwrapped CSI Phase",
+        x_label="Subcarrier",
+        y_label="Unwrapped phase",
+    ),
     "processed_phase": figure_pb2.Figure(
         uuid=str(uuid.uuid4()),
         title="Preprocessed CSI Phase",
@@ -48,9 +55,9 @@ figures = {
     ),
     "aoa_heatmap": figure_pb2.Figure(
         uuid=str(uuid.uuid4()),
-        title="Preprocessed CSI Phase",
-        x_label="Subcarrier",
-        y_label="Phase",
+        title="AoA Heatmap",
+        x_label="ToF",
+        y_label="AoA",
     ),
 }
 
@@ -104,7 +111,25 @@ def add_data(dtype: DataType, new_data: npt.NDArray[np.complex128]) -> None:
             for q in clients.get(uuid, []):
                 q.put(figure_pb2.FigureData(uuid=uuid, line=linechart))
         case DataType.HEATMAP:
+            uuid = figures["aoa_heatmap"].uuid
+            heatmap = heatmap_pb2.HeatmapData(
+                uuid=uuid,
+                data=new_data.flatten(),
+                width=new_data.shape[1],
+                height=new_data.shape[0],
+            )
+            for q in clients.get(uuid, []):
+                q.put(figure_pb2.FigureData(uuid=uuid, heatmap=heatmap))
             pass
+        case DataType.UNWRAPPED_PHASE:
+            uuid = figures["unwrapped_phase"].uuid
+            lines = [
+                line_pb2.LineChartData.Line(y=new_data[:, i, 0], label=f"Antenna {i}")
+                for i in range(new_data.shape[1])
+            ]
+            linechart = line_pb2.LineChartData(lines=lines)
+            for q in clients.get(uuid, []):
+                q.put(figure_pb2.FigureData(uuid=uuid, line=linechart))
 
 
 def listen_for_data(data_queue: "mp.Queue[VisualiserData]") -> None: