diff --git a/qim3d/tests/gui/test_annotation_tool.py b/qim3d/tests/gui/test_annotation_tool.py
index 65c02931efbcc6503a1a1bc4c1899ff0186c2d0d..2b80b9fb178df7d57f6e55633152b78e7876da20 100644
--- a/qim3d/tests/gui/test_annotation_tool.py
+++ b/qim3d/tests/gui/test_annotation_tool.py
@@ -9,14 +9,15 @@ def test_starting_class():
     assert app.title == "Annotation tool"
 
 
+def start_server(ip, port):
+    app = qim3d.gui.annotation_tool.Interface()
+    app.launch(server_name=ip, server_port=port)
+
+
 def test_app_launch():
-    ip = "0.0.0.0"
+    ip = "localhost"
     port = 65432
 
-    def start_server(ip, port):
-        app = qim3d.gui.annotation_tool.Interface()
-        app.launch(server_name=ip, server_port=port)
-
     proc = multiprocessing.Process(target=start_server, args=(ip, port))
     proc.start()
 
diff --git a/qim3d/tests/gui/test_iso3d.py b/qim3d/tests/gui/test_iso3d.py
index 4e10123db7506ea97cb7da4668d33117ad6906fe..864ede391c6fb4300ac0b3d4bb54afef44dfbd9c 100644
--- a/qim3d/tests/gui/test_iso3d.py
+++ b/qim3d/tests/gui/test_iso3d.py
@@ -9,14 +9,15 @@ def test_starting_class():
     assert app.title == "Isosurfaces for 3D visualization"
 
 
+def start_server(ip, port):
+    app = qim3d.gui.iso3d.Interface()
+    app.launch(server_name=ip, server_port=port)
+
+
 def test_app_launch():
-    ip = "0.0.0.0"
+    ip = "localhost"
     port = 65432
 
-    def start_server(ip, port):
-        app = qim3d.gui.iso3d.Interface()
-        app.launch(server_name=ip, server_port=port)
-
     proc = multiprocessing.Process(target=start_server, args=(ip, port))
     proc.start()
 
diff --git a/qim3d/utils/internal_tools.py b/qim3d/utils/internal_tools.py
index c8340d45ab9362814893d9fc707118676966cf83..a4f6ede02c206a68e8039e7aed19ad20e478d9f5 100644
--- a/qim3d/utils/internal_tools.py
+++ b/qim3d/utils/internal_tools.py
@@ -6,7 +6,6 @@ import outputformat as ouf
 import matplotlib.pyplot as plt
 import matplotlib
 import numpy as np
-import socket
 import os
 import shutil
 import requests
diff --git a/qim3d/viz/img.py b/qim3d/viz/img.py
index 83edde50b5c98af6a5b21cbd2b46e32dda8515a1..e3648f5445c7a4449c2b42057c78b2db9e48e381 100644
--- a/qim3d/viz/img.py
+++ b/qim3d/viz/img.py
@@ -1,7 +1,7 @@
 """ Provides a collection of visualization functions."""
 import matplotlib.pyplot as plt
 from matplotlib.colors import LinearSegmentedColormap
-from matplotlib import cm
+from matplotlib import colormaps
 import torch
 import numpy as np
 from qim3d.io.logger import log
@@ -56,7 +56,7 @@ def grid_overview(data, num_images=7, cmap_im="gray", cmap_segm="viridis", alpha
         num_images = len(data)
 
     # Adapt segmentation cmap so that background is transparent
-    colors_segm = cm.get_cmap(cmap_segm)(np.linspace(0, 1, 256))
+    colors_segm = colormaps.get_cmap(cmap_segm)(np.linspace(0, 1, 256))
     colors_segm[:128, 3] = 0
     custom_cmap = LinearSegmentedColormap.from_list("CustomCmap", colors_segm)
 
@@ -142,7 +142,7 @@ def grid_pred(in_targ_preds, num_images=7, cmap_im="gray", cmap_segm="viridis",
     inputs,targets,preds = [items[:num_images] for items in in_targ_preds]
     
     # Adapt segmentation cmap so that background is transparent
-    colors_segm = cm.get_cmap(cmap_segm)(np.linspace(0, 1, 256))
+    colors_segm = colormaps.get_cmap(cmap_segm)(np.linspace(0, 1, 256))
     colors_segm[:128, 3] = 0
     custom_cmap = LinearSegmentedColormap.from_list("CustomCmap", colors_segm)