Spaces:
Sleeping
Sleeping
File size: 1,816 Bytes
9266e32 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
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() |