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']
        
        depth_map = depth_map.squeeze().cpu().numpy() 
        depth_map = np.uint8(depth_map / depth_map.max() * 255)
        
        depth_image = Image.fromarray(depth_map)
    
        return depth_image