|
|
import torch |
|
|
import gradio as gr |
|
|
import numpy as np |
|
|
from PIL import Image |
|
|
import cv2 |
|
|
|
|
|
|
|
|
midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small") |
|
|
midas.eval() |
|
|
|
|
|
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms") |
|
|
transform = midas_transforms.small_transform |
|
|
|
|
|
def predict_depth(image): |
|
|
img = np.array(image) |
|
|
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
|
|
input_tensor = transform(img_rgb).unsqueeze(0) |
|
|
|
|
|
with torch.no_grad(): |
|
|
prediction = midas(input_tensor) |
|
|
prediction = torch.nn.functional.interpolate( |
|
|
prediction.unsqueeze(1), |
|
|
size=img_rgb.shape[:2], |
|
|
mode="bicubic", |
|
|
align_corners=False, |
|
|
).squeeze() |
|
|
|
|
|
depth_map = prediction.cpu().numpy() |
|
|
depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) |
|
|
depth_map = (depth_map * 255).astype(np.uint8) |
|
|
depth_img = Image.fromarray(depth_map) |
|
|
|
|
|
return depth_img |
|
|
|
|
|
|
|
|
iface = gr.Interface( |
|
|
fn=predict_depth, |
|
|
inputs=gr.Image(type="pil"), |
|
|
outputs=gr.Image(type="pil"), |
|
|
title="MiDaS Depth Estimation", |
|
|
description="Загрузите изображение и получите карту глубины." |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
iface.launch() |
