diff --git a/MultiverSeg/SegmentEditorMultiverSegLib/SegmentEditorEffect.py b/MultiverSeg/SegmentEditorMultiverSegLib/SegmentEditorEffect.py
index 1971b89..c38fda3 100644
--- a/MultiverSeg/SegmentEditorMultiverSegLib/SegmentEditorEffect.py
+++ b/MultiverSeg/SegmentEditorMultiverSegLib/SegmentEditorEffect.py
@@ -302,6 +302,7 @@ def addImageToContext(self):
         self.changeViewLayout()
         self.addImageDialog()
         self.changeViewLayout(initialLayout)  # Restore the layout
+        self.handleTaskChange(self.contextComboBox.currentIndex)
 
     def removeImageFromContext(self):
         imageToRemove = self.removeImageDialog()
@@ -309,6 +310,7 @@ def removeImageFromContext(self):
             return
 
         self.contextLogic.removeExample(imageToRemove)
+        self.handleTaskChange(self.contextComboBox.currentIndex)
 
     def addImageDialog(self):
         # Create and handle the dialog to add an image to a context