| from PIL import Image | |
| import torch | |
| # custom installation from this PR: https://github.com/huggingface/transformers/pull/34583 | |
| # !pip install git+https://github.com/geetu040/transformers.git@depth-pro-projects#egg=transformers | |
| from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation | |
| # initialize processor and model | |
| checkpoint = "geetu040/DepthPro" | |
| revision = "project" | |
| image_processor = DepthProImageProcessorFast.from_pretrained(checkpoint, revision=revision) | |
| model = DepthProForDepthEstimation.from_pretrained(checkpoint, revision=revision) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = model.to(device) | |
| def predict(image): | |
| # inference | |
| # prepare image for the model | |
| inputs = image_processor(images=image, return_tensors="pt") | |
| inputs = {k: v.to(device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| # interpolate to original size | |
| post_processed_output = image_processor.post_process_depth_estimation( | |
| outputs, target_sizes=[(image.height, image.width)], | |
| ) | |
| # visualize the prediction | |
| depth = post_processed_output[0]["predicted_depth"] | |
| depth = (depth - depth.min()) / depth.max() | |
| depth = depth * 255. | |
| depth = depth.detach().cpu().numpy() | |
| depth = Image.fromarray(depth.astype("uint8")) | |
| return depth | |