import torch
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
from PIL import Image
import requests

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load model và chuyển sang GPU
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)

# Tải ảnh
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# Chuẩn bị input
inputs = feature_extractor(images=image, return_tensors="pt").to(device)

# Dự đoán depth map
with torch.no_grad():
    outputs = model(**inputs)

# Lấy output và chuyển về dạng ảnh
depth_map = outputs.predicted_depth
depth_map = depth_map.squeeze().cpu().numpy()