File size: 10,187 Bytes
5a03db1
 
11e8a80
5a03db1
 
 
34d287c
 
 
5a03db1
34d287c
 
dc32bb2
5a03db1
34d287c
 
 
d814d5a
34d287c
 
 
c2c23af
11e8a80
d814d5a
11e8a80
d814d5a
34d287c
5a03db1
d814d5a
34d287c
d814d5a
5a03db1
7d141d7
559ce8f
d814d5a
 
 
 
 
34d287c
 
 
dc32bb2
 
34d287c
 
dc32bb2
34d287c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d814d5a
34d287c
 
 
 
 
 
 
 
 
d814d5a
 
 
34d287c
 
 
 
d814d5a
34d287c
 
 
 
 
 
 
 
 
 
d814d5a
 
 
 
34d287c
d814d5a
34d287c
d814d5a
 
 
ef862e7
 
 
34d287c
 
 
 
5a03db1
 
 
34d287c
 
 
 
 
 
5a03db1
34d287c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5a03db1
34d287c
 
 
 
 
5a03db1
34d287c
 
5a03db1
34d287c
 
 
 
 
 
 
 
 
5d6ede9
 
 
 
 
 
 
d814d5a
 
 
 
 
 
 
 
 
 
34d287c
3b22732
 
d814d5a
34d287c
 
 
d814d5a
3b22732
d814d5a
 
34d287c
5d6ede9
1bfa5fd
 
5d6ede9
 
 
 
 
757e693
5d6ede9
6eae011
5d6ede9
 
 
 
ef862e7
f6a127b
757e693
5d6ede9
 
 
 
 
 
 
 
1bfa5fd
757e693
 
 
5d6ede9
 
 
 
 
f6a127b
1bfa5fd
 
 
f6a127b
 
 
5d6ede9
 
3b22732
c5a094f
5d6ede9
 
 
 
 
 
1bfa5fd
5d6ede9
 
 
 
 
18581d3
 
34d287c
 
757e693
1bfa5fd
757e693
5d6ede9
 
 
 
 
 
34d287c
 
 
11e8a80
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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
import os
import cv2
import torch
import numpy as np
import gradio as gr

import trimesh
import sys
import os

sys.path.append('vggsfm_code/')
import shutil
from datetime import datetime

from vggsfm_code.hf_demo import demo_fn
from omegaconf import DictConfig, OmegaConf
from viz_utils.viz_fn import add_camera
import glob
# 
from scipy.spatial.transform import Rotation
import PIL
import gc

# import spaces

# @spaces.GPU
def vggsfm_demo(
    input_video,
    input_image,
    query_frame_num,
    max_query_pts=4096,
):
    gc.collect()
    torch.cuda.empty_cache()

    if input_video is not None:            
        if not isinstance(input_video, str):
            input_video = input_video["video"]["path"]
    
    cfg_file = "vggsfm_code/cfgs/demo.yaml"
    cfg = OmegaConf.load(cfg_file)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    max_input_image = 20

    target_dir = f"input_images_{timestamp}"
    if os.path.exists(target_dir): 
        shutil.rmtree(target_dir)

    os.makedirs(target_dir)
    target_dir_images = target_dir + "/images"
    os.makedirs(target_dir_images)

    if input_image is not None:
        if len(input_image)<3:
            return None, "Please input at least three frames"

        input_image = sorted(input_image)
        input_image = input_image[:max_input_image]
        
        # Copy files to the new directory
        for file_name in input_image:
            shutil.copy(file_name, target_dir_images)
    elif input_video is not None:
        vs = cv2.VideoCapture(input_video)

        fps = vs.get(cv2.CAP_PROP_FPS)


        frame_rate = 1
        frame_interval = int(fps * frame_rate)
        
        video_frame_num = 0
        count = 0 
        
        while video_frame_num<=max_input_image:
            (gotit, frame) = vs.read()
            count +=1

            if not gotit:
                break
            
            if count % frame_interval == 0:
                cv2.imwrite(target_dir_images+"/"+f"{video_frame_num:06}.png", frame)
                video_frame_num+=1
                
        if video_frame_num<3:
            return None, "Please input at least three frames"
    else:
        return None, "Input format incorrect"
        
    cfg.query_frame_num = query_frame_num
    cfg.max_query_pts = max_query_pts
    print(f"Files have been copied to {target_dir_images}")
    cfg.SCENE_DIR = target_dir
    
    # try:
    predictions = demo_fn(cfg)
    # except:
    # return None, "Something seems to be incorrect. Please verify that your inputs are formatted correctly. If the issue persists, kindly create a GitHub issue for further assistance."
    
    glbscene = vggsfm_predictions_to_glb(predictions)
    
    glbfile = target_dir + "/glbscene.glb"
    glbscene.export(file_obj=glbfile)    

    del predictions
    gc.collect()
    torch.cuda.empty_cache()
    
    print(input_image)
    print(input_video)
    return glbfile, "Success"




def vggsfm_predictions_to_glb(predictions):
    # learned from https://github.com/naver/dust3r/blob/main/dust3r/viz.py
    points3D = predictions["points3D"].cpu().numpy()
    points3D_rgb = predictions["points3D_rgb"].cpu().numpy()
    points3D_rgb = (points3D_rgb*255).astype(np.uint8)
    
    extrinsics_opencv = predictions["extrinsics_opencv"].cpu().numpy()
    intrinsics_opencv = predictions["intrinsics_opencv"].cpu().numpy()
    raw_image_paths = predictions["raw_image_paths"]
    images = predictions["images"].permute(0,2,3,1).cpu().numpy()
    images = (images*255).astype(np.uint8)
    
    glbscene = trimesh.Scene()
    point_cloud = trimesh.PointCloud(points3D, colors=points3D_rgb)
    glbscene.add_geometry(point_cloud)


    camera_edge_colors = [(255, 0, 0), (0, 0, 255), (0, 255, 0), (255, 0, 255), (255, 204, 0), (0, 204, 204),
                (128, 255, 255), (255, 128, 255), (255, 255, 128), (0, 0, 0), (128, 128, 128)]

    frame_num = len(extrinsics_opencv)
    extrinsics_opencv_4x4 = np.zeros((frame_num, 4, 4))
    extrinsics_opencv_4x4[:, :3, :4] = extrinsics_opencv
    extrinsics_opencv_4x4[:, 3, 3] = 1

    for idx in range(frame_num):
        cam_from_world = extrinsics_opencv_4x4[idx]
        cam_to_world = np.linalg.inv(cam_from_world)
        cur_cam_color = camera_edge_colors[idx % len(camera_edge_colors)]
        cur_focal = intrinsics_opencv[idx, 0, 0]

        add_camera(glbscene, cam_to_world, cur_cam_color, image=None, imsize=(1024,1024), 
                   focal=None,screen_width=0.35)

    opengl_mat = np.array([[1, 0, 0, 0],
                    [0, -1, 0, 0],
                    [0, 0, -1, 0],
                    [0, 0, 0, 1]])

    rot = np.eye(4)
    rot[:3, :3] = Rotation.from_euler('y', np.deg2rad(180)).as_matrix()
    glbscene.apply_transform(np.linalg.inv(np.linalg.inv(extrinsics_opencv_4x4[0]) @ opengl_mat @ rot))

    # Calculate the bounding box center and apply the translation
    bounding_box = glbscene.bounds
    center = (bounding_box[0] + bounding_box[1]) / 2
    translation = np.eye(4)
    translation[:3, 3] = -center

    glbscene.apply_transform(translation)
    # glbfile = "glbscene.glb"
    # glbscene.export(file_obj=glbfile)    
    return glbscene

# apple_video = "vggsfm_code/examples/videos/apple_video.mp4"
# os.path.join(os.path.dirname(__file__), "apple_video.mp4")
british_museum_video = "vggsfm_code/examples/videos/british_museum_video.mp4"

# os.path.join(os.path.dirname(__file__), "british_museum_video.mp4")
cake_video = "vggsfm_code/examples/videos/cake_video.mp4"

bonsai_video = "vggsfm_code/examples/videos/bonsai_video.mp4"

# os.path.join(os.path.dirname(__file__), "cake_video.mp4")



# apple_images = glob.glob(f'vggsfm_code/examples/apple/images/*')
bonsai_images = glob.glob(f'vggsfm_code/examples/bonsai/images/*')
cake_images = glob.glob(f'vggsfm_code/examples/cake/images/*')
british_museum_images = glob.glob(f'vggsfm_code/examples/british_museum/images/*')




with gr.Blocks() as demo:
    gr.Markdown("# 🎨 VGGSfM: Visual Geometry Grounded Deep Structure From Motion")
    
    gr.Markdown("""
    <div style="text-align: left;"> 
    <p>Welcome to <a href="https://vggsfm.github.io/" target="_blank">VGGSfM</a> demo! 
    This space demonstrates 3D reconstruction from input image frames. </p> 
    <p>To get started quickly, you can click on our <strong> examples (the bottom of the page) </strong>. If you want to reconstruct your own data, simply: </p> 
    <ul style="display: inline-block; text-align: left;"> 
        <li>upload the images (.jpg, .png, etc.), or </li> 
        <li>upload a video (.mp4, .mov, etc.) </li> 
    </ul> 
    <p>The reconstruction should normally take <strong> up to 90 second </strong>. If both images and videos are uploaded, the demo will only reconstruct the uploaded images. By default, we extract <strong> 1 image frame per second from the input video </strong>. To prevent crashes on the Hugging Face space, we currently limit reconstruction to the first 20 image frames. </p> 
    <p>SfM methods are designed for <strong> rigid/static reconstruction </strong>. When dealing with dynamic/moving inputs, these methods may still work by focusing on the rigid parts of the scene. However, to ensure high-quality results, it is better to minimize the presence of moving objects in the input data. </p> 
    <p>If you meet any problem, feel free to create an issue in our <a href="https://github.com/facebookresearch/vggsfm" target="_blank">GitHub Repo</a> ⭐</p> 
    <p>(Please note that running reconstruction on Hugging Face space is slower than on a local machine.) </p> 
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            input_video = gr.Video(label="Input video", interactive=True)
            input_images = gr.File(file_count="multiple", label="Input Images", interactive=True)
            num_query_images = gr.Slider(minimum=1, maximum=10, step=1, value=5, label="Number of query images (key frames)",
                                         info="More query images usually lead to better reconstruction at lower speeds. If the viewpoint differences between your images are minimal, you can set this value to 1. ")
            num_query_points = gr.Slider(minimum=512, maximum=4096, step=1, value=1024, label="Number of query points",
                                         info="More query points usually lead to denser reconstruction at lower speeds.")
        
        with gr.Column(scale=3):
            reconstruction_output = gr.Model3D(label="Reconstruction", height=520)
            log_output = gr.Textbox(label="Log")

    with gr.Row():
        submit_btn = gr.Button("Reconstruct", scale=1)

        # submit_btn = gr.Button("Reconstruct", scale=1, elem_attributes={"style": "background-color: blue; color: white;"})
        clear_btn = gr.ClearButton([input_video, input_images, num_query_images, num_query_points, reconstruction_output, log_output], scale=1)
    
    
    examples = [
        [british_museum_video, british_museum_images, 2, 4096],
        [bonsai_video, bonsai_images, 3, 2048],
        [cake_video, cake_images, 3, 2048],
    ]
    
    gr.Examples(examples=examples, 
                inputs=[input_video, input_images, num_query_images, num_query_points],
                outputs=[reconstruction_output, log_output],  # Provide outputs
                fn=vggsfm_demo,  # Provide the function
                cache_examples=True,
                )

    submit_btn.click(
        vggsfm_demo,
        [input_video, input_images, num_query_images, num_query_points],
        [reconstruction_output, log_output],
        concurrency_limit=1
    )

    # demo.launch(debug=True, share=True)
    demo.queue(max_size=20).launch(show_error=True, share=True)
    # demo.queue(max_size=20, concurrency_count=1).launch(debug=True, share=True)
########################################################################################################################

# else:
#     import glob
#     files = glob.glob(f'vggsfm_code/examples/cake/images/*', recursive=True)
#     vggsfm_demo(files, None, None)

    
# demo.queue(max_size=20, concurrency_count=1).launch(debug=True, share=True)