Updating

src/utils/image.py

@@ -37,7 +37,7 @@ def get_pseudo_rgb(arr: np.ndarray, gamma: float = 0.5) -> np.ndarray:
     a1[a1 > p999] = p999
     a1 /= a1.max()
 
-    if 0:
+    if 1:
         a2 = a1.copy()
         a2 = a2**gamma
         a2 /= a2.max()
@@ -47,9 +47,12 @@ def get_pseudo_rgb(arr: np.ndarray, gamma: float = 0.5) -> np.ndarray:
         a3[a3 > p9999] = p9999
         a3 /= a3.max()
 
-    return np.stack([a1, np.zeros(a1.shape), np.zeros(a1.shape)], axis=0)
+    # return np.stack([a1, np.zeros(a1.shape), np.zeros(a1.shape)], axis=0)
     # return np.stack([a2, np.zeros(a1.shape), np.zeros(a1.shape)], axis=0)
-    # return np.stack([a1, a2, a3], axis=0)
+    out = np.stack([a1, a2, a3], axis=0)
+    # print(any(np.isnan(out).flatten()))
+
+    return out
 
 
 class ImageLoadError(Exception):
@@ -122,7 +125,7 @@ class Image:
             if self.path.suffix.lower() in [".tif", ".tiff"]:
                 self._data = imread(str(self.path))
             else:
-                raise NotImplementedError("RGB is not implemented")
+                # raise NotImplementedError("RGB is not implemented")
                 # Load with OpenCV (returns BGR format)
                 self._data = cv2.imread(str(self.path), cv2.IMREAD_UNCHANGED)
 
@@ -246,20 +249,24 @@ class Image:
         if self.channels == 1:
             img = get_pseudo_rgb(self.data)
             self._dtype = img.dtype
-            return img
+            return img, True
+
+        elif self._channels == 3:
+            return cv2.cvtColor(self._data, cv2.COLOR_BGR2RGB), False
+        elif self._channels == 4:
+            return cv2.cvtColor(self._data, cv2.COLOR_BGRA2RGBA), False
+
+        else:
             raise NotImplementedError
 
-        if self._channels == 3:
-            return cv2.cvtColor(self._data, cv2.COLOR_BGR2RGB)
-        elif self._channels == 4:
-            return cv2.cvtColor(self._data, cv2.COLOR_BGRA2RGBA)
-        else:
-            return self._data
+        # else:
+        #     return self._data
 
     def get_qt_rgb(self) -> np.ascontiguousarray:
         # we keep data as (C, H, W)
-        _img = self.get_rgb()
+        _img, pseudo = self.get_rgb()
 
+        if pseudo:
             img = np.zeros((self.height, self.width, 4), dtype=np.float32)
             img[..., 0] = _img[0]  # R gradient
             img[..., 1] = _img[1]  # G gradient
@@ -267,6 +274,8 @@ class Image:
             img[..., 3] = 1.0  # A = 1.0 (opaque)
 
             return np.ascontiguousarray(img)
+        else:
+            return np.ascontiguousarray(_img)
 
     def get_grayscale(self) -> np.ndarray:
         """