diff --git a/model_server/extensions/chaeo/annotators.py b/model_server/base/annotators.py
similarity index 100%
rename from model_server/extensions/chaeo/annotators.py
rename to model_server/base/annotators.py
diff --git a/model_server/base/roiset.py b/model_server/base/roiset.py
index 4aecae346e71ef56dc805b0b2549df589038563a..0314c4f9ada9576925b6ca589b219e4364c49e1d 100644
--- a/model_server/base/roiset.py
+++ b/model_server/base/roiset.py
@@ -16,7 +16,7 @@ from sklearn.linear_model import LinearRegression
 from model_server.base.accessors import GenericImageDataAccessor, InMemoryDataAccessor, write_accessor_data_to_file
 from model_server.base.models import InstanceSegmentationModel
 from model_server.base.process import pad, rescale, resample_to_8bit, make_rgb
-from model_server.extensions.chaeo.annotators import draw_box_on_patch, draw_contours_on_patch, draw_boxes_on_3d_image
+from base.annotators import draw_box_on_patch, draw_contours_on_patch, draw_boxes_on_3d_image
 from model_server.extensions.chaeo.accessors import write_patch_to_file, MonoPatchStack, Multichannel3dPatchStack
 from base.process import mask_largest_object
 
diff --git a/tests/test_session.py b/tests/test_session.py
index dd143d5f48e6eec0057fca012a7e8e555139e016..9679aad61c699de8f07ff0a31e2cc4750cb6f9e3 100644
--- a/tests/test_session.py
+++ b/tests/test_session.py
@@ -1,6 +1,6 @@
 import pathlib
 import unittest
-from model_server.base.models import DummyImageToImageModel
+from model_server.base.models import DummySemanticSegmentationModel
 from model_server.base.session import Session
 
 class TestGetSessionObject(unittest.TestCase):
@@ -63,7 +63,7 @@ class TestGetSessionObject(unittest.TestCase):
 
     def test_session_loads_model(self):
         sesh = Session()
-        MC = DummyImageToImageModel
+        MC = DummySemanticSegmentationModel
         success = sesh.load_model(MC)
         self.assertTrue(success)
         loaded_models = sesh.describe_loaded_models()
@@ -77,7 +77,7 @@ class TestGetSessionObject(unittest.TestCase):
 
     def test_session_loads_second_instance_of_same_model(self):
         sesh = Session()
-        MC = DummyImageToImageModel
+        MC = DummySemanticSegmentationModel
         sesh.load_model(MC)
         sesh.load_model(MC)
         self.assertIn(MC.__name__ + '_00', sesh.models.keys())
@@ -86,7 +86,7 @@ class TestGetSessionObject(unittest.TestCase):
 
     def test_session_loads_model_with_params(self):
         sesh = Session()
-        MC = DummyImageToImageModel
+        MC = DummySemanticSegmentationModel
         p1 = {'p1': 'abc'}
         success = sesh.load_model(MC, params=p1)
         self.assertTrue(success)
@@ -103,7 +103,7 @@ class TestGetSessionObject(unittest.TestCase):
 
     def test_session_finds_existing_model_with_different_path_formats(self):
         sesh = Session()
-        MC = DummyImageToImageModel
+        MC = DummySemanticSegmentationModel
         p1 = {'path': 'c:\\windows\\dummy.pa'}
         p2 = {'path': 'c:/windows/dummy.pa'}
         mid = sesh.load_model(MC, params=p1)