import os
from transformers import pipeline
from PIL import Image
from tqdm import tqdm 


class DepthEstimator:
    def __init__(self, model_path="models/Depth-Anything-V2-Small-hf"):
        """
        初始化深度估计器，加载模型

        Args:
            model_path (str): 深度估计模型的路径
        """
        self.pipe = pipeline(task="depth-estimation", model=model_path)

    def estimate_depth(self, image_path):
        """
        估计图像的深度并保存深度图

        Args:
            image_path (str): 输入图像的路径

        Returns:
            PIL.Image.Image: 生成的深度图
        """
        # 打开输入图像
        image = Image.open(image_path)
        
        # 进行深度估计
        depth = self.pipe(image)["depth"]
        
        # 保存深度图
        output_depth_path = f"data/depth_image/depth-{os.path.splitext(os.path.basename(image_path))[0]}.png"
        depth.save(output_depth_path)
        
        return depth

def depth_image_list(image_paths):
    """
    处理图像列表，返回深度图像列表

    Args:
        image_paths (list): 输入图像路径列表

    Returns:
        list: 深度图像列表
    """
    depth_estimator = DepthEstimator()  # 创建深度估计器实例
    depth_images = []

    for image_path in tqdm(image_paths):
        depth_image = depth_estimator.estimate_depth(image_path)  # 估计深度
        depth_images.append(depth_image)  # 将深度图像添加到列表中

    return depth_images

# 使用示例
if __name__ == "__main__":
    #model_path = "/data/data/luosy/project/oral/models/Depth-Anything-V2-Small-hf"  # 模型路径
    input_image_path = 'masked_output_image.png'  # 输入图像路径

    depth_estimator = DepthEstimator()  # 创建深度估计器实例
    depth_image = depth_estimator.estimate_depth(input_image_path)  # 估计深度

    print(type(depth_image))