142 lines
5.2 KiB
Python
142 lines
5.2 KiB
Python
import os
|
||
import numpy as np
|
||
import cv2
|
||
from tool.lighter import adjust_highlights_shadows
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from PIL import Image
|
||
import piexif
|
||
|
||
# parser = argparse.ArgumentParser(description='Depth Anything V2')
|
||
# parser.add_argument('--input-paths', type=str, nargs='+', required=True,
|
||
# help='输入文件夹列表（多个路径，用空格分隔）')
|
||
# parser.add_argument('--output-path', type=str, default='./output',
|
||
# help="按输出路径的结构输出处理好的图片")
|
||
# parser.add_argument('--model-path', type=str, default='./model/depth_anything_v2_vitl.pth',
|
||
# help='模型路径')
|
||
# args = parser.parse_args()
|
||
|
||
|
||
def process_single_image(img_path, depth_anything):
|
||
"""处理单个图片
|
||
Args:
|
||
img_path: 图片路径
|
||
Returns:
|
||
tuple: (处理后的图片numpy数组, 原始图片路径)
|
||
"""
|
||
# 使用OpenCV读取图片
|
||
img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
|
||
if img is None:
|
||
raise ValueError(f"无法读取图片: {img_path}")
|
||
|
||
# 处理图片 - 这里假设adjust_highlights_shadows能处理OpenCV格式的图片
|
||
processed_img = adjust_highlights_shadows(img, 180, 253, depth_anything)
|
||
return processed_img, img_path
|
||
|
||
def save_image_with_exif(input_path, output_path, processed_img):
|
||
"""保存图片并保留EXIF信息
|
||
Args:
|
||
input_path: 原始图片路径
|
||
output_path: 输出图片路径
|
||
processed_img: 处理后的numpy数组图片
|
||
"""
|
||
# 确保输出目录存在
|
||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||
|
||
# 读取原始图片的EXIF信息
|
||
exif_dict = None
|
||
try:
|
||
with exif_lock: # 使用锁防止多线程同时读取EXIF
|
||
with Image.open(input_path) as img:
|
||
if 'exif' in img.info:
|
||
exif_dict = piexif.load(img.info['exif'])
|
||
except Exception as e:
|
||
print(f"Warning: 无法读取 {input_path} 的EXIF信息: {str(e)}")
|
||
|
||
# 将OpenCV格式转换为PIL格式以便保存EXIF
|
||
processed_img_rgb = cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB)
|
||
pil_img = Image.fromarray(processed_img_rgb)
|
||
|
||
# 保存图片
|
||
try:
|
||
if exif_dict:
|
||
pil_img.save(output_path, exif=piexif.dump(exif_dict))
|
||
else:
|
||
pil_img.save(output_path)
|
||
except Exception as e:
|
||
print(f"Error: 无法保存图片 {output_path}: {str(e)}")
|
||
raise
|
||
|
||
def process_single_file(input_path, output_root, input_base, depth_anything):
|
||
"""处理单个文件"""
|
||
try:
|
||
# 计算相对路径以保持文件夹结构
|
||
rel_path = os.path.relpath(os.path.dirname(input_path), start=input_base)
|
||
output_dir = os.path.join(output_root, rel_path)
|
||
output_path = os.path.join(output_dir, os.path.basename(input_path))
|
||
|
||
# 处理图片
|
||
processed_img, _ = process_single_image(input_path, depth_anything)
|
||
|
||
# 保存图片
|
||
save_image_with_exif(input_path, output_path, processed_img)
|
||
print(f"Processed and saved: {output_path}")
|
||
except Exception as e:
|
||
print(f"Error processing {input_path}: {str(e)}")
|
||
|
||
def process_images(input_paths, output_root, depth_anything, workers=4):
|
||
"""处理所有图片并保持原文件夹结构（多线程版本）
|
||
Args:
|
||
input_paths: 输入路径列表
|
||
output_root: 输出根目录
|
||
workers: 线程池大小
|
||
"""
|
||
# 收集所有需要处理的文件
|
||
all_files = []
|
||
for input_path in input_paths:
|
||
for root, dirs, files in os.walk(input_path):
|
||
for file in files:
|
||
if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
|
||
input_file_path = os.path.join(root, file)
|
||
all_files.append((input_file_path, input_path))
|
||
|
||
# 使用线程池处理文件
|
||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||
futures = []
|
||
for file_path, input_base in all_files:
|
||
futures.append(
|
||
executor.submit(
|
||
process_single_file,
|
||
file_path,
|
||
output_root,
|
||
input_base,
|
||
depth_anything
|
||
)
|
||
)
|
||
|
||
# 等待所有任务完成
|
||
for future in futures:
|
||
try:
|
||
future.result()
|
||
except Exception as e:
|
||
print(f"Error in processing: {str(e)}")
|
||
# if __name__ == '__main__':
|
||
# model_path = args.model_path
|
||
# input_paths = args.input_paths
|
||
# output_path = args.output_path
|
||
|
||
# # 线程锁，防止多线程同时访问EXIF数据时出现问题
|
||
# exif_lock = threading.Lock()
|
||
|
||
# input_paths = ["/home/dtyx/桌面/yhh/ReportGenerator/测试数据/山东国华无棣风电场叶片外部数据/A1（31301）",]
|
||
# model_path = "./model/depth_anything_v2_vitl.pth"
|
||
# output_path = "./output"
|
||
# #启动模型
|
||
# depth_anything = model_start(model_path)
|
||
# # 创建输出目录
|
||
# os.makedirs(output_path, exist_ok=True)
|
||
# # 处理所有图片
|
||
# process_images(input_paths, output_path, depth_anything)
|
||
|
||
|
||
|