diff --git a/tests/show_yolo_seg.py b/tests/show_yolo_seg.py
index 4e2b7ad..74676e5 100644
--- a/tests/show_yolo_seg.py
+++ b/tests/show_yolo_seg.py
@@ -18,6 +18,8 @@ import argparse
 from pathlib import Path
 import random
 
+from src.utils.image import Image
+
 
 def parse_label_line(line):
     parts = line.strip().split()
@@ -62,7 +64,11 @@ def poly_to_pts(coords, img_w, img_h):
 
 def random_color_for_class(cls):
     random.seed(cls)  # deterministic per class
-    return tuple(int(x) for x in np.array([random.randint(0, 255) for _ in range(3)]))
+    return (
+        0,
+        0,
+        255,
+    )  # tuple(int(x) for x in np.array([random.randint(0, 255) for _ in range(3)]))
 
 
 def draw_annotations(img, labels, alpha=0.4, draw_bbox_for_poly=True):
@@ -77,13 +83,13 @@ def draw_annotations(img, labels, alpha=0.4, draw_bbox_for_poly=True):
             color = random_color_for_class(cls)
 
             x1, y1, x2, y2 = yolo_bbox_to_xyxy(coords[:4], w, h)
-            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+            cv2.rectangle(img, (x1, y1), (x2, y2), color, 1)
 
             pts = poly_to_pts(coords[4:], w, h)
             # fill on overlay
             cv2.fillPoly(overlay, [pts], color)
             # outline on base image
-            cv2.polylines(img, [pts], isClosed=True, color=color, thickness=2)
+            cv2.polylines(img, [pts], isClosed=True, color=color, thickness=1)
             # put class text at first point
             x, y = int(pts[0, 0]), int(pts[0, 1]) - 6
             cv2.putText(
@@ -158,7 +164,9 @@ def main():
         print("Label file not found:", lbl_path)
         sys.exit(1)
 
-    img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
+    # img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
+    img = (Image(img_path).get_qt_rgb() * 255).astype(np.uint8)
+
     if img is None:
         print("Could not load image:", img_path)
         sys.exit(1)
@@ -173,6 +181,7 @@ def main():
 
     # Convert BGR -> RGB for matplotlib display
     out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
+    # out_rgb = Image()
     plt.figure(figsize=(10, 10 * out.shape[0] / out.shape[1]))
     plt.imshow(out_rgb)
     plt.axis("off")