Spaces:
Runtime error
Runtime error
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from transformers import DPTFeatureExtractor, DPTForDepthEstimation | |
| depth_estimator = None | |
| feature_extractor = None | |
| def init(): | |
| global depth_estimator, feature_extractor | |
| print("Initializing depth estimator...") | |
| depth_estimator = DPTForDepthEstimation.from_pretrained( | |
| "Intel/dpt-hybrid-midas").to("cuda") | |
| 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("cuda") | |
| with torch.no_grad(), torch.autocast("cuda"): | |
| 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 | |