273 lines
10 KiB
Python
273 lines
10 KiB
Python
import os
|
|
import argparse
|
|
import subprocess
|
|
import cv2
|
|
import numpy as np
|
|
import glob
|
|
import sys
|
|
|
|
# python '.\1_Batch_Gen_All_Depth(跑之前删除结果文件夹).py' # --img-path "./Data"
|
|
|
|
# === 配置区域 ===
|
|
# 两个项目的文件夹名称
|
|
img_path = "./Data"
|
|
V1_FOLDER = "Depth-Anything-V1-main"
|
|
V2_FOLDER = "Depth-Anything-V2-main"
|
|
|
|
# 定义各自包含的模型列表
|
|
V1_ENCODERS = ['vits', 'vitb', 'vitl']
|
|
V2_ENCODERS = ['vits', 'vitb', 'vitl', 'vitg']
|
|
# ================
|
|
|
|
def read_image_safe(path):
|
|
"""
|
|
[解决中文路径问题] 读取图片
|
|
使用 numpy 读取字节流再解码，绕过 cv2.imread 的路径编码 bug
|
|
"""
|
|
try:
|
|
# np.fromfile 读取文件内容到内存
|
|
img_array = np.fromfile(path, dtype=np.uint8)
|
|
# cv2.imdecode 解码内存数据
|
|
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
|
|
return img
|
|
except Exception as e:
|
|
print(f"Error reading image {path}: {e}")
|
|
return None
|
|
|
|
def write_image_safe(path, img):
|
|
"""
|
|
[解决中文路径问题] 保存图片
|
|
使用 cv2.imencode 编码再用 numpy 保存，绕过 cv2.imwrite 的路径编码 bug
|
|
"""
|
|
try:
|
|
# 获取文件扩展名 (例如 .png)
|
|
ext = os.path.splitext(path)[1]
|
|
if not ext:
|
|
ext = ".png"
|
|
|
|
# cv2.imencode 编码图片
|
|
success, img_array = cv2.imencode(ext, img)
|
|
if success:
|
|
# tofile 保存到文件
|
|
img_array.tofile(path)
|
|
return True
|
|
else:
|
|
print(f"Error encoding image for {path}")
|
|
return False
|
|
except Exception as e:
|
|
print(f"Error writing image {path}: {e}")
|
|
return False
|
|
|
|
def run_inference(repo_path, encoder, img_path_abs, out_dir_abs, grayscale, version_tag):
|
|
"""
|
|
调用各自文件夹下的 run.py 进行推理
|
|
注意：如果 run.py 内部也使用了 cv2.imread 且未做中文处理，这里即使传入绝对路径也可能报错。
|
|
建议尽量确保 input 路径不含特殊字符，或者修改 run.py 内部源码。
|
|
"""
|
|
if not os.path.exists(repo_path):
|
|
print(f"Error: Repository folder not found: {repo_path}")
|
|
return False
|
|
|
|
script_path = "run.py"
|
|
|
|
cmd = [
|
|
sys.executable, script_path,
|
|
'--encoder', encoder,
|
|
'--img-path', img_path_abs,
|
|
'--outdir', out_dir_abs,
|
|
'--pred-only'
|
|
]
|
|
|
|
if grayscale:
|
|
cmd.append('--grayscale')
|
|
|
|
print(f"[{version_tag}] Executing {encoder} (Gray={grayscale})...")
|
|
|
|
try:
|
|
# cwd=repo_path 是关键
|
|
subprocess.check_call(cmd, cwd=repo_path)
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Error executing {version_tag} {encoder}: {e}")
|
|
return False
|
|
return True
|
|
|
|
def add_text(img, text, is_header=False):
|
|
"""
|
|
在图片左上角添加文字标注
|
|
"""
|
|
if img is None: return None
|
|
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
font_scale = 1.0 if not is_header else 1.2
|
|
font_color = (255, 255, 255)
|
|
thickness = 2
|
|
position = (30, 50)
|
|
|
|
cv2.putText(img, text, position, font, font_scale, (0, 0, 0), thickness + 3, cv2.LINE_AA)
|
|
cv2.putText(img, text, position, font, font_scale, font_color, thickness, cv2.LINE_AA)
|
|
return img
|
|
|
|
def create_grand_comparison(image_files, root_out_dir, v1_out_base, v2_out_base, mode_suffix, out_subfolder_name):
|
|
"""
|
|
生成超大对比图： 原图 + V1结果 + V2结果
|
|
"""
|
|
compare_dir = os.path.join(root_out_dir, out_subfolder_name)
|
|
if not os.path.exists(compare_dir):
|
|
os.makedirs(compare_dir) # makedirs 支持中文路径
|
|
|
|
print(f"\n--- Generating Grand Comparisons ({out_subfolder_name}) ---")
|
|
|
|
for img_file in image_files:
|
|
file_basename = os.path.basename(img_file)
|
|
file_name_no_ext = os.path.splitext(file_basename)[0]
|
|
|
|
# 1. 读取原图 (使用 safe 读取)
|
|
raw_img = read_image_safe(img_file)
|
|
if raw_img is None:
|
|
print(f"Skipping {img_file} due to read error.")
|
|
continue
|
|
|
|
h, w = raw_img.shape[:2]
|
|
|
|
# 标注原图
|
|
raw_img_labeled = raw_img.copy()
|
|
raw_img_labeled = add_text(raw_img_labeled, "Original", is_header=True)
|
|
|
|
images_to_concat = [raw_img_labeled]
|
|
|
|
# 内部函数：读取并处理深度图
|
|
def process_result_img(base_dir, ver_tag, enc_name):
|
|
target_dir = os.path.join(base_dir, f'{ver_tag}_{enc_name}{mode_suffix}')
|
|
target_path = os.path.join(target_dir, file_name_no_ext + '.png')
|
|
|
|
# 尝试读取深度图 (使用 safe 读取)
|
|
d_img = None
|
|
if os.path.exists(target_path):
|
|
d_img = read_image_safe(target_path)
|
|
|
|
if d_img is None:
|
|
# 缺失或读取失败，生成黑色占位图
|
|
# print(f"Warning: Missing or failed to read {target_path}")
|
|
blank = np.zeros((h, w, 3), dtype=np.uint8)
|
|
return add_text(blank, f"{ver_tag}_{enc_name} (N/A)")
|
|
|
|
# 确保通道数
|
|
if len(d_img.shape) == 2:
|
|
d_img = cv2.cvtColor(d_img, cv2.COLOR_GRAY2BGR)
|
|
|
|
# 强制 Resize 到原图尺寸
|
|
if d_img.shape[:2] != (h, w):
|
|
d_img = cv2.resize(d_img, (w, h), interpolation=cv2.INTER_LINEAR)
|
|
|
|
return add_text(d_img, f"{ver_tag}_{enc_name}")
|
|
|
|
# 2. 收集 V1 结果
|
|
for enc in V1_ENCODERS:
|
|
images_to_concat.append(process_result_img(root_out_dir, "V1", enc))
|
|
|
|
# 3. 收集 V2 结果
|
|
for enc in V2_ENCODERS:
|
|
images_to_concat.append(process_result_img(root_out_dir, "V2", enc))
|
|
|
|
# 4. 拼接并保存
|
|
try:
|
|
combined_result = cv2.hconcat(images_to_concat)
|
|
save_path = os.path.join(compare_dir, f"{file_name_no_ext}_compare.png")
|
|
|
|
# 使用 safe 保存
|
|
if write_image_safe(save_path, combined_result):
|
|
print(f"Saved: {save_path}")
|
|
else:
|
|
print(f"Failed to save: {save_path}")
|
|
|
|
except Exception as e:
|
|
print(f"Error combining images for {file_name_no_ext}: {e}")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Universal Depth Anything Batch Process (V1 & V2) - Chinese Path Supported')
|
|
parser.add_argument('--img-path', type=str, default=img_path, help='Path to input images folder')
|
|
parser.add_argument('--skip-inference', action='store_true', help='Skip inference, only generate comparison')
|
|
args = parser.parse_args()
|
|
|
|
# 路径绝对化处理
|
|
input_path_abs = os.path.abspath(args.img_path)
|
|
|
|
# 定义总输出目录： 输入文件夹名 + "_Universal_Results"
|
|
root_folder_name = os.path.basename(os.path.normpath(input_path_abs)) + "_Universal_Results"
|
|
root_out_abs = os.path.abspath(root_folder_name)
|
|
|
|
if not os.path.exists(root_out_abs):
|
|
os.makedirs(root_out_abs)
|
|
|
|
print(f"Input: {input_path_abs}")
|
|
print(f"Output: {root_out_abs}")
|
|
|
|
# ==========================
|
|
# 1. 运行推理 (Inference)
|
|
# ==========================
|
|
# 注意：如果 Depth-Anything 原版 run.py 内部不支持中文路径，
|
|
# 这里即使调度成功，内部 subprocess 也可能报错。
|
|
# 解决方法：确保 run.py 内部也使用了类似 imdecode 的方法，或者仅处理中文路径的对比图生成。
|
|
if not args.skip_inference:
|
|
# --- Run V1 ---
|
|
for enc in V1_ENCODERS:
|
|
out_dir = os.path.join(root_out_abs, f'V1_{enc}_gray_result')
|
|
if not os.path.exists(out_dir):
|
|
run_inference(V1_FOLDER, enc, input_path_abs, out_dir, True, "V1")
|
|
|
|
out_dir = os.path.join(root_out_abs, f'V1_{enc}_colorful_result')
|
|
if not os.path.exists(out_dir):
|
|
run_inference(V1_FOLDER, enc, input_path_abs, out_dir, False, "V1")
|
|
|
|
# --- Run V2 ---
|
|
for enc in V2_ENCODERS:
|
|
out_dir = os.path.join(root_out_abs, f'V2_{enc}_gray_result')
|
|
if not os.path.exists(out_dir):
|
|
run_inference(V2_FOLDER, enc, input_path_abs, out_dir, True, "V2")
|
|
|
|
out_dir = os.path.join(root_out_abs, f'V2_{enc}_colorful_result')
|
|
if not os.path.exists(out_dir):
|
|
run_inference(V2_FOLDER, enc, input_path_abs, out_dir, False, "V2")
|
|
else:
|
|
print("Skipping inference step as requested...")
|
|
|
|
# ==========================
|
|
# 2. 生成对比图 (Comparison)
|
|
# ==========================
|
|
|
|
# 收集图片列表 (支持 glob 中文路径查找，但在 Windows 下有时需要注意编码)
|
|
exts = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tif', '*.tiff']
|
|
image_files = []
|
|
|
|
# glob.glob 在较新版本的 Python (3.10+) 中通常能较好处理中文路径
|
|
# 如果 glob 找不到文件，请尝试将路径转为 glob.escape(input_path_abs)
|
|
for ext in exts:
|
|
image_files.extend(glob.glob(os.path.join(input_path_abs, ext)))
|
|
# image_files.extend(glob.glob(os.path.join(input_path_abs, '**', ext), recursive=True))
|
|
|
|
image_files = sorted(list(set(image_files)))
|
|
|
|
if not image_files:
|
|
print("No images found in input path. (Check if path encoding is correct)")
|
|
# 调试信息
|
|
print(f"Glob pattern used: {os.path.join(input_path_abs, '*.jpg')} etc.")
|
|
return
|
|
|
|
# 生成灰度大对比
|
|
create_grand_comparison(
|
|
image_files, root_out_abs, "V1", "V2",
|
|
mode_suffix='_gray_result',
|
|
out_subfolder_name='Comparison_All_Gray'
|
|
)
|
|
|
|
# 生成彩色大对比
|
|
create_grand_comparison(
|
|
image_files, root_out_abs, "V1", "V2",
|
|
mode_suffix='_colorful_result',
|
|
out_subfolder_name='Comparison_All_Colorful'
|
|
)
|
|
|
|
print(f"\nAll tasks completed. Results stored in: {root_out_abs}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |