Spaces:
Sleeping
Sleeping
File size: 1,786 Bytes
5d92054 |
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 |
"""
Created on Sat Apr 9 04:08:02 2022
@author: Admin_with ODD Team
Edited by our team : Sat Oct 5 10:00 2024
references: https://github.com/vinvino02/GLPDepth
"""
import torch
from transformers import GLPNForDepthEstimation, GLPNFeatureExtractor
from PIL import Image
from config import CONFIG
# GLPDepth Model Class
class GLPDepth:
def __init__(self):
self.feature_extractor = GLPNFeatureExtractor.from_pretrained(CONFIG['glpn_model_path'])
self.model = GLPNForDepthEstimation.from_pretrained(CONFIG['glpn_model_path'])
self.model.to(CONFIG['device'])
self.model.eval()
def predict(self, img: Image.Image, img_shape : tuple):
"""Predict the depth map of the input image.
Args:
img (PIL.Image): Input image for depth estimation.
img_shape (tuple): Original image size (height, width).
Returns:
np.ndarray: The predicted depth map in numpy array format.
"""
with torch.no_grad():
# Preprocess image and move to the appropriate device
pixel_values = self.feature_extractor(img, return_tensors="pt").pixel_values.to(CONFIG['device'])
# Get model output
outputs = self.model(pixel_values)
predicted_depth = outputs.predicted_depth
# Resize depth prediction to original image size
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=img_shape[:2], # Interpolate to original image size
mode="bicubic",
align_corners=False,
)
prediction = prediction.squeeze().cpu().numpy() # Convert to numpy array (shape: (H, W))
return prediction |