File size: 6,318 Bytes
57746f1 33ab518 57746f1 33ab518 57746f1 33ab518 57746f1 |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import os, subprocess, shlex, sys, gc
import time
import torch
import numpy as np
import shutil
import argparse
import gradio as gr
import uuid
import spaces
#
subprocess.run(shlex.split("pip install wheel/torch_scatter-2.1.2+pt21cu121-cp310-cp310-linux_x86_64.whl"))
subprocess.run(shlex.split("pip install wheel/flash_attn-2.6.3+cu123torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"))
subprocess.run(shlex.split("pip install wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl"))
subprocess.run(shlex.split("pip install wheel/simple_knn-0.0.0-cp310-cp310-linux_x86_64.whl"))
subprocess.run(shlex.split("pip install wheel/curope-0.0.0-cp310-cp310-linux_x86_64.whl"))
subprocess.run(shlex.split("pip install wheel/pointops-1.0-cp310-cp310-linux_x86_64.whl"))
from src.utils.visualization_utils import render_video_from_file
from src.model import LSM_MASt3R
model = LSM_MASt3R.from_pretrained("checkpoints/pretrained_model/checkpoint-40.pth")
model = model.eval()
@spaces.GPU(duration=80)
def process(inputfiles, input_path=None):
# 创建唯一的缓存目录
cache_dir = os.path.join('outputs', str(uuid.uuid4()))
os.makedirs(cache_dir, exist_ok=True)
if input_path is not None:
imgs_path = './assets/examples/' + input_path
imgs_names = sorted(os.listdir(imgs_path))
inputfiles = []
for imgs_name in imgs_names:
file_path = os.path.join(imgs_path, imgs_name)
print(file_path)
inputfiles.append(file_path)
print(inputfiles)
filelist = inputfiles
if len(filelist) != 2:
gr.Warning("Please select 2 images")
shutil.rmtree(cache_dir) # 清理缓存目录
return None, None, None, None, None, None
ply_path = os.path.join(cache_dir, 'gaussians.ply')
# render_video_from_file(filelist, model, output_path=cache_dir, resolution=224)
render_video_from_file(filelist, model, output_path=cache_dir, resolution=512)
rgb_video_path = os.path.join(cache_dir, 'moved', 'output_images_video.mp4')
depth_video_path = os.path.join(cache_dir, 'moved', 'output_depth_video.mp4')
feature_video_path = os.path.join(cache_dir, 'moved', 'output_fmap_video.mp4')
return filelist, rgb_video_path, depth_video_path, feature_video_path, ply_path, ply_path
_TITLE = 'LargeSpatialModel'
_DESCRIPTION = '''
<div style="display: flex; justify-content: center; align-items: center;">
<div style="width: 100%; text-align: center; font-size: 30px;">
<strong>Large Spatial Model: End-to-end Unposed Images to Semantic 3D</strong>
</div>
</div>
<p></p>
<div align="center">
<a style="display:inline-block" href="https://arxiv.org/abs/2410.18956"><img src="https://img.shields.io/badge/ArXiv-2410.18956-b31b1b?logo=arxiv" alt='arxiv'></a>
<a style="display:inline-block" href="https://largespatialmodel.github.io/"><img src='https://img.shields.io/badge/Project_Page-ff7512?logo=lightning'></a>
<a title="Social" href="https://x.com/WayneINR" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
<img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
</a>
</div>
<p></p>
* Official demo of: [LargeSpatialModel: End-to-end Unposed Images to Semantic 3D](https://largespatialmodel.github.io/).
* Examples for direct viewing: you can simply click the examples (in the bottom of the page), to quickly view the results on representative data.
'''
block = gr.Blocks().queue()
with block:
gr.Markdown(_DESCRIPTION)
with gr.Column(variant="panel"):
with gr.Tab("Input"):
with gr.Row():
with gr.Column(scale=1):
inputfiles = gr.File(file_count="multiple", label="Load Images")
input_path = gr.Textbox(visible=False, label="example_path")
with gr.Column(scale=1):
image_gallery = gr.Gallery(
label="Gallery",
show_label=False,
elem_id="gallery",
columns=[2],
height=300, # 固定高度
object_fit="cover" # 确保图片填满空间
)
button_gen = gr.Button("Start Reconstruction", elem_id="button_gen")
processing_msg = gr.Markdown("Processing...", visible=False, elem_id="processing_msg")
with gr.Column(variant="panel"):
with gr.Tab("Output"):
with gr.Row():
with gr.Column(scale=1):
rgb_video = gr.Video(label="RGB Video", autoplay=True)
with gr.Column(scale=1):
feature_video = gr.Video(label="Feature Video", autoplay=True)
with gr.Column(scale=1):
depth_video = gr.Video(label="Depth Video", autoplay=True)
with gr.Row():
with gr.Group():
output_model = gr.Model3D(
label="3D Dense Model under Gaussian Splats Formats, need more time to visualize",
interactive=False,
camera_position=[0.5, 0.5, 1], # 稍微偏移一点，以便更好地查看模型
height=600,
)
gr.Markdown(
"""
<div class="model-description">
Use the left mouse button to rotate, the scroll wheel to zoom, and the right mouse button to move.
</div>
"""
)
with gr.Row():
output_file = gr.File(label="PLY File")
examples = gr.Examples(
examples=[
"sofa",
],
inputs=[input_path],
outputs=[image_gallery, rgb_video, depth_video, feature_video, output_model, output_file],
fn=lambda x: process(inputfiles=None, input_path=x),
cache_examples=True,
label="Examples"
)
button_gen.click(
process,
inputs=[inputfiles],
outputs=[image_gallery, rgb_video, depth_video, feature_video, output_model, output_file],
)
block.launch(server_name="0.0.0.0", share=False)
|