Grayscale RGB conversion modified

src/model/yolo_wrapper.py

@@ -10,6 +10,7 @@ import torch
 from PIL import Image
 import tempfile
 import os
+import numpy as np
 from src.utils.logger import get_logger
 
 
@@ -188,7 +189,7 @@ class YOLOWrapper:
             logger.error(f"Error during inference: {e}")
             raise
         finally:
-            if cleanup_path:
+            if 0:  # cleanup_path:
                 try:
                     os.remove(cleanup_path)
                 except OSError as cleanup_error:
@@ -233,7 +234,26 @@ class YOLOWrapper:
                 try:
                     with Image.open(source_path) as img:
                         if len(img.getbands()) == 1:
-                            rgb_img = img.convert("RGB")
+                            grayscale = np.array(img)
+                            if grayscale.ndim == 3:
+                                grayscale = grayscale[:, :, 0]
+                            original_dtype = grayscale.dtype
+                            grayscale = grayscale.astype(np.float32)
+
+                            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
@@ -243,7 +263,7 @@ class YOLOWrapper:
                             rgb_img.save(tmp_path)
                             cleanup_path = tmp_path
                             logger.info(
-                                f"Converted single-channel image {source_path} to RGB for inference"
+                                f"Converted single-channel image {source_path} to RGB for inference at {tmp_path}"
                             )
                             return tmp_path, cleanup_path
                 except Exception as convert_error: