File size: 1,401 Bytes
537fd2c 76d24ba 537fd2c 76d24ba 537fd2c 76d24ba 537fd2c 76d24ba 537fd2c 76d24ba 537fd2c |
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 |
import torch
import numpy as np
from PIL import Image
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
device = None
depth_estimator = None
feature_extractor = None
def init():
global device, depth_estimator, feature_extractor
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Initializing depth estimator...")
depth_estimator = DPTForDepthEstimation.from_pretrained(
"Intel/dpt-hybrid-midas").to(device)
feature_extractor = DPTFeatureExtractor.from_pretrained(
"Intel/dpt-hybrid-midas")
def get_depth_map(image):
original_size = image.size
image = feature_extractor(
images=image, return_tensors="pt").pixel_values.to(device)
with torch.no_grad(), torch.autocast(device):
depth_map = depth_estimator(image).predicted_depth
depth_map = torch.nn.functional.interpolate(
depth_map.unsqueeze(1),
size=original_size[::-1],
mode="bicubic",
align_corners=False,
)
depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
depth_map = (depth_map - depth_min) / (depth_max - depth_min)
image = torch.cat([depth_map] * 3, dim=1)
image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
return image
|