|
|
# -*- coding: utf-8 -*-
|
|
|
"""
|
|
|
图片加载模块
|
|
|
提供图片加载功能，支持单张图片、多张图片和目录批量加载
|
|
|
"""
|
|
|
|
|
|
import cv2
|
|
|
import numpy as np
|
|
|
from pathlib import Path
|
|
|
from typing import List, Optional, Generator, Union, Tuple
|
|
|
from dataclasses import dataclass
|
|
|
|
|
|
|
|
|
# 支持的图片格式
|
|
|
SUPPORTED_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp'}
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
class ImageInfo:
|
|
|
"""
|
|
|
图片信息数据结构
|
|
|
|
|
|
Attributes:
|
|
|
path: 图片文件路径
|
|
|
image: 图片数据 (numpy array, BGR 格式)
|
|
|
width: 图片宽度
|
|
|
height: 图片高度
|
|
|
filename: 文件名（不含路径）
|
|
|
"""
|
|
|
path: str
|
|
|
image: np.ndarray
|
|
|
width: int
|
|
|
height: int
|
|
|
filename: str
|
|
|
|
|
|
@classmethod
|
|
|
def from_image(cls, path: str, image: np.ndarray) -> "ImageInfo":
|
|
|
"""
|
|
|
从图片数据创建 ImageInfo
|
|
|
|
|
|
Args:
|
|
|
path: 图片路径
|
|
|
image: 图片数据
|
|
|
|
|
|
Returns:
|
|
|
ImageInfo 实例
|
|
|
"""
|
|
|
h, w = image.shape[:2]
|
|
|
return cls(
|
|
|
path=path,
|
|
|
image=image,
|
|
|
width=w,
|
|
|
height=h,
|
|
|
filename=Path(path).name
|
|
|
)
|
|
|
|
|
|
|
|
|
class ImageLoader:
|
|
|
"""
|
|
|
图片加载器
|
|
|
支持单张图片、多张图片列表和目录批量加载
|
|
|
"""
|
|
|
|
|
|
def __init__(self, supported_extensions: Optional[set] = None):
|
|
|
"""
|
|
|
初始化图片加载器
|
|
|
|
|
|
Args:
|
|
|
supported_extensions: 支持的图片扩展名集合，默认使用 SUPPORTED_EXTENSIONS
|
|
|
"""
|
|
|
self._extensions = supported_extensions or SUPPORTED_EXTENSIONS
|
|
|
|
|
|
def load(self, path: Union[str, Path]) -> Optional[ImageInfo]:
|
|
|
"""
|
|
|
加载单张图片
|
|
|
|
|
|
Args:
|
|
|
path: 图片文件路径
|
|
|
|
|
|
Returns:
|
|
|
ImageInfo 对象，加载失败返回 None
|
|
|
"""
|
|
|
path = Path(path)
|
|
|
|
|
|
if not path.exists():
|
|
|
print(f"[ERROR] 文件不存在: {path}")
|
|
|
return None
|
|
|
|
|
|
if not path.is_file():
|
|
|
print(f"[ERROR] 路径不是文件: {path}")
|
|
|
return None
|
|
|
|
|
|
if path.suffix.lower() not in self._extensions:
|
|
|
print(f"[ERROR] 不支持的图片格式: {path.suffix}")
|
|
|
return None
|
|
|
|
|
|
# 使用 cv2.imdecode 处理中文路径
|
|
|
image = self._read_image(str(path))
|
|
|
|
|
|
if image is None:
|
|
|
print(f"[ERROR] 无法读取图片: {path}")
|
|
|
return None
|
|
|
|
|
|
return ImageInfo.from_image(str(path), image)
|
|
|
|
|
|
def load_batch(self, paths: List[Union[str, Path]]) -> Generator[ImageInfo, None, None]:
|
|
|
"""
|
|
|
批量加载多张图片
|
|
|
|
|
|
Args:
|
|
|
paths: 图片路径列表
|
|
|
|
|
|
Yields:
|
|
|
ImageInfo 对象
|
|
|
"""
|
|
|
for path in paths:
|
|
|
info = self.load(path)
|
|
|
if info is not None:
|
|
|
yield info
|
|
|
|
|
|
def load_directory(
|
|
|
self,
|
|
|
dir_path: Union[str, Path],
|
|
|
pattern: Optional[str] = None,
|
|
|
recursive: bool = False
|
|
|
) -> Generator[ImageInfo, None, None]:
|
|
|
"""
|
|
|
加载目录中的所有图片
|
|
|
|
|
|
Args:
|
|
|
dir_path: 目录路径
|
|
|
pattern: 文件匹配模式（如 "*.jpg"），None 表示加载所有支持的格式
|
|
|
recursive: 是否递归搜索子目录
|
|
|
|
|
|
Yields:
|
|
|
ImageInfo 对象
|
|
|
"""
|
|
|
dir_path = Path(dir_path)
|
|
|
|
|
|
if not dir_path.exists():
|
|
|
print(f"[ERROR] 目录不存在: {dir_path}")
|
|
|
return
|
|
|
|
|
|
if not dir_path.is_dir():
|
|
|
print(f"[ERROR] 路径不是目录: {dir_path}")
|
|
|
return
|
|
|
|
|
|
# 获取文件列表
|
|
|
if pattern:
|
|
|
if recursive:
|
|
|
files = list(dir_path.rglob(pattern))
|
|
|
else:
|
|
|
files = list(dir_path.glob(pattern))
|
|
|
else:
|
|
|
# 加载所有支持的格式
|
|
|
files = []
|
|
|
for ext in self._extensions:
|
|
|
if recursive:
|
|
|
files.extend(dir_path.rglob(f"*{ext}"))
|
|
|
files.extend(dir_path.rglob(f"*{ext.upper()}"))
|
|
|
else:
|
|
|
files.extend(dir_path.glob(f"*{ext}"))
|
|
|
files.extend(dir_path.glob(f"*{ext.upper()}"))
|
|
|
|
|
|
# 按文件名排序
|
|
|
files = sorted(set(files), key=lambda p: p.name)
|
|
|
|
|
|
print(f"[INFO] 在目录 {dir_path} 中找到 {len(files)} 张图片")
|
|
|
|
|
|
for file_path in files:
|
|
|
info = self.load(file_path)
|
|
|
if info is not None:
|
|
|
yield info
|
|
|
|
|
|
def get_image_paths(
|
|
|
self,
|
|
|
dir_path: Union[str, Path],
|
|
|
pattern: Optional[str] = None,
|
|
|
recursive: bool = False
|
|
|
) -> List[str]:
|
|
|
"""
|
|
|
获取目录中所有图片的路径列表（不加载图片）
|
|
|
|
|
|
Args:
|
|
|
dir_path: 目录路径
|
|
|
pattern: 文件匹配模式
|
|
|
recursive: 是否递归搜索
|
|
|
|
|
|
Returns:
|
|
|
图片路径列表
|
|
|
"""
|
|
|
dir_path = Path(dir_path)
|
|
|
|
|
|
if not dir_path.exists() or not dir_path.is_dir():
|
|
|
return []
|
|
|
|
|
|
if pattern:
|
|
|
if recursive:
|
|
|
files = list(dir_path.rglob(pattern))
|
|
|
else:
|
|
|
files = list(dir_path.glob(pattern))
|
|
|
else:
|
|
|
files = []
|
|
|
for ext in self._extensions:
|
|
|
if recursive:
|
|
|
files.extend(dir_path.rglob(f"*{ext}"))
|
|
|
files.extend(dir_path.rglob(f"*{ext.upper()}"))
|
|
|
else:
|
|
|
files.extend(dir_path.glob(f"*{ext}"))
|
|
|
files.extend(dir_path.glob(f"*{ext.upper()}"))
|
|
|
|
|
|
return sorted([str(f) for f in set(files)], key=lambda p: Path(p).name)
|
|
|
|
|
|
def _read_image(self, path: str) -> Optional[np.ndarray]:
|
|
|
"""
|
|
|
读取图片，支持中文路径
|
|
|
|
|
|
Args:
|
|
|
path: 图片路径
|
|
|
|
|
|
Returns:
|
|
|
图片数据，读取失败返回 None
|
|
|
"""
|
|
|
# 使用 numpy 和 imdecode 处理中文路径
|
|
|
try:
|
|
|
with open(path, 'rb') as f:
|
|
|
data = np.frombuffer(f.read(), dtype=np.uint8)
|
|
|
image = cv2.imdecode(data, cv2.IMREAD_COLOR)
|
|
|
return image
|
|
|
except Exception as e:
|
|
|
print(f"[ERROR] 读取图片失败: {path}, 错误: {e}")
|
|
|
return None
|
|
|
|
|
|
@property
|
|
|
def supported_extensions(self) -> set:
|
|
|
"""获取支持的图片扩展名"""
|
|
|
return self._extensions.copy()
|
|
|
|
|
|
|
|
|
def load_image(path: Union[str, Path]) -> Optional[np.ndarray]:
|
|
|
"""
|
|
|
便捷函数：加载单张图片
|
|
|
|
|
|
Args:
|
|
|
path: 图片路径
|
|
|
|
|
|
Returns:
|
|
|
图片数据 (numpy array, BGR 格式)，加载失败返回 None
|
|
|
"""
|
|
|
loader = ImageLoader()
|
|
|
info = loader.load(path)
|
|
|
return info.image if info else None
|
|
|
|
|
|
|
|
|
def load_images(
|
|
|
paths: Optional[List[Union[str, Path]]] = None,
|
|
|
directory: Optional[Union[str, Path]] = None,
|
|
|
pattern: Optional[str] = None,
|
|
|
recursive: bool = False
|
|
|
) -> Generator[ImageInfo, None, None]:
|
|
|
"""
|
|
|
便捷函数：批量加载图片
|
|
|
|
|
|
Args:
|
|
|
paths: 图片路径列表
|
|
|
directory: 图片目录
|
|
|
pattern: 文件匹配模式
|
|
|
recursive: 是否递归搜索
|
|
|
|
|
|
Yields:
|
|
|
ImageInfo 对象
|
|
|
"""
|
|
|
loader = ImageLoader()
|
|
|
|
|
|
if paths:
|
|
|
yield from loader.load_batch(paths)
|
|
|
elif directory:
|
|
|
yield from loader.load_directory(directory, pattern, recursive)
|