File size: 3,193 Bytes
5c1bb37 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """
IO 工具模块
图像和 JSON 文件的读写。
"""
import json
import numpy as np
from pathlib import Path
from typing import Dict, Any, Union
import cv2
def load_image(path: Union[str, Path]) -> np.ndarray:
"""
加载图像
Args:
path: 图像路径
Returns:
img: (H, W, 3) RGB uint8 图像
"""
path = Path(path)
if not path.exists():
raise FileNotFoundError(f"Image not found: {path}")
# OpenCV 读取 BGR
img = cv2.imread(str(path), cv2.IMREAD_COLOR)
if img is None:
raise ValueError(f"Failed to load image: {path}")
# BGR -> RGB
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
def save_image(img: np.ndarray, path: Union[str, Path]) -> None:
"""
保存图像
Args:
img: (H, W, 3) RGB uint8 图像
path: 输出路径
"""
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
# RGB -> BGR
if len(img.shape) == 3 and img.shape[2] == 3:
img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
else:
img_bgr = img
cv2.imwrite(str(path), img_bgr)
def load_json(path: Union[str, Path]) -> Dict[str, Any]:
"""
加载 JSON 文件
Args:
path: JSON 文件路径
Returns:
data: 解析后的字典
"""
path = Path(path)
if not path.exists():
raise FileNotFoundError(f"JSON not found: {path}")
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def save_json(data: Dict[str, Any], path: Union[str, Path], indent: int = 2) -> None:
"""
保存 JSON 文件
Args:
data: 要保存的数据
path: 输出路径
indent: 缩进
"""
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
json.dump(data, f, indent=indent, ensure_ascii=False)
def load_depth(path: Union[str, Path]) -> np.ndarray:
"""
加载深度图（.npy 格式）
Args:
path: 深度图路径
Returns:
depth: (H, W) float32 深度图
"""
path = Path(path)
if not path.exists():
raise FileNotFoundError(f"Depth not found: {path}")
depth = np.load(str(path))
return depth.astype(np.float32)
def save_depth(depth: np.ndarray, path: Union[str, Path]) -> None:
"""
保存深度图（.npy 格式）
Args:
depth: (H, W) 深度图
path: 输出路径
"""
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
np.save(str(path), depth.astype(np.float32))
def list_files(
directory: Union[str, Path],
pattern: str = "*",
sort: bool = True,
) -> list:
"""
列出目录中的文件
Args:
directory: 目录路径
pattern: glob 模式
sort: 是否排序
Returns:
files: 文件路径列表
"""
directory = Path(directory)
if not directory.exists():
return []
files = list(directory.glob(pattern))
if sort:
files = sorted(files)
return files
|