martin/object-segmentation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Making image manipulations thru one class

Browse Source
This commit is contained in:
Martin Laasmaa
2025-12-11 16:59:56 +02:00
parent 32a6a122bd
commit 9a221acb63
3 changed files with 30 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
src/model/yolo_wrapper.py
View File

@@ -7,10 +7,9 @@ from ultralytics import YOLO
from pathlib import Path
from typing import Optional, List, Dict, Callable, Any
import torch
from PIL import Image
import tempfile
import os
from src.utils.image import convert_grayscale_to_rgb_preserve_range
from src.utils.image import Image, convert_grayscale_to_rgb_preserve_range
from src.utils.logger import get_logger
@@ -232,22 +231,23 @@ class YOLOWrapper:
source_path = Path(source)
if source_path.is_file():
try:
with Image.open(source_path) as img:
if len(img.getbands()) == 1:
rgb_img = convert_grayscale_to_rgb_preserve_range(img)
else:
rgb_img = img.convert("RGB")
img_obj = Image(source_path)
pil_img = img_obj.pil_image
if len(pil_img.getbands()) == 1:
rgb_img = convert_grayscale_to_rgb_preserve_range(pil_img)
else:
rgb_img = pil_img.convert("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 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}"