Initial media depth project backup
This commit is contained in:
213
Depth-Anything-V2-main/batch_process_and_compare_V2.py
Normal file
213
Depth-Anything-V2-main/batch_process_and_compare_V2.py
Normal file
@@ -0,0 +1,213 @@
|
||||
import os
|
||||
import argparse
|
||||
import subprocess
|
||||
import cv2
|
||||
import numpy as np
|
||||
import glob
|
||||
import sys
|
||||
|
||||
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_encoder(encoder, img_path, out_dir, grayscale):
|
||||
"""
|
||||
调用 run.py 生成指定 encoder 的深度图
|
||||
"""
|
||||
cmd = [
|
||||
sys.executable, 'run.py',
|
||||
'--encoder', encoder,
|
||||
'--img-path', img_path,
|
||||
'--outdir', out_dir,
|
||||
'--pred-only'
|
||||
]
|
||||
|
||||
if grayscale:
|
||||
cmd.append('--grayscale')
|
||||
|
||||
print(f"Executing: {' '.join(cmd)}")
|
||||
try:
|
||||
subprocess.check_call(cmd)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Error occurred while running {encoder} (Gray={grayscale}): {e}")
|
||||
return False
|
||||
return True
|
||||
|
||||
def add_text(img, text):
|
||||
"""
|
||||
在图片左上角添加文字标注
|
||||
"""
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
font_scale = 1.0
|
||||
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_comparison(image_files, folder_name, encoders, mode_suffix, out_folder_name):
|
||||
"""
|
||||
生成对比图的核心逻辑
|
||||
"""
|
||||
compare_dir = os.path.join(folder_name, out_folder_name)
|
||||
os.makedirs(compare_dir, exist_ok=True)
|
||||
|
||||
print(f"\n--- Generating Comparisons ({out_folder_name}) -> {compare_dir} ---")
|
||||
|
||||
for img_file in image_files:
|
||||
file_basename = os.path.basename(img_file)
|
||||
file_name_no_ext = os.path.splitext(file_basename)[0]
|
||||
|
||||
# 1. 读取原图
|
||||
raw_img = read_image_safe(img_file)
|
||||
if raw_img is None:
|
||||
continue
|
||||
|
||||
# 获取原图尺寸 (高度, 宽度)
|
||||
h, w = raw_img.shape[:2]
|
||||
|
||||
# 标注原图
|
||||
raw_img_labeled = raw_img.copy()
|
||||
raw_img_labeled = add_text(raw_img_labeled, "Original")
|
||||
|
||||
images_to_concat = [raw_img_labeled]
|
||||
valid_group = True
|
||||
|
||||
# 2. 依次读取 4 个 encoder 的结果
|
||||
for encoder in encoders:
|
||||
result_dir = os.path.join(folder_name, f'V2_{encoder}{mode_suffix}')
|
||||
result_path = os.path.join(result_dir, file_name_no_ext + '.png')
|
||||
|
||||
if not os.path.exists(result_path):
|
||||
print(f"Missing result for {file_basename} in {encoder}, skipping.")
|
||||
valid_group = False
|
||||
break
|
||||
|
||||
depth_img = read_image_safe(result_path)
|
||||
|
||||
if depth_img is None:
|
||||
print(f"Failed to read {result_path}, skipping.")
|
||||
valid_group = False
|
||||
break
|
||||
|
||||
# 确保通道数正确 (转为3通道以便拼接)
|
||||
if len(depth_img.shape) == 2:
|
||||
depth_img = cv2.cvtColor(depth_img, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
# --- 关键修改:强制 Resize 到原图尺寸 ---
|
||||
# 无论原始输出尺寸如何,都强制缩放为 (w, h)
|
||||
if depth_img.shape[:2] != (h, w):
|
||||
depth_img = cv2.resize(depth_img, (w, h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# 标注文字
|
||||
label = f"Depth_Anything_V2_{encoder}"
|
||||
depth_img = add_text(depth_img, label)
|
||||
|
||||
images_to_concat.append(depth_img)
|
||||
|
||||
if valid_group:
|
||||
# 水平拼接
|
||||
try:
|
||||
combined_result = cv2.hconcat(images_to_concat)
|
||||
# 保存
|
||||
save_path = os.path.join(compare_dir, f"{file_name_no_ext}_compare.png")
|
||||
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='Batch process depth estimation (Gray & Color) and compare.')
|
||||
parser.add_argument('--img-path', type=str, required=True, help='Path to input images folder')
|
||||
args = parser.parse_args()
|
||||
|
||||
# 1. 确定根目录名称: 输入文件夹名 + "_results"
|
||||
folder_name = os.path.basename(os.path.normpath(args.img_path)) + "_results"
|
||||
|
||||
encoders = ['vits', 'vitb', 'vitl', 'vitg']
|
||||
|
||||
# 2. 批量生成深度图 (Run Inference)
|
||||
for encoder in encoders:
|
||||
# A. 灰度图
|
||||
gray_out_dir = os.path.join(folder_name, f'V2_{encoder}_gray_result')
|
||||
if not os.path.exists(gray_out_dir):
|
||||
run_encoder(encoder, args.img_path, gray_out_dir, grayscale=True)
|
||||
|
||||
# B. 彩色图
|
||||
color_out_dir = os.path.join(folder_name, f'V2_{encoder}_colorful_result')
|
||||
if not os.path.exists(color_out_dir):
|
||||
run_encoder(encoder, args.img_path, color_out_dir, grayscale=False)
|
||||
|
||||
# 3. 准备原图列表
|
||||
exts = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tif', '*.tiff']
|
||||
image_files = []
|
||||
for ext in exts:
|
||||
image_files.extend(glob.glob(os.path.join(args.img_path, ext)))
|
||||
image_files.extend(glob.glob(os.path.join(args.img_path, '**', ext), recursive=True))
|
||||
image_files = sorted(list(set(image_files)))
|
||||
|
||||
if not image_files:
|
||||
print("No images found in input path.")
|
||||
return
|
||||
|
||||
# 4. 生成灰度对比图
|
||||
create_comparison(
|
||||
image_files=image_files,
|
||||
folder_name=folder_name,
|
||||
encoders=encoders,
|
||||
mode_suffix='_gray_result',
|
||||
out_folder_name='compare_all_gray'
|
||||
)
|
||||
|
||||
# 5. 生成彩色对比图
|
||||
create_comparison(
|
||||
image_files=image_files,
|
||||
folder_name=folder_name,
|
||||
encoders=encoders,
|
||||
mode_suffix='_colorful_result',
|
||||
out_folder_name='compare_all_colorful'
|
||||
)
|
||||
|
||||
print(f"\nAll done! Results are in: {folder_name}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user