diff --git a/callbacks/image_viewer.py b/callbacks/image_viewer.py
index eb9abe0..e8e14fe 100644
--- a/callbacks/image_viewer.py
+++ b/callbacks/image_viewer.py
@@ -125,24 +125,7 @@ def render_image(
         tf = np.zeros((500, 500))
     fig = px.imshow(tf, binary_string=True)
     if toggle_seg_result:
-        unique_segmentation_values = np.unique(result)
-        normalized_range = np.linspace(
-            0, 1, len(unique_segmentation_values)
-        )  # heatmap requires a normalized range
-        color_list = (
-            px.colors.qualitative.Plotly
-        )  # TODO placeholder - replace with user defined classess
-        colorscale = [
-            [normalized_range[i], color_list[i % len(color_list)]]
-            for i in range(len(unique_segmentation_values))
-        ]
-        fig.add_trace(
-            go.Heatmap(
-                z=result,
-                colorscale=colorscale,
-                showscale=False,
-            )
-        )
+        fig.add_trace(go.Heatmap(z=result, showscale=False, colorscale=None))
         fig["data"][1]["opacity"] = opacity / 100
 
     fig.update_layout(