Spaces:
Runtime error
Runtime error
Create app.py (#1)
Browse files- Create app.py (f7a645b9b4a002761b191caf58fe3c202404a42c)
app.py
ADDED
@@ -0,0 +1,658 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os, sys
|
2 |
+
import subprocess
|
3 |
+
from huggingface_hub import snapshot_download
|
4 |
+
|
5 |
+
is_local_run = False
|
6 |
+
|
7 |
+
code_dir = snapshot_download("One-2-3-45/code") if not is_local_run else "../code" # , token=os.environ['TOKEN']
|
8 |
+
|
9 |
+
sys.path.append(code_dir)
|
10 |
+
|
11 |
+
elev_est_dir = os.path.abspath(os.path.join(code_dir, "one2345_elev_est"))
|
12 |
+
sys.path.append(elev_est_dir)
|
13 |
+
|
14 |
+
if not is_local_run:
|
15 |
+
# import pip
|
16 |
+
# pip.main(['install', elev_est_dir])
|
17 |
+
# export TORCH_CUDA_ARCH_LIST="7.0;7.2;8.0;8.6"
|
18 |
+
# export IABN_FORCE_CUDA=1
|
19 |
+
os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
|
20 |
+
os.environ["IABN_FORCE_CUDA"] = "1"
|
21 |
+
os.environ["FORCE_CUDA"] = "1"
|
22 |
+
# pip.main(["install", "inplace_abn"])
|
23 |
+
subprocess.run(['pip', 'install', 'inplace_abn'])
|
24 |
+
# FORCE_CUDA=1 pip install --no-cache-dir git+https://github.com/mit-han-lab/[email protected]
|
25 |
+
# pip.main(["install", "--no-cache-dir", "git+https://github.com/mit-han-lab/[email protected]"])
|
26 |
+
subprocess.run(['pip', 'install', '--no-cache-dir', 'git+https://github.com/mit-han-lab/[email protected]'])
|
27 |
+
|
28 |
+
import shutil
|
29 |
+
import torch
|
30 |
+
import fire
|
31 |
+
import gradio as gr
|
32 |
+
import numpy as np
|
33 |
+
import plotly.graph_objects as go
|
34 |
+
from functools import partial
|
35 |
+
|
36 |
+
import cv2
|
37 |
+
from PIL import Image
|
38 |
+
import trimesh
|
39 |
+
import tempfile
|
40 |
+
from zero123_utils import init_model, predict_stage1_gradio, zero123_infer
|
41 |
+
from sam_utils import sam_init, sam_out_nosave
|
42 |
+
from utils import image_preprocess_nosave, gen_poses
|
43 |
+
from one2345_elev_est.tools.estimate_wild_imgs import estimate_elev
|
44 |
+
from rembg import remove
|
45 |
+
|
46 |
+
_GPU_INDEX = 0
|
47 |
+
|
48 |
+
_TITLE = '''One-2-3-45: Any Single Image to 3D Mesh in 45 Seconds without Per-Shape Optimization'''
|
49 |
+
|
50 |
+
|
51 |
+
_DESCRIPTION = '''
|
52 |
+
<div>
|
53 |
+
<a style="display:inline-block" href="http://one-2-3-45.com"><img src="https://img.shields.io/badge/Project_Homepage-f9f7f7?logo=data:image/webp;base64,UklGRmIRAABXRUJQVlA4IFYRAABQPwCdASrIAMgAPm00lEekpiolqDvpMIANiWJu3pE7maI+vTDkhN5f7PfmGT7nS6p8nKBr0I+YBzr/ML+2/rG/8j1Sf3/1AP6v/sOsW/bn2GP2W9Zv/zeyf/cf+Z+63tReoB/6uCx2p5E/iUkPwG1FO9t/XgHPL7RH7TzI8EvEA4Mr1T2CP0J6xH+d5J/rz2F/LG9h37l+x9+4xQ3m86D2Te/zeVV/tWyTw7s85XZ0ABD4N2CpzWHt8feKiWkqdTkRjojREWrbUDAKXlYsV7EGU9rWR2gCxVXnstqpNVhwra603swvYRlMyRzKc5nJHEEeLuV8EDee/MpPVIq2DNUcXpCZXvFLHgXBWvZWzSZCFo4iub8df+Yu9q7rw5qemOe2Nt1IIoyjBmXdjCunMherehhPjQIQGiI6PDcriy/zhhwHE7O+0gmpUsYmcDR+ixOfLPY0yjnQosZkIoK1pfttDGtirMbSMqndDVi73JMxcSlNb0MNFtgdAAXNk5Z77wgsPz9RRj9oWO/KRpXn5ra4gUt+mMgSCvFG86zgghSehTRD54z10sNxqnG3/rpKDifOvT4EQU1uA9ZckUZcUt5L5C0+dOdj1I56uLJEAsn0432gHD5wRG7dgSfYusXvhGl2uMaczlXSJ0JfX+Z0e9q7sHywvyEkWC+iJREwvtWi1K+NQAD+/JSRhGP+QTeW9xU73vXKZO+JaR/TAb6vV9dNzIjket6jYZdxK0qCcaf95QeouegLeSQL/9WeH5l2/DE2AKdYjhEYzzIefp7c6cTfM3D3q3kSFxAF/xP/f/3YUFjjOzfzl5xrD3XaWz0TAehn6+ze5pANq6t5GDX8ZOfpIBGUplJj6UZXd76ropLkDdM+d/F2Megl53hry7QvtcUGNlKgjLd7/txvzvkYIPre5sKVvAJzj9DEml706Piekk2NTtBnCMQtQAPO7/Soo3p3QbqLnMIY2PKCq3jFUkeMDAB6uvaHy7e8G/yi+LlFCfYgju+h+ha+jj6NYh6xUx/9TpQoQ1VFrpEw7pCAaQ2NbzVcj/EfBLQUWQBwliZd6FG70L3ATK7AS/cu+Pm/ASndDhIDTx08uveDvY2kW7Mqproq8D4ImWzJ7ZwM8JfrvyN9/wh0Iu00O3UbTDU58dYfWzxI1gDb2Yt6+AyvgjRY/WUM8aikx5MTFi6ZEWeffMc8ruwWeKmfwJtpDxNYhJgSN5gZoOS+XedZmwoYfiuaf9hhPdDtJCM429liA9mZQ2GNfMOPtcLJV/4xTUuWJx4/d43remtcIdsy1GlD79SNNSlfWuSCF8LIopGEcQwZnBVOBmJ7O2bQvBHNAQ6dlz+Pc8zL7MgsN7uff5PDGyXHqV4lCh2Q/xbeZaYVv1agSO4QU5FvEX/6AQpNxQqgtvrXM+XsXTWYJpiO7+ucPvPudldDswT7/ITpp7AdSJ9OjPYr3cMRBVy5sXXkyY8SVv0z//QqGJbxMA3IV81dfN5dUNNlMvl+EBv6Qrnq42ZAEXMEbW/zcmuvIO+539I0BKM+COuGTuEmhulQBlMmdlBNVaII5lFuENjHGpIPNULKARA/PhM9jOvRG2xs8SFCjLZ9ZNLyWznJxq3sZWeGUNXcyJPKDr3WAWCmP8tNyjKQk6aKOD1s/+2MCNQ9b4Zb2XJJLW4kBOS6P10n42Scz8D1K6GTeouELbGApoVNEYjw2jdbNqsNLZiJi6XfSs7Kz5ozupOLJsYOzfLQTjDz7BZOzEYFoB+acySl5Qs3mO84Mw6pIb9MqHhzFqtmrsD4V6uwssP9SUvJTMA4orRWZAMdCA9lMHQi6XBPR9CBidQdudyeYtUj5gWMyzdcSkG8l/2zbzpz8THi23y/8+Ijg5naj6GtYnpJna+6QpI6dlY+E2KF7bOK2ctKwBjClUvAjHXa1162i6DsToHLQE4stmDJdvI1POb9Hj0Mq+azo1wrOfqVFcAS5XNc37IJeYBs/cQYZ08mg2vXWWJYVWz648jTHABHf+LiHsy4WRaVo4oOOSyeampoUYSM9WUJ3iOlTMis5U2DCrGoAiATOAyyuwMcYgTni5FGSpdE5BnoS6ORUiYapPetM/XmsvikTkKNn4z4jhiLFFcU+bH1pZ2DseVK9vCgY5s9ZDjNb9Ky+8fwn9dJtsZ6M7opvXhqde9Ljos6KWQ/8hj3pswa2lLZ7WRc9xaxTjq1sytCxfOd+J+VrsXedNuuYDMwumYIzF1Xsbz1VCURDw6C1twAPizF49s4JfToWsGhgG4wtBE5NAU4KvnGleFGzv54AwBR9qqPVD9bcN7ZmZhphTcAGnR2oOcvT98FmknnVXqJYdHSeP9nWG6A8YUUgDmy7rYLtbxgpid5ysrxvLeyMOeTaQXuNZXn5kQeqDGELMfQ5U2PAk+/DhnbWTsirG5NAP0tgGbSNGM5cv9+trgSk6sXdw1lZOLrfqLGZ8Dt19DWcxmjsbDU30CoSc1alYxX5G+uIHy72tQxjzsot1O4iZeNO34PItzBoeg0Fq+YQZGsHdknwJkAbffRl96xFsDxM6l4g22OaMEHxLMC9uFFE8Ee/xf+krkjv7YCfJrCM3Nw6xfyrhtxN3x2GxSg4YTu2dtjb3zVI/6jYNwgGiaDWoh5I29uQ8ZvES8Ros5jgxDzeKB1tJ3HtDM9SFGNJfQiLiSyYZQLBjCcGbi3+vlythB3k6af+P5rDqah2oPFl29Ngnw/tmpkmRIvri5i55FPeY9J4nXfvWYHTHdoB0oVA2NEk2nropP+T7GXhAxA2NgyGtzHaVU2yxiSju87w8MLIo1eac26wOnbEo/oD6Zcb8vyu0x7ug9iERQ5FlppDnIktT6QC6Kk3qBxovLzOPdEvYQoytaN256n2dmkxAaq78klv6PnU7HiH3e/I9RC27VOP0j6JDW19KvC9/uN9tfOi6WMr0IGKpTsZAUZXm+Ukyk/Rpu9ZPIH5/3CL+yfj3ROts+BWIZNj8lpFHfmYhmN/J0+/lDIGmbRVMbvmif9tqr53fqb8EkFpdMHnK8jc0oIYu2Io5SWOzHc7GMdwt5RB8LR5jUjV6Xv+rR7N4IcTQRphe7WarFsxHmihpNr8sLroqsVxBH+6dOjC5DPhRV6aJB9ZB0NjpLtbjRsEKe1By1huo8rJa+DS73fTUfxWPaJjQsasBOcc6xwuob3OBjTFjUOxfiBbhMDNUFcamlMphrkbmTe2smHz0hrScXZjoHxphV537e8PNenBpI//N58bUOcmV4Lo1H1BLLjNTw1gK+rKFgaU/WOZQ0DZ1kRRqCa86XYnFposmkLgDNooS/yeW/RGfvopRDH40d2TeW8t1+2fDHQcwocSXolq+dxC6JMGsu2rCrhdjzhqd1KPMp5EVGQuCyLc8LfjUhQ8fSs63P9aVDYZBDhO8oWSI2Lbk7cRpKJ38ww9dD0b0OjvucHkJl1zIwyQFqKKEfIN7RPvV8Q1Xxot6Y5f8/UqOCOVZRt+IM1JFcJ4AstPMOXs9hAyZzEs1EY9lv3976/18LNNvL8K7RPNH1uz3qwAajMXLaOTEK7IzCjex3YZQ0LCICPzWVKMNbkSFpmy5ow1A54fK4F45T0apL1FE8dc/Jy6ERymiJ8ZvT+BJHUtbS5oB72w8NeIb0zTuqTzYwMQiKeCI+DlJTd6R3dgbvDETb7XtLT1L5quVxBiyJLxgARoeWU1DY3eWTFJkicFp/UqIFCYgLUhQgGm/1gAxylWf4wZmbQy6RGlY3/pfn5qxqFq8Xmza7Unght3AckydGZ6u6yWcooxZwILsHaklA/Bu2HRlCLzLer57IQWfvHUjJ8pqEoZ/TE0WqZc4SF6CBVC4KGEIqyPnH/+chaIQRfGuKg0rKAAc5tB+7vGl4ck72A+dA9iW0UUwXqD6Y333q9MEdov8hbXuiRkRMv1CEm0h1N8yhxOEe1SLWxlCmvUHcVvhojM6S4ODYr2rxHxOqx63MVVCk6PpQAB2gn4D/9+QHVBBqAxLV8Pggh6aRlEmPuHNEc+b/1Zqh4lvwxUgyMFngOgTAhqZAZqBpRRD41KfU7wEbzruYQhOIxxPMbGIe93LRFgYZLz21sLmS/02JhQ1eY6cSu/iumOWleWujzjCcB9qxDpOBlsugpiveHwunOO9Lu66uGNeKuw6Aqo+zTcdEX+BOlserNtSeyYmhQrwLA+mnEqaAtKv5eAyTC03krdlSEI+++xVMU+kqsGF+6H9yNBQj5aZxmOzd7BejBdBBInEjlj868zR80jlgVKb+yQ7XkdiFIvQl/XvaFpPGqYb4UR70U0jNe/I9UuFggu8M90wyOi7Ihm3t9FZTBPv4zmee4ue5pVKpdsOOZLwBSODTdpUJb6ctU3I9n0KAUBM4+RkNoQzCyb+iXoXl22CL2eQWlOlBi8IG84Y2bMIiLnPs5qeUth9zlniN14oQNTtVJibuIgkylT50ExHyuqz1ra2+wW3QDltErT6yyrKnL8rmkPesI3aPAL880z4U6TWXqcU6hkryL8W5gdI94KYuDTBEim0GM6IAAKf8JZNX3sM/OIB9h3XbFUuNXRocJY9iqQAGdinm3YPLbRBxP5S5EWwlTdIVK5yjUpV+tCN0HXOVf7xj9pnyIMPDz/Znf4zufz+0zonywFQLgAdKiVwC5a6EDC2rmxYC4L82QIO17SKc8NCAJZPTWwwrPGgb0nhQdi3g32QzHUAqE2qhq2jyM7WINI34P28PN5IE50uRx/XFn2a8h2Qgla55PnsIT7KbDBo0Nd4XUkCWINxReQK0/NZEDZrUuZghyZYnnIuIi0pTpecJWliTLvfxyiRkIsb9t2mT6VzM8H2HN8nq0rF7BC27r0JoLl/5YgZQZmw763cQ625wkmPOX0vr1M35fZYv06zKm1ux/L+W6O3ju3VdFudKgEgRIeT+bIOQKoKaT+knRugmPDGt1JAt6bKTT2bvIYnf5OvZs9id5x+qy5UeotL3uxYiBj7SyGxTCHdovbak2BG5hGmuVWxRojEJS9IEqUKwy133zg24keiFy0bXsG125D1XgQ/uI3IM8dijJ4N6jHObWneJl3zHvKb+cX97XFAv5VV5ySEfm0Iglkir9QaTOXP9SuTeCao7Q3fULO6Jcp+sOU6x+jCjlsmHiF7oDCAb/RYITB9oIcAGGQJkT6ccTNJPlyPiNf4/nefe5V4j5RWLTDm7Wb/kt426CIGzE2ekBjrvExlch914MGYjMJcBllUj/LTfXTYLSqPwPzU/xBVSUOR9o7wFnFBZTaI0ApgY11rRsoEgTu9yRgBxw0h71O/RjpN5Ku/U/er87C9/jHzucfXpRDcP1JxOOxoJziSE01YxnPjmyDigmCcur63bY/xXdZeNQNprWvE3mAIP14fFkdJ4+0vwkAP+BXokPPQBkZuEWFAUEz1H/YQf4Q9bCQZXl/WSpUpG/TjBo8EpZLTJ2Jwa1G3H2hVIUlifUnV/SvKDYbpUvl6mKuwdgglJxkJOXjtf84FjvjeHUOzf8ZhLw3PH53rUrDz0INySaGJ/n4a/iuvLMaL146Ii98kND4sM0nElTIxnJe+LJF/8aimynAAiTshwnKc7MHqCtuDaUFfEQCGw0tmys5yKZM5zzawr6LW7MdQs9XUDiyTrX+YcI0uPZZ43oMnO737u5Tmc/sAeKCNGIWt8kw87EMQ+BP5NMrf8b8wDvSD2XVEZu7xqwNCizeSYQJGVJVSAdJ27XwXrFfHtdHHrlojW+3BFzE5rOzDsUsA00zYHxt+e9zo9Yn0sImcxGhbDFBGD892Rgz9G+eor3huRF8h4p1qYpjTe/ykVkhWyvHRjNNevOV7Gk1jhgiwOajzrXwNsIJNUvAQQB017GRYgey7MAEeBoAx5RuxYU+oMH6DNk5eYcrJDxo48XGbO4QhCMRgAA"></a>
|
54 |
+
<a style="display:inline-block; margin-left: .5em" href="https://arxiv.org/abs/2306.16928"><img src="https://img.shields.io/badge/2306.16928-f9f7f7?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADcAAABMCAYAAADJPi9EAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAAuIwAALiMBeKU/dgAAABl0RVh0U29mdHdhcmUAd3d3Lmlua3NjYXBlLm9yZ5vuPBoAAAa2SURBVHja3Zt7bBRFGMAXUCDGF4rY7m7bAwuhlggKStFgLBgFEkCIIRJEEoOBYHwRFYKilUgEReVNJEGCJJpehHI3M9vZvd3bUP1DjNhEIRQQsQgSHiJgQZ5dv7krWEvvdmZ7d7vHJN+ft/f99pv5XvOtJMFCqvoCUpTdIEeRLC+L9Ox5i3Q9LACaCeK0kXoSChVcD3C/tQPHpAEsquQ73IkUcEz2kcLCknyGW5MGjkljRFVL8xJOKyi4CwCOuQAeAkfTP1+tNxLkogvgEbDgffkJqKqvuMA5ifOpqg/5qWecRstNg7xoUTI1Fovdxg8oy2s5AP8CGeYHmGngeZaOL4I4LXLcpHg4149/GDz4xqgsb+UAbMKKUpkrqHA43MUyyJpWUK0EHeG2YKRXr7tB+QMcgGewLD+ebTDbtrtbBt7UPlhS4rV4IvcDI7J8P1OeA/AcAI7LHljN7aB8XTowJmZt9EFRD/o0SDMH4HlwMhMyDWZZSAHFf3YDs3RS49WDLuaAY3IJq+qzmQKLxXAZKN7oDoYbdV3v5elPqiSpMyiOuAEVZVqHXb1OhloUH+MA+ztO0cAO/RkrfyBE7OAEbAZvO8vzVtTRWFD6DAfY5biBM3PWiaL0a4lvXICwnV8WjmE6ntYmhqX2jjp5LbMZjCw/wbYeN6CizOa2GMVzQOlmHjB4Ceuyk6LJ8huccEmR5Xddg7OOV/NAtchW+E3XbOag60QA4Qwuarca0bRuEJyr+cFQwzcY98huxhAKdQelt4kAQpj4qJ3gvFXAYn+aJumXk1yPlpQUgtIHhbYoFMUstNRRWgjnpl4A7IKlayNymqFHFaWCpV9CFry3LGxR1CgA5kB5M8OX2goApwpaz6mdOMGxtAgXWJySxb4WuQD4qTDgU+N5AAnzpr7ChSWpCyisiQJqY0Y7FtmSKpbV23b45kC0KHBxcQ9QeI8w4KgnHRPVtIU7rOtbioLVg5Hl/qDwSVFAMqLSMSObroCdZYlzIJtMRFVHCaRo/wFWPgaAXzdbBpkc2A4aKzCNd97+URQuESYGDDhIVfWOQIKZJu4D2+oXlgDTV1865gUQZDts756BArMNMoR1oa46BYqbyPixZz1ZUFV3sgwoGBajuBKATl3btIn8QYYMuezRgrsiRUWyr2BxA40EkPMpA/Hm6gbUu7fjEXA3azP6AsbKD9bxdUuhjM9W7fII52BF+daRpE4+WA3P501+jbfmHvQKyFqMuXf7Ot4mkN2fr50y+bRH61X7AXdUpHSxaPQ4GVbR5AGw3g+434XgQGKfr72I+vQRhfsu92dOx7WicInzt3CBg1RVpMm0NveWo2SqFzgmdNZMbriILD+S+zoueWf2vSdAipzacWN5nMl6XxNlUHa/J8DoJodUDE0HR8Ll5V0lPxcrLEHZPV4AzS83OLis7FowVa3RSku7BSNxJqQAlN3hBTC2apmDSkpaw22wJemGQFUG7J4MlP3JC6A+f96V7vRyX9It3nzT/GrjIU8edM7rMSnIi10f476lzbE1K7yEiEuWro0OJBguLCwDuFOJc1Na6sRWL/cCeMIwUN9ggSVbe3v/5/EgzTKWLvEAiBrYRUkgwNI2ZaFQNT75UDxEUEx97zYnzpmiLEmbaYCbNxYtFAb0/Z4AztgUrhyxuNgxPnhfHFDHz/vTgFWUQZxTRkkJhQ6YNdVUEPAfO6ZV5BRss6LcCVb7VaAma9giy0XJZBt9IQh42NY0NSdgbLIPlLUF6rEdrdt0CUCK1wsCbkcI3ZSLc7ZSwGLbmJXbPsNxnE5xilYKAobZ77LpGZ8TAIun+/iCKQoF71IxQDI3K2CCd+ARNvXg9sykBcnHAoCZG4u66hlDoQLe6QV4CRtFSxZQ+D0BwNO2jgdkzoGoah1nj3FVlSR19taTSYxI8QLut23U8dsgzqHulJNCQpcqBnpTALCuQ6NSYLHpmR5i42gZzuIdcrMMvMJbQlxe3jXxyZnLACl7ARm/FjPIDOY8ODtpM71sxwfcZpvBeUzKWmfNINM5AS+wO0Khh7dMqKccu4+qatarZjYAwDlgetzStHtEt+XedsBOQtU9XMrRgjg4KTnc5nr+dmqadit/4C4uLm8DuA9koJTj1TL7fI5nDL+qqoo/FLGAzL7dYT17PzvAcQONYSUQRxW/QMrHZVIyik0ZuQA2mzp+Ji8BW4YM3Mbzm9inaHkJCGfrUZZjujiYailfFwA8DHIy3acwUj4v9vUVa+SmgNsl5fuyDTKovW9/IAmfLV0Pi2UncA515kjYdrwC9i9rpuHiq3JwtAAAAABJRU5ErkJggg=="></a>
|
55 |
+
<a style="display:inline-block; margin-left: .5em" href='https://github.com/One-2-3-45/One-2-3-45'><img src='https://img.shields.io/github/stars/One-2-3-45/One-2-3-45?style=social' /></a>
|
56 |
+
</div>
|
57 |
+
We reconstruct a 3D textured mesh from a single image by initially predicting multi-view images and then lifting them to 3D.
|
58 |
+
'''
|
59 |
+
_USER_GUIDE = "Please upload an image in the block above (or choose an example above) and click **Run Generation**."
|
60 |
+
_BBOX_1 = "Predicting bounding box for the input image..."
|
61 |
+
_BBOX_2 = "Bounding box adjusted. Continue adjusting or **Run Generation**."
|
62 |
+
_BBOX_3 = "Bounding box predicted. Adjust it using sliders or **Run Generation**."
|
63 |
+
_SAM = "Preprocessing the input image... (safety check, SAM segmentation, *etc*.)"
|
64 |
+
_GEN_1 = "Predicting multi-view images... (may take \~13 seconds) <br> Images will be shown in the bottom right blocks."
|
65 |
+
_GEN_2 = "Predicting nearby views and generating mesh... (may take \~33 seconds) <br> Mesh will be shown on the right."
|
66 |
+
_DONE = "Done! Mesh is shown on the right. <br> If it is not satisfactory, please select **Retry view** checkboxes for inaccurate views and click **Regenerate selected view(s)** at the bottom."
|
67 |
+
_REGEN_1 = "Selected view(s) are regenerated. You can click **Regenerate nearby views and mesh**. <br> Alternatively, if the regenerated view(s) are still not satisfactory, you can repeat the previous step (select the view and regenerate)."
|
68 |
+
_REGEN_2 = "Regeneration done. Mesh is shown on the right."
|
69 |
+
|
70 |
+
|
71 |
+
def calc_cam_cone_pts_3d(polar_deg, azimuth_deg, radius_m, fov_deg):
|
72 |
+
'''
|
73 |
+
:param polar_deg (float).
|
74 |
+
:param azimuth_deg (float).
|
75 |
+
:param radius_m (float).
|
76 |
+
:param fov_deg (float).
|
77 |
+
:return (5, 3) array of float with (x, y, z).
|
78 |
+
'''
|
79 |
+
polar_rad = np.deg2rad(polar_deg)
|
80 |
+
azimuth_rad = np.deg2rad(azimuth_deg)
|
81 |
+
fov_rad = np.deg2rad(fov_deg)
|
82 |
+
polar_rad = -polar_rad # NOTE: Inverse of how used_x relates to x.
|
83 |
+
|
84 |
+
# Camera pose center:
|
85 |
+
cam_x = radius_m * np.cos(azimuth_rad) * np.cos(polar_rad)
|
86 |
+
cam_y = radius_m * np.sin(azimuth_rad) * np.cos(polar_rad)
|
87 |
+
cam_z = radius_m * np.sin(polar_rad)
|
88 |
+
|
89 |
+
# Obtain four corners of camera frustum, assuming it is looking at origin.
|
90 |
+
# First, obtain camera extrinsics (rotation matrix only):
|
91 |
+
camera_R = np.array([[np.cos(azimuth_rad) * np.cos(polar_rad),
|
92 |
+
-np.sin(azimuth_rad),
|
93 |
+
-np.cos(azimuth_rad) * np.sin(polar_rad)],
|
94 |
+
[np.sin(azimuth_rad) * np.cos(polar_rad),
|
95 |
+
np.cos(azimuth_rad),
|
96 |
+
-np.sin(azimuth_rad) * np.sin(polar_rad)],
|
97 |
+
[np.sin(polar_rad),
|
98 |
+
0.0,
|
99 |
+
np.cos(polar_rad)]])
|
100 |
+
|
101 |
+
# Multiply by corners in camera space to obtain go to space:
|
102 |
+
corn1 = [-1.0, np.tan(fov_rad / 2.0), np.tan(fov_rad / 2.0)]
|
103 |
+
corn2 = [-1.0, -np.tan(fov_rad / 2.0), np.tan(fov_rad / 2.0)]
|
104 |
+
corn3 = [-1.0, -np.tan(fov_rad / 2.0), -np.tan(fov_rad / 2.0)]
|
105 |
+
corn4 = [-1.0, np.tan(fov_rad / 2.0), -np.tan(fov_rad / 2.0)]
|
106 |
+
corn1 = np.dot(camera_R, corn1)
|
107 |
+
corn2 = np.dot(camera_R, corn2)
|
108 |
+
corn3 = np.dot(camera_R, corn3)
|
109 |
+
corn4 = np.dot(camera_R, corn4)
|
110 |
+
|
111 |
+
# Now attach as offset to actual 3D camera position:
|
112 |
+
corn1 = np.array(corn1) / np.linalg.norm(corn1, ord=2)
|
113 |
+
corn_x1 = cam_x + corn1[0]
|
114 |
+
corn_y1 = cam_y + corn1[1]
|
115 |
+
corn_z1 = cam_z + corn1[2]
|
116 |
+
corn2 = np.array(corn2) / np.linalg.norm(corn2, ord=2)
|
117 |
+
corn_x2 = cam_x + corn2[0]
|
118 |
+
corn_y2 = cam_y + corn2[1]
|
119 |
+
corn_z2 = cam_z + corn2[2]
|
120 |
+
corn3 = np.array(corn3) / np.linalg.norm(corn3, ord=2)
|
121 |
+
corn_x3 = cam_x + corn3[0]
|
122 |
+
corn_y3 = cam_y + corn3[1]
|
123 |
+
corn_z3 = cam_z + corn3[2]
|
124 |
+
corn4 = np.array(corn4) / np.linalg.norm(corn4, ord=2)
|
125 |
+
corn_x4 = cam_x + corn4[0]
|
126 |
+
corn_y4 = cam_y + corn4[1]
|
127 |
+
corn_z4 = cam_z + corn4[2]
|
128 |
+
|
129 |
+
xs = [cam_x, corn_x1, corn_x2, corn_x3, corn_x4]
|
130 |
+
ys = [cam_y, corn_y1, corn_y2, corn_y3, corn_y4]
|
131 |
+
zs = [cam_z, corn_z1, corn_z2, corn_z3, corn_z4]
|
132 |
+
|
133 |
+
return np.array([xs, ys, zs]).T
|
134 |
+
|
135 |
+
class CameraVisualizer:
|
136 |
+
def __init__(self, gradio_plot):
|
137 |
+
self._gradio_plot = gradio_plot
|
138 |
+
self._fig = None
|
139 |
+
self._polar = 0.0
|
140 |
+
self._azimuth = 0.0
|
141 |
+
self._radius = 0.0
|
142 |
+
self._raw_image = None
|
143 |
+
self._8bit_image = None
|
144 |
+
self._image_colorscale = None
|
145 |
+
|
146 |
+
def encode_image(self, raw_image, elev=90):
|
147 |
+
'''
|
148 |
+
:param raw_image (H, W, 3) array of uint8 in [0, 255].
|
149 |
+
'''
|
150 |
+
# https://stackoverflow.com/questions/60685749/python-plotly-how-to-add-an-image-to-a-3d-scatter-plot
|
151 |
+
|
152 |
+
dum_img = Image.fromarray(np.ones((3, 3, 3), dtype='uint8')).convert('P', palette='WEB')
|
153 |
+
idx_to_color = np.array(dum_img.getpalette()).reshape((-1, 3))
|
154 |
+
|
155 |
+
self._raw_image = raw_image
|
156 |
+
self._8bit_image = Image.fromarray(raw_image).convert('P', palette='WEB', dither=None)
|
157 |
+
# self._8bit_image = Image.fromarray(raw_image.clip(0, 254)).convert(
|
158 |
+
# 'P', palette='WEB', dither=None)
|
159 |
+
self._image_colorscale = [
|
160 |
+
[i / 255.0, 'rgb({}, {}, {})'.format(*rgb)] for i, rgb in enumerate(idx_to_color)]
|
161 |
+
self._elev = elev
|
162 |
+
# return self.update_figure()
|
163 |
+
|
164 |
+
def update_figure(self):
|
165 |
+
fig = go.Figure()
|
166 |
+
|
167 |
+
if self._raw_image is not None:
|
168 |
+
(H, W, C) = self._raw_image.shape
|
169 |
+
|
170 |
+
x = np.zeros((H, W))
|
171 |
+
(y, z) = np.meshgrid(np.linspace(-1.0, 1.0, W), np.linspace(1.0, -1.0, H) * H / W)
|
172 |
+
|
173 |
+
angle_deg = self._elev-90
|
174 |
+
angle = np.radians(90-self._elev)
|
175 |
+
rotation_matrix = np.array([
|
176 |
+
[np.cos(angle), 0, np.sin(angle)],
|
177 |
+
[0, 1, 0],
|
178 |
+
[-np.sin(angle), 0, np.cos(angle)]
|
179 |
+
])
|
180 |
+
# Assuming x, y, z are the original 3D coordinates of the image
|
181 |
+
coordinates = np.stack((x, y, z), axis=-1) # Combine x, y, z into a single array
|
182 |
+
# Apply the rotation matrix
|
183 |
+
rotated_coordinates = np.matmul(coordinates, rotation_matrix)
|
184 |
+
# Extract the new x, y, z coordinates from the rotated coordinates
|
185 |
+
x, y, z = rotated_coordinates[..., 0], rotated_coordinates[..., 1], rotated_coordinates[..., 2]
|
186 |
+
|
187 |
+
fig.add_trace(go.Surface(
|
188 |
+
x=x, y=y, z=z,
|
189 |
+
surfacecolor=self._8bit_image,
|
190 |
+
cmin=0,
|
191 |
+
cmax=255,
|
192 |
+
colorscale=self._image_colorscale,
|
193 |
+
showscale=False,
|
194 |
+
lighting_diffuse=1.0,
|
195 |
+
lighting_ambient=1.0,
|
196 |
+
lighting_fresnel=1.0,
|
197 |
+
lighting_roughness=1.0,
|
198 |
+
lighting_specular=0.3))
|
199 |
+
|
200 |
+
scene_bounds = 3.5
|
201 |
+
base_radius = 2.5
|
202 |
+
zoom_scale = 1.5 # Note that input radius offset is in [-0.5, 0.5].
|
203 |
+
fov_deg = 50.0
|
204 |
+
edges = [(0, 1), (0, 2), (0, 3), (0, 4), (1, 2), (2, 3), (3, 4), (4, 1)]
|
205 |
+
|
206 |
+
input_cone = calc_cam_cone_pts_3d(
|
207 |
+
angle_deg, 0.0, base_radius, fov_deg) # (5, 3).
|
208 |
+
output_cone = calc_cam_cone_pts_3d(
|
209 |
+
self._polar, self._azimuth, base_radius + self._radius * zoom_scale, fov_deg) # (5, 3).
|
210 |
+
output_cones = []
|
211 |
+
for i in range(1,4):
|
212 |
+
output_cones.append(calc_cam_cone_pts_3d(
|
213 |
+
angle_deg, i*90, base_radius + self._radius * zoom_scale, fov_deg))
|
214 |
+
delta_deg = 30 if angle_deg <= -15 else -30
|
215 |
+
for i in range(4):
|
216 |
+
output_cones.append(calc_cam_cone_pts_3d(
|
217 |
+
angle_deg+delta_deg, 30+i*90, base_radius + self._radius * zoom_scale, fov_deg))
|
218 |
+
|
219 |
+
cones = [(input_cone, 'rgb(174, 54, 75)', 'Input view (Predicted view 1)')]
|
220 |
+
for i in range(len(output_cones)):
|
221 |
+
cones.append((output_cones[i], 'rgb(32, 77, 125)', f'Predicted view {i+2}'))
|
222 |
+
|
223 |
+
for idx, (cone, clr, legend) in enumerate(cones):
|
224 |
+
|
225 |
+
for (i, edge) in enumerate(edges):
|
226 |
+
(x1, x2) = (cone[edge[0], 0], cone[edge[1], 0])
|
227 |
+
(y1, y2) = (cone[edge[0], 1], cone[edge[1], 1])
|
228 |
+
(z1, z2) = (cone[edge[0], 2], cone[edge[1], 2])
|
229 |
+
fig.add_trace(go.Scatter3d(
|
230 |
+
x=[x1, x2], y=[y1, y2], z=[z1, z2], mode='lines',
|
231 |
+
line=dict(color=clr, width=3),
|
232 |
+
name=legend, showlegend=(i == 1) and (idx <= 1)))
|
233 |
+
|
234 |
+
# Add label.
|
235 |
+
if cone[0, 2] <= base_radius / 2.0:
|
236 |
+
fig.add_trace(go.Scatter3d(
|
237 |
+
x=[cone[0, 0]], y=[cone[0, 1]], z=[cone[0, 2] - 0.05], showlegend=False,
|
238 |
+
mode='text', text=legend, textposition='bottom center'))
|
239 |
+
else:
|
240 |
+
fig.add_trace(go.Scatter3d(
|
241 |
+
x=[cone[0, 0]], y=[cone[0, 1]], z=[cone[0, 2] + 0.05], showlegend=False,
|
242 |
+
mode='text', text=legend, textposition='top center'))
|
243 |
+
|
244 |
+
# look at center of scene
|
245 |
+
fig.update_layout(
|
246 |
+
# width=640,
|
247 |
+
# height=480,
|
248 |
+
# height=400,
|
249 |
+
height=450,
|
250 |
+
autosize=True,
|
251 |
+
hovermode=False,
|
252 |
+
margin=go.layout.Margin(l=0, r=0, b=0, t=0),
|
253 |
+
showlegend=False,
|
254 |
+
legend=dict(
|
255 |
+
yanchor='bottom',
|
256 |
+
y=0.01,
|
257 |
+
xanchor='right',
|
258 |
+
x=0.99,
|
259 |
+
),
|
260 |
+
scene=dict(
|
261 |
+
aspectmode='manual',
|
262 |
+
aspectratio=dict(x=1, y=1, z=1.0),
|
263 |
+
camera=dict(
|
264 |
+
eye=dict(x=base_radius - 1.6, y=0.0, z=0.6),
|
265 |
+
center=dict(x=0.0, y=0.0, z=0.0),
|
266 |
+
up=dict(x=0.0, y=0.0, z=1.0)),
|
267 |
+
xaxis_title='',
|
268 |
+
yaxis_title='',
|
269 |
+
zaxis_title='',
|
270 |
+
xaxis=dict(
|
271 |
+
range=[-scene_bounds, scene_bounds],
|
272 |
+
showticklabels=False,
|
273 |
+
showgrid=True,
|
274 |
+
zeroline=False,
|
275 |
+
showbackground=True,
|
276 |
+
showspikes=False,
|
277 |
+
showline=False,
|
278 |
+
ticks=''),
|
279 |
+
yaxis=dict(
|
280 |
+
range=[-scene_bounds, scene_bounds],
|
281 |
+
showticklabels=False,
|
282 |
+
showgrid=True,
|
283 |
+
zeroline=False,
|
284 |
+
showbackground=True,
|
285 |
+
showspikes=False,
|
286 |
+
showline=False,
|
287 |
+
ticks=''),
|
288 |
+
zaxis=dict(
|
289 |
+
range=[-scene_bounds, scene_bounds],
|
290 |
+
showticklabels=False,
|
291 |
+
showgrid=True,
|
292 |
+
zeroline=False,
|
293 |
+
showbackground=True,
|
294 |
+
showspikes=False,
|
295 |
+
showline=False,
|
296 |
+
ticks='')))
|
297 |
+
|
298 |
+
self._fig = fig
|
299 |
+
return fig
|
300 |
+
|
301 |
+
|
302 |
+
def stage1_run(models, device, cam_vis, tmp_dir,
|
303 |
+
input_im, scale, ddim_steps, elev=None, rerun_all=[],
|
304 |
+
*btn_retrys):
|
305 |
+
is_rerun = True if cam_vis is None else False
|
306 |
+
model = models['turncam'].half()
|
307 |
+
|
308 |
+
stage1_dir = os.path.join(tmp_dir, "stage1_8")
|
309 |
+
if not is_rerun:
|
310 |
+
os.makedirs(stage1_dir, exist_ok=True)
|
311 |
+
output_ims = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(4)), device=device, ddim_steps=ddim_steps, scale=scale)
|
312 |
+
stage2_steps = 50 # ddim_steps
|
313 |
+
zero123_infer(model, tmp_dir, indices=[0], device=device, ddim_steps=stage2_steps, scale=scale)
|
314 |
+
try:
|
315 |
+
elev_output = estimate_elev(tmp_dir)
|
316 |
+
except:
|
317 |
+
print("Failed to estimate polar angle")
|
318 |
+
elev_output = 90
|
319 |
+
print("Estimated polar angle:", elev_output)
|
320 |
+
gen_poses(tmp_dir, elev_output)
|
321 |
+
show_in_im1 = np.asarray(input_im, dtype=np.uint8)
|
322 |
+
cam_vis.encode_image(show_in_im1, elev=elev_output)
|
323 |
+
new_fig = cam_vis.update_figure()
|
324 |
+
|
325 |
+
flag_lower_cam = elev_output <= 75
|
326 |
+
if flag_lower_cam:
|
327 |
+
output_ims_2 = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(4,8)), device=device, ddim_steps=ddim_steps, scale=scale)
|
328 |
+
else:
|
329 |
+
output_ims_2 = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(8,12)), device=device, ddim_steps=ddim_steps, scale=scale)
|
330 |
+
torch.cuda.empty_cache()
|
331 |
+
return (90-elev_output, new_fig, *output_ims, *output_ims_2)
|
332 |
+
else:
|
333 |
+
rerun_idx = [i for i in range(len(btn_retrys)) if btn_retrys[i]]
|
334 |
+
if 90-int(elev["label"]) > 75:
|
335 |
+
rerun_idx_in = [i if i < 4 else i+4 for i in rerun_idx]
|
336 |
+
else:
|
337 |
+
rerun_idx_in = rerun_idx
|
338 |
+
for idx in rerun_idx_in:
|
339 |
+
if idx not in rerun_all:
|
340 |
+
rerun_all.append(idx)
|
341 |
+
print("rerun_idx", rerun_all)
|
342 |
+
output_ims = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=rerun_idx_in, device=device, ddim_steps=ddim_steps, scale=scale)
|
343 |
+
outputs = [gr.update(visible=True)] * 8
|
344 |
+
for idx, view_idx in enumerate(rerun_idx):
|
345 |
+
outputs[view_idx] = output_ims[idx]
|
346 |
+
reset = [gr.update(value=False)] * 8
|
347 |
+
torch.cuda.empty_cache()
|
348 |
+
return (rerun_all, *reset, *outputs)
|
349 |
+
|
350 |
+
def stage2_run(models, device, tmp_dir,
|
351 |
+
elev, scale, is_glb=False, rerun_all=[], stage2_steps=50):
|
352 |
+
flag_lower_cam = 90-int(elev["label"]) <= 75
|
353 |
+
is_rerun = True if rerun_all else False
|
354 |
+
model = models['turncam'].half()
|
355 |
+
if not is_rerun:
|
356 |
+
if flag_lower_cam:
|
357 |
+
zero123_infer(model, tmp_dir, indices=list(range(1,8)), device=device, ddim_steps=stage2_steps, scale=scale)
|
358 |
+
else:
|
359 |
+
zero123_infer(model, tmp_dir, indices=list(range(1,4))+list(range(8,12)), device=device, ddim_steps=stage2_steps, scale=scale)
|
360 |
+
else:
|
361 |
+
print("rerun_idx", rerun_all)
|
362 |
+
zero123_infer(model, tmp_dir, indices=rerun_all, device=device, ddim_steps=stage2_steps, scale=scale)
|
363 |
+
|
364 |
+
dataset = tmp_dir
|
365 |
+
main_dir_path = os.path.dirname(__file__)
|
366 |
+
torch.cuda.empty_cache()
|
367 |
+
os.chdir(os.path.join(code_dir, 'SparseNeuS_demo_v1/'))
|
368 |
+
|
369 |
+
bash_script = f'CUDA_VISIBLE_DEVICES={_GPU_INDEX} python exp_runner_generic_blender_val.py --specific_dataset_name {dataset} --mode export_mesh --conf confs/one2345_lod0_val_demo.conf'
|
370 |
+
print(bash_script)
|
371 |
+
os.system(bash_script)
|
372 |
+
os.chdir(main_dir_path)
|
373 |
+
|
374 |
+
ply_path = os.path.join(tmp_dir, f"meshes_val_bg/lod0/mesh_00215000_gradio_lod0.ply")
|
375 |
+
mesh_ext = ".glb" if is_glb else ".obj"
|
376 |
+
mesh_path = os.path.join(tmp_dir, f"mesh{mesh_ext}")
|
377 |
+
# Read the textured mesh from .ply file
|
378 |
+
mesh = trimesh.load_mesh(ply_path)
|
379 |
+
rotation_matrix = trimesh.transformations.rotation_matrix(np.pi/2, [1, 0, 0])
|
380 |
+
mesh.apply_transform(rotation_matrix)
|
381 |
+
rotation_matrix = trimesh.transformations.rotation_matrix(np.pi, [0, 0, 1])
|
382 |
+
mesh.apply_transform(rotation_matrix)
|
383 |
+
# flip x
|
384 |
+
mesh.vertices[:, 0] = -mesh.vertices[:, 0]
|
385 |
+
mesh.faces = np.fliplr(mesh.faces)
|
386 |
+
# Export the mesh as .obj file with colors
|
387 |
+
if not is_glb:
|
388 |
+
mesh.export(mesh_path, file_type='obj', include_color=True)
|
389 |
+
else:
|
390 |
+
mesh.export(mesh_path, file_type='glb')
|
391 |
+
torch.cuda.empty_cache()
|
392 |
+
|
393 |
+
if not is_rerun:
|
394 |
+
return (mesh_path)
|
395 |
+
else:
|
396 |
+
return (mesh_path, gr.update(value=[]), gr.update(visible=False), gr.update(visible=False))
|
397 |
+
|
398 |
+
def nsfw_check(models, raw_im, device='cuda'):
|
399 |
+
safety_checker_input = models['clip_fe'](raw_im, return_tensors='pt').to(device)
|
400 |
+
(_, has_nsfw_concept) = models['nsfw'](
|
401 |
+
images=np.ones((1, 3)), clip_input=safety_checker_input.pixel_values)
|
402 |
+
del safety_checker_input
|
403 |
+
if np.any(has_nsfw_concept):
|
404 |
+
print('NSFW content detected.')
|
405 |
+
return Image.open("unsafe.png")
|
406 |
+
else:
|
407 |
+
print('Safety check passed.')
|
408 |
+
return False
|
409 |
+
|
410 |
+
def preprocess_run(predictor, models, raw_im, preprocess, *bbox_sliders):
|
411 |
+
raw_im.thumbnail([512, 512], Image.Resampling.LANCZOS)
|
412 |
+
check_results = nsfw_check(models, raw_im, device=predictor.device)
|
413 |
+
if check_results:
|
414 |
+
return check_results
|
415 |
+
image_sam = sam_out_nosave(predictor, raw_im.convert("RGB"), *bbox_sliders)
|
416 |
+
input_256 = image_preprocess_nosave(image_sam, lower_contrast=preprocess, rescale=True)
|
417 |
+
torch.cuda.empty_cache()
|
418 |
+
return input_256
|
419 |
+
|
420 |
+
def on_coords_slider(image, x_min, y_min, x_max, y_max, color=(88, 191, 131, 255)):
|
421 |
+
"""Draw a bounding box annotation for an image."""
|
422 |
+
print("Slider adjusted, drawing bbox...")
|
423 |
+
image.thumbnail([512, 512], Image.Resampling.LANCZOS)
|
424 |
+
image_size = image.size
|
425 |
+
if max(image_size) > 224:
|
426 |
+
image.thumbnail([224, 224], Image.Resampling.LANCZOS)
|
427 |
+
shrink_ratio = max(image.size) / max(image_size)
|
428 |
+
x_min = int(x_min * shrink_ratio)
|
429 |
+
y_min = int(y_min * shrink_ratio)
|
430 |
+
x_max = int(x_max * shrink_ratio)
|
431 |
+
y_max = int(y_max * shrink_ratio)
|
432 |
+
image = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2BGRA)
|
433 |
+
image = cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, int(max(max(image.shape) / 400*2, 2)))
|
434 |
+
return cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA) # image[:, :, ::-1]
|
435 |
+
|
436 |
+
def init_bbox(image):
|
437 |
+
image.thumbnail([512, 512], Image.Resampling.LANCZOS)
|
438 |
+
width, height = image.size
|
439 |
+
image_rem = image.convert('RGBA')
|
440 |
+
image_nobg = remove(image_rem, alpha_matting=True)
|
441 |
+
arr = np.asarray(image_nobg)[:,:,-1]
|
442 |
+
x_nonzero = np.nonzero(arr.sum(axis=0))
|
443 |
+
y_nonzero = np.nonzero(arr.sum(axis=1))
|
444 |
+
x_min = int(x_nonzero[0].min())
|
445 |
+
y_min = int(y_nonzero[0].min())
|
446 |
+
x_max = int(x_nonzero[0].max())
|
447 |
+
y_max = int(y_nonzero[0].max())
|
448 |
+
image_mini = image.copy()
|
449 |
+
image_mini.thumbnail([224, 224], Image.Resampling.LANCZOS)
|
450 |
+
shrink_ratio = max(image_mini.size) / max(width, height)
|
451 |
+
x_min_shrink = int(x_min * shrink_ratio)
|
452 |
+
y_min_shrink = int(y_min * shrink_ratio)
|
453 |
+
x_max_shrink = int(x_max * shrink_ratio)
|
454 |
+
y_max_shrink = int(y_max * shrink_ratio)
|
455 |
+
|
456 |
+
return [on_coords_slider(image_mini, x_min_shrink, y_min_shrink, x_max_shrink, y_max_shrink),
|
457 |
+
gr.update(value=x_min, maximum=width),
|
458 |
+
gr.update(value=y_min, maximum=height),
|
459 |
+
gr.update(value=x_max, maximum=width),
|
460 |
+
gr.update(value=y_max, maximum=height)]
|
461 |
+
|
462 |
+
|
463 |
+
def run_demo(
|
464 |
+
device_idx=_GPU_INDEX,
|
465 |
+
ckpt='zero123-xl.ckpt'):
|
466 |
+
|
467 |
+
device = f"cuda:{device_idx}" if torch.cuda.is_available() else "cpu"
|
468 |
+
models = init_model(device, os.path.join(code_dir, ckpt))
|
469 |
+
|
470 |
+
# init sam model
|
471 |
+
predictor = sam_init(device_idx)
|
472 |
+
|
473 |
+
with open('instructions_12345.md', 'r') as f:
|
474 |
+
article = f.read()
|
475 |
+
|
476 |
+
# NOTE: Examples must match inputs
|
477 |
+
example_folder = os.path.join(os.path.dirname(__file__), 'demo_examples')
|
478 |
+
example_fns = os.listdir(example_folder)
|
479 |
+
example_fns.sort()
|
480 |
+
examples_full = [os.path.join(example_folder, x) for x in example_fns if x.endswith('.png')]
|
481 |
+
|
482 |
+
# Compose demo layout & data flow.
|
483 |
+
with gr.Blocks(title=_TITLE, css="style.css") as demo:
|
484 |
+
with gr.Row():
|
485 |
+
with gr.Column(scale=1):
|
486 |
+
gr.Markdown('# ' + _TITLE)
|
487 |
+
with gr.Column(scale=0):
|
488 |
+
gr.DuplicateButton(value='Duplicate Space for private use',
|
489 |
+
elem_id='duplicate-button')
|
490 |
+
gr.Markdown(_DESCRIPTION)
|
491 |
+
|
492 |
+
with gr.Row(variant='panel'):
|
493 |
+
with gr.Column(scale=6):
|
494 |
+
image_block = gr.Image(type='pil', image_mode='RGBA', height=290, label='Input image', tool=None)
|
495 |
+
|
496 |
+
gr.Examples(
|
497 |
+
examples=examples_full, # NOTE: elements must match inputs list!
|
498 |
+
inputs=[image_block],
|
499 |
+
outputs=[image_block],
|
500 |
+
cache_examples=False,
|
501 |
+
label='Examples (click one of the images below to start)',
|
502 |
+
examples_per_page=40
|
503 |
+
)
|
504 |
+
preprocess_chk = gr.Checkbox(
|
505 |
+
False, label='Reduce image contrast (mitigate shadows on the backside)')
|
506 |
+
with gr.Accordion('Advanced options', open=False):
|
507 |
+
scale_slider = gr.Slider(0, 30, value=3, step=1,
|
508 |
+
label='Diffusion guidance scale')
|
509 |
+
steps_slider = gr.Slider(5, 200, value=75, step=5,
|
510 |
+
label='Number of diffusion inference steps')
|
511 |
+
glb_chk = gr.Checkbox(
|
512 |
+
False, label='Export the mesh in .glb format')
|
513 |
+
|
514 |
+
run_btn = gr.Button('Run Generation', variant='primary', interactive=False)
|
515 |
+
guide_text = gr.Markdown(_USER_GUIDE, visible=True)
|
516 |
+
|
517 |
+
with gr.Column(scale=4):
|
518 |
+
with gr.Row():
|
519 |
+
bbox_block = gr.Image(type='pil', label="Bounding box", height=290, interactive=False)
|
520 |
+
sam_block = gr.Image(type='pil', label="SAM output", interactive=False)
|
521 |
+
max_width = max_height = 256
|
522 |
+
with gr.Row():
|
523 |
+
x_min_slider = gr.Slider(label="X min", interactive=True, value=0, minimum=0, maximum=max_width, step=1)
|
524 |
+
y_min_slider = gr.Slider(label="Y min", interactive=True, value=0, minimum=0, maximum=max_height, step=1)
|
525 |
+
with gr.Row():
|
526 |
+
x_max_slider = gr.Slider(label="X max", interactive=True, value=max_width, minimum=0, maximum=max_width, step=1)
|
527 |
+
y_max_slider = gr.Slider(label="Y max", interactive=True, value=max_height, minimum=0, maximum=max_height, step=1)
|
528 |
+
bbox_sliders = [x_min_slider, y_min_slider, x_max_slider, y_max_slider]
|
529 |
+
|
530 |
+
mesh_output = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="One-2-3-45's Textured Mesh", elem_id="model-3d-out")
|
531 |
+
|
532 |
+
with gr.Row(variant='panel'):
|
533 |
+
with gr.Column(scale=85):
|
534 |
+
elev_output = gr.Label(label='Estimated elevation (degree, w.r.t. the horizontal plane)')
|
535 |
+
vis_output = gr.Plot(label='Camera poses of the input view (red) and predicted views (blue)', elem_id="plot-out")
|
536 |
+
|
537 |
+
with gr.Column(scale=115):
|
538 |
+
gr.Markdown('Predicted multi-view images')
|
539 |
+
with gr.Row():
|
540 |
+
view_1 = gr.Image(interactive=False, height=200, show_label=False)
|
541 |
+
view_2 = gr.Image(interactive=False, height=200, show_label=False)
|
542 |
+
view_3 = gr.Image(interactive=False, height=200, show_label=False)
|
543 |
+
view_4 = gr.Image(interactive=False, height=200, show_label=False)
|
544 |
+
with gr.Row():
|
545 |
+
btn_retry_1 = gr.Checkbox(label='Retry view 1')
|
546 |
+
btn_retry_2 = gr.Checkbox(label='Retry view 2')
|
547 |
+
btn_retry_3 = gr.Checkbox(label='Retry view 3')
|
548 |
+
btn_retry_4 = gr.Checkbox(label='Retry view 4')
|
549 |
+
with gr.Row():
|
550 |
+
view_5 = gr.Image(interactive=False, height=200, show_label=False)
|
551 |
+
view_6 = gr.Image(interactive=False, height=200, show_label=False)
|
552 |
+
view_7 = gr.Image(interactive=False, height=200, show_label=False)
|
553 |
+
view_8 = gr.Image(interactive=False, height=200, show_label=False)
|
554 |
+
with gr.Row():
|
555 |
+
btn_retry_5 = gr.Checkbox(label='Retry view 5')
|
556 |
+
btn_retry_6 = gr.Checkbox(label='Retry view 6')
|
557 |
+
btn_retry_7 = gr.Checkbox(label='Retry view 7')
|
558 |
+
btn_retry_8 = gr.Checkbox(label='Retry view 8')
|
559 |
+
with gr.Row():
|
560 |
+
regen_view_btn = gr.Button('1. Regenerate selected view(s)', variant='secondary', visible=False)
|
561 |
+
regen_mesh_btn = gr.Button('2. Regenerate nearby views and mesh', variant='secondary', visible=False)
|
562 |
+
|
563 |
+
gr.Markdown(article)
|
564 |
+
gr.HTML("""
|
565 |
+
<div class="footer">
|
566 |
+
<p>
|
567 |
+
One-2-3-45 Demo by <a style="text-decoration:none" href="https://chaoxu.xyz" target="_blank">Chao Xu</a>
|
568 |
+
</p>
|
569 |
+
</div>
|
570 |
+
""")
|
571 |
+
|
572 |
+
update_guide = lambda GUIDE_TEXT: gr.update(value=GUIDE_TEXT)
|
573 |
+
|
574 |
+
views = [view_1, view_2, view_3, view_4, view_5, view_6, view_7, view_8]
|
575 |
+
btn_retrys = [btn_retry_1, btn_retry_2, btn_retry_3, btn_retry_4, btn_retry_5, btn_retry_6, btn_retry_7, btn_retry_8]
|
576 |
+
|
577 |
+
rerun_idx = gr.State([])
|
578 |
+
tmp_dir = gr.State('./demo_tmp/tmp_dir')
|
579 |
+
|
580 |
+
def refresh(tmp_dir):
|
581 |
+
if os.path.exists(tmp_dir):
|
582 |
+
shutil.rmtree(tmp_dir)
|
583 |
+
tmp_dir = tempfile.TemporaryDirectory(dir=os.path.join(os.path.dirname(__file__), 'demo_tmp'))
|
584 |
+
print("create tmp_dir", tmp_dir.name)
|
585 |
+
clear = [gr.update(value=[])] + [None] * 5 + [gr.update(visible=False)] * 2 + [None] * 8 + [gr.update(value=False)] * 8
|
586 |
+
return (tmp_dir.name, *clear)
|
587 |
+
|
588 |
+
placeholder = gr.Image(visible=False)
|
589 |
+
tmp_func = lambda x: False if not x else gr.update(visible=False)
|
590 |
+
disable_func = lambda x: gr.update(interactive=False)
|
591 |
+
enable_func = lambda x: gr.update(interactive=True)
|
592 |
+
image_block.change(disable_func, inputs=run_btn, outputs=run_btn, queue=False
|
593 |
+
).success(fn=refresh,
|
594 |
+
inputs=[tmp_dir],
|
595 |
+
outputs=[tmp_dir, rerun_idx, bbox_block, sam_block, elev_output, vis_output, mesh_output, regen_view_btn, regen_mesh_btn, *views, *btn_retrys],
|
596 |
+
queue=False
|
597 |
+
).success(fn=tmp_func, inputs=[image_block], outputs=[placeholder], queue=False
|
598 |
+
).success(fn=partial(update_guide, _BBOX_1), outputs=[guide_text], queue=False
|
599 |
+
).success(fn=init_bbox,
|
600 |
+
inputs=[image_block],
|
601 |
+
outputs=[bbox_block, *bbox_sliders], queue=False
|
602 |
+
).success(fn=partial(update_guide, _BBOX_3), outputs=[guide_text], queue=False
|
603 |
+
).success(enable_func, inputs=run_btn, outputs=run_btn, queue=False)
|
604 |
+
|
605 |
+
|
606 |
+
for bbox_slider in bbox_sliders:
|
607 |
+
bbox_slider.release(fn=on_coords_slider,
|
608 |
+
inputs=[image_block, *bbox_sliders],
|
609 |
+
outputs=[bbox_block],
|
610 |
+
queue=False
|
611 |
+
).success(fn=partial(update_guide, _BBOX_2), outputs=[guide_text], queue=False)
|
612 |
+
|
613 |
+
cam_vis = CameraVisualizer(vis_output)
|
614 |
+
|
615 |
+
# Define the function to be called when any of the btn_retry buttons are clicked
|
616 |
+
def on_retry_button_click(*btn_retrys):
|
617 |
+
any_checked = any([btn_retry for btn_retry in btn_retrys])
|
618 |
+
print('any_checked:', any_checked, [btn_retry for btn_retry in btn_retrys])
|
619 |
+
if any_checked:
|
620 |
+
return (gr.update(visible=True), gr.update(visible=True))
|
621 |
+
else:
|
622 |
+
return (gr.update(), gr.update())
|
623 |
+
# make regen_btn visible when any of the btn_retry is checked
|
624 |
+
for btn_retry in btn_retrys:
|
625 |
+
# Add the event handlers to the btn_retry buttons
|
626 |
+
btn_retry.change(fn=on_retry_button_click, inputs=[*btn_retrys], outputs=[regen_view_btn, regen_mesh_btn], queue=False)
|
627 |
+
|
628 |
+
|
629 |
+
run_btn.click(fn=partial(update_guide, _SAM), outputs=[guide_text], queue=False
|
630 |
+
).success(fn=partial(preprocess_run, predictor, models),
|
631 |
+
inputs=[image_block, preprocess_chk, *bbox_sliders],
|
632 |
+
outputs=[sam_block]
|
633 |
+
).success(fn=partial(update_guide, _GEN_1), outputs=[guide_text], queue=False
|
634 |
+
).success(fn=partial(stage1_run, models, device, cam_vis),
|
635 |
+
inputs=[tmp_dir, sam_block, scale_slider, steps_slider],
|
636 |
+
outputs=[elev_output, vis_output, *views]
|
637 |
+
).success(fn=partial(update_guide, _GEN_2), outputs=[guide_text], queue=False
|
638 |
+
).success(fn=partial(stage2_run, models, device),
|
639 |
+
inputs=[tmp_dir, elev_output, scale_slider, glb_chk],
|
640 |
+
outputs=[mesh_output]
|
641 |
+
).success(fn=partial(update_guide, _DONE), outputs=[guide_text], queue=False)
|
642 |
+
|
643 |
+
|
644 |
+
regen_view_btn.click(fn=partial(stage1_run, models, device, None),
|
645 |
+
inputs=[tmp_dir, sam_block, scale_slider, steps_slider, elev_output, rerun_idx, *btn_retrys],
|
646 |
+
outputs=[rerun_idx, *btn_retrys, *views]
|
647 |
+
).success(fn=partial(update_guide, _REGEN_1), outputs=[guide_text], queue=False)
|
648 |
+
regen_mesh_btn.click(fn=partial(stage2_run, models, device),
|
649 |
+
inputs=[tmp_dir, elev_output, scale_slider, glb_chk, rerun_idx],
|
650 |
+
outputs=[mesh_output, rerun_idx, regen_view_btn, regen_mesh_btn]
|
651 |
+
).success(fn=partial(update_guide, _REGEN_2), outputs=[guide_text], queue=False)
|
652 |
+
|
653 |
+
|
654 |
+
demo.queue().launch(share=False, max_threads=80) # auth=("admin", os.environ['PASSWD'])
|
655 |
+
|
656 |
+
|
657 |
+
if __name__ == '__main__':
|
658 |
+
fire.Fire(run_demo)
|