diff --git a/NN_server/Models/ensemble_1200_v44.py b/NN_server/Models/ensemble_1200_v44.py
index db22dfa..bc285e2 100644
--- a/NN_server/Models/ensemble_1200_v44.py
+++ b/NN_server/Models/ensemble_1200_v44.py
@@ -10,18 +10,16 @@ import os
 import re
 
 
-def _as_display_image(image):
+def _as_training_colormap_image(image):
     arr = np.asarray(image)
     if arr.ndim == 3 and arr.shape[0] in {1, 3, 4}:
         arr = np.moveaxis(arr, 0, -1)
 
-    if np.issubdtype(arr.dtype, np.floating):
-        arr = np.nan_to_num(arr, nan=0.0, posinf=255.0, neginf=0.0)
-        if arr.size and (arr.max() > 1.0 or arr.min() < 0.0):
-            return np.clip(arr, 0, 255).astype(np.uint8)
-        return np.clip(arr, 0.0, 1.0)
+    if arr.ndim == 3:
+        arr = arr[..., :3].astype(np.float32)
+        arr = 0.299 * arr[..., 0] + 0.587 * arr[..., 1] + 0.114 * arr[..., 2]
 
-    return arr
+    return np.nan_to_num(arr.astype(np.float32), nan=0.0, posinf=255.0, neginf=0.0)
 
 
 def _render_signal_channel(values, figsize=(16, 16), dpi=16, resize=(256, 256)):
@@ -246,7 +244,7 @@ def post_func_ensemble(src="", model_type="", prediction="", model_id=0, ind_inf
 
         if isinstance(data, (list, tuple)) and len(data) >= 2:
             fig, ax = plt.subplots()
-            ax.imshow(_as_display_image(data[0]))
+            ax.imshow(_as_training_colormap_image(data[0]), cmap="viridis")
             plt.savefig(src + "_inference_" + str(ind_inference) + "_" + prediction + "_real_" + str(model_id) + "_" + model_type + ".png")
             plt.clf()
             plt.cla()
@@ -255,7 +253,7 @@ def post_func_ensemble(src="", model_type="", prediction="", model_id=0, ind_inf
             gc.collect()
 
             fig, ax = plt.subplots()
-            ax.imshow(_as_display_image(data[1]))
+            ax.imshow(_as_training_colormap_image(data[1]), cmap="viridis")
             plt.savefig(src + "_inference_" + str(ind_inference) + "_" + prediction + "_mod_" + str(model_id) + "_" + model_type + ".png")
             plt.clf()
             plt.cla()
diff --git a/NN_server/Models/ensemble_2400_v44.py b/NN_server/Models/ensemble_2400_v44.py
index db22dfa..bc285e2 100644
--- a/NN_server/Models/ensemble_2400_v44.py
+++ b/NN_server/Models/ensemble_2400_v44.py
@@ -10,18 +10,16 @@ import os
 import re
 
 
-def _as_display_image(image):
+def _as_training_colormap_image(image):
     arr = np.asarray(image)
     if arr.ndim == 3 and arr.shape[0] in {1, 3, 4}:
         arr = np.moveaxis(arr, 0, -1)
 
-    if np.issubdtype(arr.dtype, np.floating):
-        arr = np.nan_to_num(arr, nan=0.0, posinf=255.0, neginf=0.0)
-        if arr.size and (arr.max() > 1.0 or arr.min() < 0.0):
-            return np.clip(arr, 0, 255).astype(np.uint8)
-        return np.clip(arr, 0.0, 1.0)
+    if arr.ndim == 3:
+        arr = arr[..., :3].astype(np.float32)
+        arr = 0.299 * arr[..., 0] + 0.587 * arr[..., 1] + 0.114 * arr[..., 2]
 
-    return arr
+    return np.nan_to_num(arr.astype(np.float32), nan=0.0, posinf=255.0, neginf=0.0)
 
 
 def _render_signal_channel(values, figsize=(16, 16), dpi=16, resize=(256, 256)):
@@ -246,7 +244,7 @@ def post_func_ensemble(src="", model_type="", prediction="", model_id=0, ind_inf
 
         if isinstance(data, (list, tuple)) and len(data) >= 2:
             fig, ax = plt.subplots()
-            ax.imshow(_as_display_image(data[0]))
+            ax.imshow(_as_training_colormap_image(data[0]), cmap="viridis")
             plt.savefig(src + "_inference_" + str(ind_inference) + "_" + prediction + "_real_" + str(model_id) + "_" + model_type + ".png")
             plt.clf()
             plt.cla()
@@ -255,7 +253,7 @@ def post_func_ensemble(src="", model_type="", prediction="", model_id=0, ind_inf
             gc.collect()
 
             fig, ax = plt.subplots()
-            ax.imshow(_as_display_image(data[1]))
+            ax.imshow(_as_training_colormap_image(data[1]), cmap="viridis")
             plt.savefig(src + "_inference_" + str(ind_inference) + "_" + prediction + "_mod_" + str(model_id) + "_" + model_type + ".png")
             plt.clf()
             plt.cla()
diff --git a/NN_server/Models/ensemble_915_v44.py b/NN_server/Models/ensemble_915_v44.py
index 88bc250..2f3b53b 100644
--- a/NN_server/Models/ensemble_915_v44.py
+++ b/NN_server/Models/ensemble_915_v44.py
@@ -10,18 +10,16 @@ import os
 import re
 
 
-def _as_display_image(image):
+def _as_training_colormap_image(image):
     arr = np.asarray(image)
     if arr.ndim == 3 and arr.shape[0] in {1, 3, 4}:
         arr = np.moveaxis(arr, 0, -1)
 
-    if np.issubdtype(arr.dtype, np.floating):
-        arr = np.nan_to_num(arr, nan=0.0, posinf=255.0, neginf=0.0)
-        if arr.size and (arr.max() > 1.0 or arr.min() < 0.0):
-            return np.clip(arr, 0, 255).astype(np.uint8)
-        return np.clip(arr, 0.0, 1.0)
+    if arr.ndim == 3:
+        arr = arr[..., :3].astype(np.float32)
+        arr = 0.299 * arr[..., 0] + 0.587 * arr[..., 1] + 0.114 * arr[..., 2]
 
-    return arr
+    return np.nan_to_num(arr.astype(np.float32), nan=0.0, posinf=255.0, neginf=0.0)
 
 
 def _prune_old_inference_images(src, model_type, model_id, keep_last=200):
@@ -218,7 +216,7 @@ def post_func_ensemble(src="", model_type="", prediction="", model_id=0, ind_inf
 
         if isinstance(data, (list, tuple)) and len(data) >= 2:
             fig, ax = plt.subplots()
-            ax.imshow(_as_display_image(data[0]))
+            ax.imshow(_as_training_colormap_image(data[0]), cmap="viridis")
             plt.savefig(src + "_inference_" + str(ind_inference) + "_" + prediction + "_real_" + str(model_id) + "_" + model_type + ".png")
             plt.clf()
             plt.cla()
@@ -227,7 +225,7 @@ def post_func_ensemble(src="", model_type="", prediction="", model_id=0, ind_inf
             gc.collect()
 
             fig, ax = plt.subplots()
-            ax.imshow(_as_display_image(data[1]))
+            ax.imshow(_as_training_colormap_image(data[1]), cmap="viridis")
             plt.savefig(src + "_inference_" + str(ind_inference) + "_" + prediction + "_mod_" + str(model_id) + "_" + model_type + ".png")
             plt.clf()
             plt.cla()