from PIL import Image
from transformers import pipeline
import base64
import requests
import io
import numpy as np
import os

def convert_img_base64(img):
    with open(img, 'rb') as file:
        base64_image = base64.b64encode(file.read())
    return base64_image



class DepthEstimation():
    def __init__(self, img_base64):
        self.img_base64 = img_base64

    def convert_from_base64(self):
        image_data = base64.b64decode(self.img_base64)
        image = Image.open(io.BytesIO(image_data))
        return image

    def run_depth(self):
        image_data = self.convert_from_base64()
        
        pipe = pipeline("depth-estimation", model="depth-anything/Depth-Anything-V2-Base-hf")
        depth = pipe(image_data)
        depth_map = depth['predicted_depth']
        
        # Convert the tensor to a NumPy array and normalize
        depth_map = depth_map.squeeze().cpu().numpy()  # Remove batch dimension
        depth_map = np.uint8(depth_map / depth_map.max() * 255)  # Normalize to 0-255
        
        # Convert depth map to PIL Image
        depth_image = Image.fromarray(depth_map)
        
        # Save the depth image to a bytes buffer
        buffer = io.BytesIO()
        depth_image.save(buffer, format='PNG')
        image_bytes = buffer.getvalue()
        
        # Encode the bytes to base64
        base64_string = base64.b64encode(image_bytes).decode()
        
        # Decode and display for verification
        # image_data = base64.b64decode(base64_string)
        # image = Image.open(io.BytesIO(image_data))
        # image.show()
        # image.save('output_image.png')

        return base64_string


# if __name__ == "__main__":
#     img_path = 'img3.jpg'
#     base64_image = convert_img_base64(img_path)
#     DepthEstimation(base64_image).run_depth()