| import gradio as gr | |
| import torch | |
| from utils import colorize | |
| from utils import get_min_and_max | |
| from PIL import Image | |
| import tempfile | |
| import numpy as np | |
| import json | |
| DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| model = torch.hub.load('isl-org/ZoeDepth', "ZoeD_N", pretrained=True).to(DEVICE).eval() | |
| def split_numpy_16_bit(numpy_image_16bit): | |
| result_numpy_split = np.zeros((numpy_image_16bit.shape[0], numpy_image_16bit.shape[1], 3)); | |
| i = 0; | |
| for row in numpy_image_16bit: | |
| j = 0 | |
| for pixel in row: | |
| firstPart = np.array(pixel).astype(np.uint8) | |
| secondPart = np.array(pixel>>8).astype(np.uint8) | |
| result_numpy_split[i, j, 0] = firstPart | |
| result_numpy_split[i, j, 1] = secondPart | |
| j += 1 | |
| i += 1 | |
| return result_numpy_split | |
| def predict(image): | |
| image.thumbnail((1024,1024)) | |
| depth = model.infer_pil(image) | |
| colored_depth, vmin, vmax = colorize(depth, cmap='gray_r') | |
| tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False) | |
| raw_depth = Image.fromarray((depth*256).astype('uint16')) | |
| raw_depth.save(tmp.name) | |
| depth_uint16 = (depth*256).astype('uint16') | |
| numpy_image_split = split_numpy_16_bit(depth_uint16) | |
| depth_image_two_components = Image.fromarray(np.uint8(numpy_image_split)).convert('RGB') | |
| return colored_depth, tmp.name, depth_image_two_components | |
| iface = gr.Interface(fn=predict, inputs=gr.Image(label="Input Image", type='pil'), outputs=["image", "file", "image"]) | |
| iface.launch() |