midas / app.py
adpro's picture
Update app.py
a310278 verified
raw
history blame
811 Bytes
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()