Fixing grayscale conversion

src/model/yolo_wrapper.py

@@ -10,7 +10,7 @@ import torch
 from PIL import Image
 import tempfile
 import os
-import numpy as np
+from src.utils.image import convert_grayscale_to_rgb_preserve_range
 from src.utils.logger import get_logger
 
 
@@ -234,38 +234,20 @@ class YOLOWrapper:
                 try:
                     with Image.open(source_path) as img:
                         if len(img.getbands()) == 1:
-                            grayscale = np.array(img)
-                            if grayscale.ndim == 3:
-                                grayscale = grayscale[:, :, 0]
-                            original_dtype = grayscale.dtype
-                            grayscale = grayscale.astype(np.float32)
+                            rgb_img = convert_grayscale_to_rgb_preserve_range(img)
+                        else:
+                            rgb_img = img.convert("RGB")
 
-                            if np.issubdtype(original_dtype, np.integer):
-                                dtype_info = np.iinfo(original_dtype)
-                                denom = float(max(dtype_info.max, 1))
-                            else:
-                                max_val = (
-                                    float(grayscale.max()) if grayscale.size else 0.0
-                                )
-                                denom = max(max_val, 1.0)
-
-                            grayscale = np.clip(grayscale / denom, 0.0, 1.0)
-                            grayscale_u8 = (grayscale * 255.0).round().astype(np.uint8)
-                            rgb_arr = np.repeat(grayscale_u8[:, :, None], 3, axis=2)
-                            rgb_img = Image.fromarray(rgb_arr, mode="RGB")
-
-                            suffix = source_path.suffix or ".png"
-                            tmp = tempfile.NamedTemporaryFile(
-                                suffix=suffix, delete=False
-                            )
-                            tmp_path = tmp.name
-                            tmp.close()
-                            rgb_img.save(tmp_path)
-                            cleanup_path = tmp_path
-                            logger.info(
-                                f"Converted single-channel image {source_path} to RGB for inference at {tmp_path}"
-                            )
-                            return tmp_path, cleanup_path
+                        suffix = source_path.suffix or ".png"
+                        tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
+                        tmp_path = tmp.name
+                        tmp.close()
+                        rgb_img.save(tmp_path)
+                        cleanup_path = tmp_path
+                        logger.info(
+                            f"Converted image {source_path} to RGB for inference at {tmp_path}"
+                        )
+                        return tmp_path, cleanup_path
                 except Exception as convert_error:
                     logger.warning(
                         f"Failed to preprocess {source_path} as RGB, continuing with original file: {convert_error}"