diff --git a/src/nimbus_inference/cell_analyzer.py b/src/nimbus_inference/cell_analyzer.py
new file mode 100644
index 0000000..de0152a
--- /dev/null
+++ b/src/nimbus_inference/cell_analyzer.py
@@ -0,0 +1,303 @@
+import os
+import ipywidgets as widgets
+from IPython.core.display import HTML
+from IPython.display import display
+from io import BytesIO
+from skimage import io
+import numpy as np
+from natsort import natsorted
+from skimage.segmentation import find_boundaries
+
+
+class CellAnalyzer(object):
+    def __init__(
+            self, input_dir, cell_df, output_dir,
+            segmentation_naming_convention, img_width='1000px', context=None
+        ):
+        """Viewer for Nimbus application.
+        Args:
+            input_dir (str): Path to directory containing individual channels of multiplexed images
+            cell_df (pd.DataFrame): DataFrame with columns 'pixie_ct', 'nimbus_ct', 'fov', 'label'
+            output_dir (str): Path to save output annotations
+            segmentation_naming_convention (fn): Function that maps input path to segmentation path
+            img_width (str): Width of images in viewer.
+            context (int): area around the cell to display
+        """
+        self.image_width = img_width
+        self.input_dir = input_dir
+        self.segmentation_naming_convention = segmentation_naming_convention
+        self.cell_df = cell_df
+        if "annotations" not in self.cell_df.columns:
+            self.cell_df["annotations"] = [None]*len(self.cell_df)
+        self.cell_ids = self.cell_df["Cell ID"].values.tolist()
+        self.fov_names = self.cell_df["fov"].values
+        self.output_dir = output_dir
+        
+        self.context = context
+        display(HTML(
+            "<style>.bottom_spacing_class {margin-bottom: 10px;}</style>"
+        ))
+        display(HTML(
+            "<style>.left_spacing_class {margin-left: 60px;}</style>"
+        ))
+        display(HTML("<style>.font_size_class {line-height:30px;font-weight:500;font-size:200%;\
+                     margin-left: 30px;}</style>"))
+        self.update_button = widgets.Button(description="Update Image")
+        self.update_button.on_click(self.update_button_click)
+        self.update_button.add_class(
+              "left_spacing_class"            
+        )
+        self.overlay_checkbox = widgets.Checkbox(
+            value=True,
+            description='Overlay segmentations',
+            disabled=False
+        )
+        self.model_button = widgets.RadioButtons(
+            options=["Nimbus", "Pixie", "None"], description="Cell type",
+            style=dict(font_size="18px")
+        )
+        self.model_button.add_class("font_size_class")
+        
+        self.save_annotations_button = widgets.Button(description="Save and next cell")
+        self.save_annotations_button.on_click(self.save_annotations_button_click)
+        self.save_annotations_button.add_class("left_spacing_class")
+
+        self.cell_id_select = widgets.Select(
+            options=self.cell_ids,
+            description='Cell ID:',
+            disabled=False
+        )
+        self.cell_id_select.observe(self.cell_id_select_fn, names='value')
+        
+        self.red_select = widgets.Select(
+            options=[],
+            description='Red:',
+            disabled=False,
+            layout=widgets.Layout(height='120px')
+        )
+        self.green_select = widgets.Select(
+            options=[],
+            description='Green:',
+            disabled=False,
+            layout=widgets.Layout(height='120px')
+        )
+        self.blue_select = widgets.Select(
+            options=[],
+            description='Blue:',
+            disabled=False,
+            layout=widgets.Layout(height='120px')
+        )
+        self.input_image = widgets.Image()
+        self.output_image = widgets.Image()
+
+    def cell_id_select_fn(self, change):
+        """Selects fov to display.
+        Args:
+            change (dict): Change dictionary from ipywidgets.
+        """
+        fov = self.cell_df[self.cell_df["Cell ID"] == self.cell_id_select.value]["fov"].values[0]
+        fov_path = os.path.join(self.input_dir, fov)
+        channels = [
+            ch for ch in os.listdir(fov_path) if os.path.isfile(os.path.join(fov_path, ch))
+        ]
+        channels = [ch.split(".")[0] for ch in channels]
+        self.red_select.options = natsorted(channels)
+        self.green_select.options = natsorted(channels)
+        self.blue_select.options = natsorted(channels)
+        nimbus_ct = self.cell_df[
+            self.cell_df["Cell ID"] == self.cell_id_select.value
+        ]["nimbus_ct"].values[0]
+        pixie_ct = self.cell_df[
+            self.cell_df["Cell ID"] == self.cell_id_select.value
+        ]["pixie_ct"].values[0]
+        self.model_button.options = [nimbus_ct, pixie_ct, "None of the above"]
+
+    def create_composite_image(self, path_dict, add_boundaries=True):
+        """Creates composite image from input paths.
+        Args:
+            path_dict (dict): Dictionary of paths to images.
+            add_boundaries (bool): Whether to add boundaries to multiplex image.
+        Returns:
+            composite_image (np.array): Composite image.
+        """
+        for k in ["red", "green", "blue"]:
+            if k not in path_dict.keys():
+                path_dict[k] = None
+        output_image = []
+        for k, p in path_dict.items():
+            if p:
+                img = io.imread(p)
+                output_image.append(img)
+            else:
+                non_none = [p for p in path_dict.values() if p]
+                img = io.imread(non_none[0])
+                output_image.append(img*0)
+        # add overlay of instances
+        composite_image = np.stack(output_image, axis=-1)
+        if self.segmentation_naming_convention and add_boundaries:
+            fov_path = os.path.split(list(path_dict.values())[0])[0]
+            seg_path = self.segmentation_naming_convention(fov_path)
+            seg_img = io.imread(seg_path)
+            seg_boundaries = find_boundaries(seg_img, mode='inner')
+            individual_cell_boundary = seg_boundaries.copy()
+            cell_mask = seg_img == self.cell_df[
+                self.cell_df["Cell ID"] == self.cell_id_select.value
+            ]["label"].values[0]
+            individual_cell_boundary = np.logical_and(cell_mask, seg_boundaries)
+            val = (np.max(composite_image, axis=(0,1))*0.5).astype(composite_image.dtype)
+            val = np.min(val[val>0])
+        else:
+            seg_boundaries = None
+            individual_cell_boundary = None
+        return composite_image, seg_boundaries, individual_cell_boundary
+    
+    def create_instance_image(self, path_dict, cell_id):
+        """Creates composite image from input paths.
+        Args:
+            path_dict (dict): Dictionary of paths to images.
+            cell_id (int): id of cell to highlight
+        Returns:
+            composite_image (np.array): Composite image.
+        """
+        # add overlay of instances
+        fov_path = os.path.split(list(path_dict.values())[0])[0]
+        seg_path = self.segmentation_naming_convention(fov_path)
+        seg_img = io.imread(seg_path)
+        seg_boundaries = find_boundaries(seg_img, mode='inner')
+        seg_img[seg_boundaries] = 0
+        seg_img_clipped = np.clip(seg_img, 0, 1) * 0.2
+        seg_img_clipped[seg_img == cell_id] = 1
+        return seg_img_clipped
+
+    def layout(self):
+        """Creates layout for viewer."""
+        channel_selectors = widgets.VBox([
+            self.red_select,
+            self.green_select,
+            self.blue_select
+        ])
+        self.input_image.layout.width = self.image_width
+        self.output_image.layout.width = self.image_width
+        self.input_image.layout.height = self.image_width
+        self.output_image.layout.height = self.image_width
+
+        layout = widgets.HBox([
+            widgets.VBox([
+                self.cell_id_select,
+                channel_selectors,
+                self.overlay_checkbox,
+                self.update_button,
+            widgets.VBox([
+                self.model_button,
+                self.save_annotations_button
+            ]),
+            ], layout=widgets.Layout(width='30%')),
+        widgets.HBox([
+            # input_html,
+            self.input_image
+        ]),
+        widgets.HBox([
+            # output_html,
+            self.output_image
+        ])
+        ])
+        display(layout)
+
+    def search_for_similar(self, select_value):
+        """Searches for similar filename in input directory.
+        Args:
+            select_value (str): Filename to search for.
+        Returns:
+            similar_path (str): Path to similar filename.
+        """
+        fov = self.cell_df[self.cell_df["Cell ID"] == self.cell_id_select.value]["fov"].values[0]
+        in_f_path = os.path.join(self.input_dir, fov)
+        # search for similar filename in in_f_path
+        in_f_files = [
+            f for f in os.listdir(in_f_path) if os.path.isfile(os.path.join(in_f_path, f))
+        ]
+        similar_path = None
+        for f in in_f_files:
+            if select_value + "." in f:
+                similar_path = os.path.join(self.input_dir, fov, f)
+        return similar_path
+
+    def update_img(self, image_viewer, composite_image):
+        """Updates image in viewer by saving it as png and loading it with the viewer widget.
+        Args:
+            image_viewer (ipywidgets.Image): Image widget to update.
+            composite_image (np.array): Composite image to display.
+        """
+        # Convert composite image to bytes and assign it to the output_image widget
+        with BytesIO() as output_buffer:
+            io.imsave(output_buffer, composite_image, format="png")
+            output_buffer.seek(0)
+            image_viewer.value = output_buffer.read()
+
+    def update_composite(self):
+        """Updates composite image in viewer."""
+        in_path_dict = {
+            "red": None,
+            "green": None,
+            "blue": None
+        }
+        if self.red_select.value:
+            in_path_dict["red"] = self.search_for_similar(self.red_select.value)
+        if self.green_select.value:
+            in_path_dict["green"] = self.search_for_similar(self.green_select.value)
+        if self.blue_select.value:
+            in_path_dict["blue"] = self.search_for_similar(self.blue_select.value)
+        non_none = [p for p in in_path_dict.values() if p]
+        if not non_none:
+            return
+        in_composite_image, seg_boundaries, individual_cell_boundary = self.create_composite_image(
+            in_path_dict, add_boundaries=self.overlay_checkbox.value
+        )
+        in_composite_image = in_composite_image / np.quantile(
+            in_composite_image, 0.999, axis=(0,1)
+        )
+        in_composite_image = np.clip(in_composite_image*255, 0, 255).astype(np.uint8)
+        if seg_boundaries is not None:
+            in_composite_image[seg_boundaries] = [75, 75, 75]
+            in_composite_image[individual_cell_boundary] = [175, 175, 175]
+        cell_label = self.cell_df[
+            self.cell_df["Cell ID"] == self.cell_id_select.value
+        ]["label"].values[0]
+        seg_image = self.create_instance_image(in_path_dict, cell_label)
+        seg_image = (seg_image * 255).astype(np.uint8)
+        if self.context:
+            h, w = np.where(seg_image == seg_image.max())
+            h, w = np.mean(h).astype(np.uint16), np.mean(w).astype(np.uint16)
+            h_max, w_max = seg_image.shape
+            h_0, h_1 = np.clip(h-self.context, 0, h_max), np.clip(h+self.context, 0, h_max)
+            w_0, w_1 = np.clip(w-self.context, 0, w_max), np.clip(w+self.context, 0, w_max)
+            seg_image = seg_image[h_0:h_1, w_0:w_1]
+            in_composite_image = in_composite_image[h_0:h_1, w_0:w_1]
+        # update image viewers
+        self.update_img(self.input_image, in_composite_image)
+        self.update_img(self.output_image, seg_image)
+
+    def update_button_click(self, button):
+        """Updates composite image in viewer when update button is clicked."""
+        self.update_composite()
+    
+    def save_annotations_button_click(self, button):
+        """Updates composite image in viewer when pixie button is clicked."""
+        self.cell_df.loc[
+            self.cell_df["Cell ID"] == self.cell_id_select.value, "annotations"
+        ] = self.model_button.value
+        self.cell_df.to_csv(
+            os.path.join(self.output_dir, "annotations.csv"), index=False
+        )
+        # bump to next cell
+        current_index = self.cell_id_select.options.index(self.cell_id_select.value)
+        max_index = len(self.cell_id_select.options)
+        if current_index + 1 < max_index:
+            self.cell_id_select.value = self.cell_id_select.options[current_index + 1]
+        self.update_composite()
+
+    def display(self):
+        """Displays viewer."""
+        self.cell_id_select_fn(None)
+        self.layout()
+        self.update_composite() 
\ No newline at end of file
diff --git a/src/nimbus_inference/viewer_widget.py b/src/nimbus_inference/viewer_widget.py
index eef7899..e406e79 100644
--- a/src/nimbus_inference/viewer_widget.py
+++ b/src/nimbus_inference/viewer_widget.py
@@ -6,24 +6,34 @@
 from copy import copy
 import numpy as np
 from natsort import natsorted
+from skimage.segmentation import find_boundaries
 
 
 class NimbusViewer(object):
-    def __init__(self, input_dir, output_dir, img_width='600px'):
+    def __init__(
+            self, input_dir, output_dir, segmentation_naming_convention=None, img_width='600px'
+        ):
         """Viewer for Nimbus application.
         Args:
             input_dir (str): Path to directory containing individual channels of multiplexed images
             output_dir (str): Path to directory containing output of Nimbus application.
+            segmentation_naming_convention (fn): Function that maps input path to segmentation path
             img_width (str): Width of images in viewer.
         """
         self.image_width = img_width
         self.input_dir = input_dir
         self.output_dir = output_dir
+        self.segmentation_naming_convention = segmentation_naming_convention
         self.fov_names = [os.path.basename(p) for p in os.listdir(output_dir) if \
                           os.path.isdir(os.path.join(output_dir, p))]
         self.fov_names = natsorted(self.fov_names)
         self.update_button = widgets.Button(description="Update Image")
         self.update_button.on_click(self.update_button_click)
+        self.overlay_checkbox = widgets.Checkbox(
+            value=True,
+            description='Overlay segmentations',
+            disabled=False
+        )
 
         self.fov_select = widgets.Select(
             options=self.fov_names,
@@ -63,7 +73,7 @@ def select_fov(self, change):
         self.green_select.options = natsorted(channels)
         self.blue_select.options = natsorted(channels)
 
-    def create_composite_image(self, path_dict):
+    def create_composite_image(self, path_dict, add_overlay=True, add_boundaries=False):
         """Creates composite image from input paths.
         Args:
             path_dict (dict): Dictionary of paths to images.
@@ -82,9 +92,31 @@ def create_composite_image(self, path_dict):
                 non_none = [p for p in path_dict.values() if p]
                 img = io.imread(non_none[0])
                 output_image.append(img*0)
-
+        # add overlay of instances
         composite_image = np.stack(output_image, axis=-1)
-        return composite_image
+        if self.segmentation_naming_convention and add_overlay:
+            fov_path = os.path.split(list(path_dict.values())[0])[0]
+            seg_path = self.segmentation_naming_convention(fov_path)
+            seg_img = io.imread(seg_path)
+            seg_boundaries = find_boundaries(seg_img, mode='inner')
+            seg_img[seg_boundaries] = 0
+            seg_img = np.clip(seg_img, 0, 1)
+            seg_img = np.repeat(seg_img[..., np.newaxis], 3, axis=-1) * np.max(composite_image)
+            background_mask = composite_image < np.max(composite_image) * 0.2
+            composite_image[background_mask] += (seg_img[background_mask] * 0.2).astype(
+                composite_image.dtype
+            )
+        elif self.segmentation_naming_convention and add_boundaries:
+            fov_path = os.path.split(list(path_dict.values())[0])[0]
+            seg_path = self.segmentation_naming_convention(fov_path)
+            seg_img = io.imread(seg_path)
+            seg_boundaries = find_boundaries(seg_img, mode='inner')
+            val = (np.max(composite_image, axis=(0,1))*0.5).astype(composite_image.dtype)
+            val = np.min(val[val>0])
+            composite_image[seg_boundaries] = [val]*3 
+        else:
+            seg_boundaries = None
+        return composite_image, seg_boundaries
 
     def layout(self):
         """Creates layout for viewer."""
@@ -95,6 +127,8 @@ def layout(self):
         ])
         self.input_image.layout.width = self.image_width
         self.output_image.layout.width = self.image_width
+        self.input_image.layout.height = self.image_width
+        self.output_image.layout.height = self.image_width
         viewer_html = widgets.HTML("<h2>Select files</h2>")
         input_html = widgets.HTML("<h2>Input</h2>")
         output_html = widgets.HTML("<h2>Nimbus Output</h2>")
@@ -104,6 +138,7 @@ def layout(self):
                 viewer_html,
                 self.fov_select,
                 channel_selectors,
+                self.overlay_checkbox,
                 self.update_button
             ]),
         widgets.VBox([
@@ -173,12 +208,16 @@ def update_composite(self):
         non_none = [p for p in path_dict.values() if p]
         if not non_none:
             return
-        composite_image = self.create_composite_image(path_dict)
-        in_composite_image = self.create_composite_image(in_path_dict)
+        composite_image, _ = self.create_composite_image(path_dict)
+        in_composite_image, seg_boundaries = self.create_composite_image(
+            in_path_dict, add_overlay=False, add_boundaries=self.overlay_checkbox.value
+        )
         in_composite_image = in_composite_image / np.quantile(
             in_composite_image, 0.999, axis=(0,1)
         )
         in_composite_image = np.clip(in_composite_image*255, 0, 255).astype(np.uint8)
+        if seg_boundaries is not None:
+            in_composite_image[seg_boundaries] = [127, 127, 127]
         # update image viewers
         self.update_img(self.input_image, in_composite_image)
         self.update_img(self.output_image, composite_image)
diff --git a/templates/3_analyze_individual_cells.ipynb b/templates/3_analyze_individual_cells.ipynb
new file mode 100644
index 0000000..3b3d3e1
--- /dev/null
+++ b/templates/3_analyze_individual_cells.ipynb
@@ -0,0 +1,118 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Cell analysis notebook\n",
+    "- Extract cells where the pure pixie annotation and the nimbus-pixie annotation disagree.\n",
+    "- Build a viewer widget with channels on the left and an instance map on the right with all cells in gray, except the one in question.\n",
+    "- Add a dropdown on the left to select the cell in question.\n",
+    "- Add a dropdown with the correct celltype and a button to submit the correction."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from IPython.display import display, HTML\n",
+    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))\n",
+    "from nimbus_inference.nimbus import prep_naming_convention\n",
+    "from nimbus_inference.cell_analyzer import CellAnalyzer\n",
+    "import pandas as pd\n",
+    "import os"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Get cells from confusion matrix off-diagonal\n",
+    "Prepare a dataframe that contains the following columns:\n",
+    "- `Cell ID`: arbitrary and unique cell id\n",
+    "- `pixie_ct`: the cell type from the pixie annotation\n",
+    "- `nimbus_ct`: the cell type from the nimbus-pixie annotation\n",
+    "- `fov`: the fov of the cell\n",
+    "- `label`: the instance label of the cell within the fov\n",
+    "\n",
+    "Change `base_dir` to your local data path. Below `base_dir` we expect the following folder structure:\n",
+    "```\n",
+    "|-- base_dir\n",
+    "|   |-- image_data\n",
+    "|   |   |-- fov_1\n",
+    "|   |   |-- fov_2\n",
+    "|   |-- segmentation\n",
+    "|   |   |-- deepcell_output\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "off_diagonal_df = pd.DataFrame(\n",
+    "    {\n",
+    "     \"Cell ID\": [1, 2, 3, 4],\n",
+    "     \"pixie_ct\": [\"Eosinophil\", \"Neutrophil\", \"Lymphocyte\", \"Monocyte\"],\n",
+    "     \"nimbus_ct\": [\"Neutrophil\", \"Lymphocyte\", \"Monocyte\", \"Eosinophil\"],\n",
+    "     \"fov\": [\"TMA32_R4C6\", \"TMA32_R4C6\", \"TMA32_R5C4\", \"TMA32_R5C4\"],\n",
+    "     \"label\": [451, 123, 234, 345],\n",
+    "     } \n",
+    ")\n",
+    "base_dir = os.path.normpath(\"C:/Users/lorenz/Desktop/angelo_lab/data/SPAIN_TNBC_fov_subset\")\n",
+    "tiff_dir = os.path.join(base_dir, \"image_data\")\n",
+    "deepcell_output_dir = os.path.join(base_dir, \"segmentation\", \"deepcell_output\")\n",
+    "segmentation_naming_convention = prep_naming_convention(deepcell_output_dir)\n",
+    "\n",
+    "off_diagonal_df"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Use the CellAnalyzer to annotate cells"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "viewer = CellAnalyzer(input_dir=tiff_dir,\n",
+    "                      cell_df=off_diagonal_df,\n",
+    "                      output_dir=base_dir,\n",
+    "                      segmentation_naming_convention=segmentation_naming_convention,\n",
+    "                      img_width='600px',\n",
+    "                      context=200)\n",
+    "viewer.display()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/tests/test_cell_analyzer.py b/tests/test_cell_analyzer.py
new file mode 100644
index 0000000..e34d96d
--- /dev/null
+++ b/tests/test_cell_analyzer.py
@@ -0,0 +1,65 @@
+from nimbus_inference.cell_analyzer import CellAnalyzer
+from nimbus_inference.nimbus import Nimbus, prep_naming_convention
+from tests.test_utils import prepare_ome_tif_data, prepare_tif_data
+import pandas as pd
+import numpy as np
+import tempfile
+import os
+
+cell_df = pd.DataFrame(
+    {
+     "Cell ID": [1, 2, 3, 4],
+     "pixie_ct": ["Eosinophil", "Neutrophil", "Lymphocyte", "Monocyte"],
+     "nimbus_ct": ["Neutrophil", "Lymphocyte", "Monocyte", "Eosinophil"],
+     "fov": ["fov_0", "fov_0", "fov_1", "fov_1"],
+     "label": [1, 2, 3, 4],
+     } 
+)
+
+def test_CellAnalyzer():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        _ = prepare_tif_data(
+            num_samples=2, temp_dir=temp_dir, selected_markers=["CD4", "CD11c", "CD56"]
+        )
+        naming_convention = prep_naming_convention(os.path.join(temp_dir, "deepcell_output"))
+
+        viewer_widget = CellAnalyzer(
+            temp_dir, cell_df=cell_df, segmentation_naming_convention=naming_convention,
+            output_dir=temp_dir
+        )
+        assert isinstance(viewer_widget, CellAnalyzer)
+
+
+def test_composite_image():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        _ = prepare_tif_data(
+            num_samples=2, temp_dir=temp_dir, selected_markers=["CD4", "CD11c", "CD56"]
+        )
+        naming_convention = prep_naming_convention(os.path.join(temp_dir, "deepcell_output"))
+        viewer_widget = CellAnalyzer(
+            temp_dir, cell_df=cell_df, segmentation_naming_convention=naming_convention,
+            output_dir=temp_dir
+        )
+        path_dict = {
+            "red": os.path.join(temp_dir, "fov_0", "CD4.tiff"),
+            "green": os.path.join(temp_dir, "fov_0", "CD11c.tiff"),
+        }
+        composite_image, _, _ = viewer_widget.create_composite_image(path_dict)
+        assert isinstance(composite_image, np.ndarray)
+        assert composite_image.shape == (256, 256, 3)
+
+        path_dict["blue"] = os.path.join(temp_dir, "fov_0", "CD56.tiff")
+        composite_image, _, _ = viewer_widget.create_composite_image(path_dict)
+        assert composite_image.shape == (256, 256, 3)
+        # test if segmentation gets added
+        viewer_widget = CellAnalyzer(
+            temp_dir, cell_df=cell_df, segmentation_naming_convention=naming_convention,
+            output_dir=temp_dir
+        )
+        composite_image, seg_boundaries, cell_boundaries = viewer_widget.create_composite_image(
+            path_dict
+        )
+        assert composite_image.shape == (256, 256, 3)
+        assert seg_boundaries.shape == (256, 256)
+        assert cell_boundaries.shape == (256, 256)
+        assert np.unique(seg_boundaries).tolist() == [0, 1]
diff --git a/tests/test_viewer_widget.py b/tests/test_viewer_widget.py
index 163271c..9e37108 100644
--- a/tests/test_viewer_widget.py
+++ b/tests/test_viewer_widget.py
@@ -1,4 +1,5 @@
 from nimbus_inference.viewer_widget import NimbusViewer
+from nimbus_inference.nimbus import Nimbus, prep_naming_convention
 from tests.test_utils import prepare_ome_tif_data, prepare_tif_data
 import numpy as np
 import tempfile
@@ -24,10 +25,19 @@ def test_composite_image():
             "red": os.path.join(temp_dir, "fov_0", "CD4.tiff"),
             "green": os.path.join(temp_dir, "fov_0", "CD11c.tiff"),
         }
-        composite_image = viewer_widget.create_composite_image(path_dict)
+        composite_image, _ = viewer_widget.create_composite_image(path_dict)
         assert isinstance(composite_image, np.ndarray)
         assert composite_image.shape == (256, 256, 3)
 
         path_dict["blue"] = os.path.join(temp_dir, "fov_0", "CD56.tiff")
-        composite_image = viewer_widget.create_composite_image(path_dict)
-        assert composite_image.shape == (256, 256, 3)
\ No newline at end of file
+        composite_image, _ = viewer_widget.create_composite_image(path_dict)
+        assert composite_image.shape == (256, 256, 3)
+        # test if segmentation gets added
+        naming_convention = prep_naming_convention(os.path.join(temp_dir, "deepcell_output"))
+        viewer_widget = NimbusViewer(
+            temp_dir, temp_dir, segmentation_naming_convention=naming_convention
+        )
+        composite_image, seg_boundaries = viewer_widget.create_composite_image(path_dict)
+        assert composite_image.shape == (256, 256, 3)
+        assert seg_boundaries.shape == (256, 256)
+        assert np.unique(seg_boundaries).tolist() == [0, 1]