diff --git a/src/client/dcp_client/gui/_filesystem_wig.py b/src/client/dcp_client/gui/_filesystem_wig.py
index dbade63..add73f9 100644
--- a/src/client/dcp_client/gui/_filesystem_wig.py
+++ b/src/client/dcp_client/gui/_filesystem_wig.py
@@ -1,7 +1,9 @@
+import os
 import numpy as np
 from skimage.io import imread
 from skimage.color import label2rgb
 
+
 from PyQt5.QtWidgets import QFileSystemModel
 from PyQt5.QtCore import Qt, QVariant, QDir
 from PyQt5.QtGui import QImage
@@ -71,7 +73,7 @@ def data(self, index, role=Qt.DisplayRole):
             if filepath_img.endswith(settings.accepted_types):
                 
                 # if a mask make sure it is displayed properly
-                if "_seg" in filepath_img:
+                if "_seg" in filepath_img and os.path.exists(filepath_img):
                     img = imread(filepath_img)
                     if img.ndim > 2: img = img[0]
                     img = label2rgb(img)
diff --git a/src/client/dcp_client/gui/napari_window.py b/src/client/dcp_client/gui/napari_window.py
index c1e6ab8..5739c32 100644
--- a/src/client/dcp_client/gui/napari_window.py
+++ b/src/client/dcp_client/gui/napari_window.py
@@ -32,8 +32,8 @@ def __init__(self, app: Application) -> None:
         self.app = app
         self.setWindowTitle("napari viewer")
         self.setStyleSheet("background-color: #262930;")
-        #screen_size = QGuiApplication.primaryScreen().geometry()
-        #self.resize(int(screen_size.width()*0.9), int(screen_size.height()*0.8))
+        screen_size = QGuiApplication.primaryScreen().geometry()
+        self.resize(int(screen_size.width()*0.8), int(screen_size.height()*0.8))
 
         # Load image and get corresponding segmentation filenames
         img = self.app.load_image()
@@ -203,16 +203,25 @@ def on_remove_from_dataset_button_clicked(self) -> None:
         """
         Defines what happens when the "Remove from dataset" button is clicked.
         """
-        seg_name_to_remove = self.viewer.layers.selection.active.name
-        if seg_name_to_remove:
-            # Delete the image and corresponding masks from the dataset
-            image_name = self.app.cur_selected_img
-            seg_files_to_remove = [seg_name_to_remove + '.tiff']
-            self.app.delete_images([image_name] + seg_files_to_remove)
+        '''
+        try:
+            # get image name
+            files_to_remove = [self.viewer.layers.selection.active.name]
+        except AttributeError:
+            message_text = "Please first select the image in the layer list."
+            _ = self.create_warning_box(message_text, message_title="Warning")
+            return
+        '''
+        rmv_files = [self.app.cur_selected_img]
+        self.app.search_segs()
+        rmv_files.extend(self.app.seg_filepaths)
+        # Delete the image and corresponding masks from the dataset
+        self.app.delete_images(rmv_files)
+        self.app.cur_selected_img = ""
+        self.app.seg_filepaths = []
+        self.viewer.close()
+        self.close()
 
-            self.viewer.close()
-            self.close()
-    
     def set_editable_mask(self) -> None:
         """
         This function is not implemented. In theory the use can choose between which mask to edit.
@@ -241,6 +250,9 @@ def axis_changed(self, event) -> None:
         Is triggered each time the user switches the viewer between the mask channels. At this point the class mask
         needs to be updated according to the changes made tot the instance segmentation mask.
         """
+
+        if self.app.cur_selected_img=="": return # because this also gets triggered when removing outlier
+
         self.active_mask_index = self.viewer.dims.current_step[0]
         masks = deepcopy(self.layer.data)