File size: 1,500 Bytes
99b3515 8bed08d 53d983f 8bed08d bfeb983 99b3515 33fefc7 99b3515 bfeb983 b0b87bf 4cbb052 b0b87bf fef8be5 53d983f bfeb983 85ba2c9 bfeb983 6132da7 bfeb983 81f0805 bfeb983 9b4b602 |
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 |
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() |