| from depth import MidasDepth | |
| import gradio as gr | |
| import numpy as np | |
| import tempfile | |
| depth_estimator = MidasDepth() | |
| def get_depth(rgb): | |
| print("Estimating depth...") | |
| rgb = rgb.convert("RGB") | |
| depth = depth_estimator.get_depth(rgb) | |
| print("Creating mesh...") | |
| w, h = rgb.size | |
| grid = np.mgrid[0:h, 0:w].transpose(1, 2, 0 | |
| ).reshape(-1, 2)[..., ::-1] | |
| flat_grid = grid[:, 1] * w + grid[:, 0] | |
| positions = np.concatenate(((grid - np.array([[w, h]]) | |
| / 2) / w * 2, | |
| depth.flatten()[flat_grid][..., np.newaxis]), | |
| axis=-1) | |
| positions[:, :-1] *= positions[:, -1:] | |
| positions[:, :2] *= -1 | |
| pick_edges = depth < 0 | |
| y, x = (t.flatten() for t in np.mgrid[0:h, 0:w]) | |
| faces = np.concatenate(( | |
| np.stack((y * w + x, | |
| (y - 1) * w + x, | |
| y * w + (x - 1)), axis=-1) | |
| [(~pick_edges.flatten()) * (x > 0) * (y > 0)], | |
| np.stack((y * w + x, | |
| (y + 1) * w + x, | |
| y * w + (x + 1)), axis=-1) | |
| [(~pick_edges.flatten()) * (x < w - 1) * (y < h - 1)] | |
| )) | |
| print("Writing...") | |
| tf = tempfile.NamedTemporaryFile(suffix=".obj").name | |
| save_obj(positions, np.asarray(rgb).reshape(-1, 3) / 255., faces, tf) | |
| return rgb, (depth.clip(0, 64) * 1024).astype("uint16"), tf | |
| def save_obj(positions, rgb, faces, filename): | |
| with open(filename, "w") as f: | |
| for position, color in zip(positions, rgb): | |
| f.write( | |
| f"v {' '.join(map(str, position))} {' '.join(map(str, color))}\n") | |
| for face in faces: | |
| f.write(f"f {' '.join(map(str, face))}\n") | |
| gr.Interface(fn=get_depth, inputs=[ | |
| gr.components.Image(label="rgb", type="pil"), | |
| ], outputs=[ | |
| gr.components.Image(type="pil", label="image"), | |
| gr.components.Image(type="numpy", label="depth"), | |
| gr.components.Model3D(label="3d model") | |
| ]).launch(share=True) | |