File size: 8,226 Bytes
0f8b4e3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a19de17
0f8b4e3
a19de17
0f8b4e3
 
 
 
 
 
 
 
3afeb70
ae2071d
3afeb70
 
 
 
ae2071d
 
 
3afeb70
bffa280
ae2071d
bffa280
3afeb70
 
 
ae2071d
8287c2f
3afeb70
ae2071d
3afeb70
ae2071d
 
0f8b4e3
 
 
 
 
 
 
 
 
 
 
 
 
 
3afeb70
 
 
 
 
 
0f8b4e3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8287c2f
0b13fd6
 
 
0f8b4e3
 
 
 
 
 
 
 
 
 
 
 
 
 
3afeb70
 
8287c2f
3afeb70
0f8b4e3
 
 
 
 
 
 
 
 
3afeb70
 
 
 
 
0f8b4e3
 
 
 
 
 
 
 
 
 
 
 
 
 
3afeb70
0f8b4e3
 
 
 
0b13fd6
 
 
3afeb70
0b13fd6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0f8b4e3
 
 
3afeb70
0f8b4e3
 
 
 
 
 
 
 
a19de17
0f8b4e3
 
 
 
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
import sys
sys.path.append("flash3d")

from omegaconf import OmegaConf
import gradio as gr
import spaces
import torch
import torchvision.transforms as TT
import torchvision.transforms.functional as TTF
from huggingface_hub import hf_hub_download

from networks.gaussian_predictor import GaussianPredictor
from util.vis3d import save_ply


def main():
    if torch.cuda.is_available():
        device = "cuda:0"
    else:
        device = "cpu"

    model_cfg_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
                                     filename="config_re10k_v1.yaml")
    model_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
                                 filename="model_re10k_v1.pth")

    cfg = OmegaConf.load(model_cfg_path)
    model = GaussianPredictor(cfg)
    device = torch.device(device)
    model.load_model(model_path)
    model.to(device)

    pad_border_fn = TT.Pad((cfg.dataset.pad_border_aug, cfg.dataset.pad_border_aug))
    to_tensor = TT.ToTensor()

    def check_input_image(input_image):
        if input_image is None:
            raise gr.Error("No image uploaded!")

    def preprocess(image, dynamic_size=False, padding=True):
        h, w = image.size
        size = 32
        if dynamic_size:
            while max(h, w) // size > 20:
                size *=2
            crop_image = TTF.center_crop(image, (w // size * size, h // size * size))
            resize_image = TTF.resize(crop_image, (w // size * 32, h // size * 32), interpolation=TT.InterpolationMode.BICUBIC)
            model.cfg.dataset.width, model.cfg.dataset.height = resize_image.size
        else:
            model.cfg.dataset.height, model.cfg.dataset.width = 256, 384
            resize_image = TTF.resize(
                image, (model.cfg.dataset.height, model.cfg.dataset.width), 
                interpolation=TT.InterpolationMode.BICUBIC
            )
        if padding:
            input_image = pad_border_fn(resize_image)
            model.cfg.dataset.pad_border_aug = 32
        else:
            input_image = resize_image
            model.cfg.dataset.pad_border_aug = 0
        model.set_backproject()
        return input_image

    @spaces.GPU()
    def reconstruct_and_export(image):
        """
        Passes image through model, outputs reconstruction in form of a dict of tensors.
        """
        image = to_tensor(image).to(device).unsqueeze(0)
        inputs = {
            ("color_aug", 0, 0): image,
        }

        outputs = model(inputs)

        # export reconstruction to ply
        save_ply(outputs, 
                 ply_out_path,
                 num_gauss=model.cfg.model.gaussians_per_pixel,
                 h=model.cfg.dataset.height,
                 w=model.cfg.dataset.width,
                 pad=model.cfg.dataset.pad_border_aug)

        return ply_out_path
    
    ply_out_path = f'./mesh.ply'

    css = """
        h1 {
            text-align: center;
            display:block;
        }
        """

    with gr.Blocks(css=css) as demo:
        gr.Markdown(
            """
            # Flash3D
            **Flash3D** [[project page](https://www.robots.ox.ac.uk/~vgg/research/flash3d/)] is a fast, super efficient, trainable on a single GPU in a day for scene 3D reconstruction from a single image.
            The model used in the demo was trained on only **RealEstate10k dataset on a single A6000 GPU within 1 day**.
            Upload an image of a scene or click on one of the provided examples to see how the Flash3D does.
            The 3D viewer will render a .ply scene exported from the 3D Gaussians, which is only an approximation.
            """
            )
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                with gr.Row():
                    input_image = gr.Image(
                        label="Input Image",
                        image_mode="RGBA",
                        sources="upload",
                        type="pil",
                        elem_id="content_image",
                    )
                with gr.Row():
                    submit = gr.Button("Generate", elem_id="generate", variant="primary")
                
                with gr.Row():
                        dynamic_size = gr.Checkbox(True, interactive=True, label='Use the original image ratio')
                        padding = gr.Checkbox(True, interactive=True, label='add padding to the image')

                with gr.Row(variant="panel"): 
                    gr.Examples(
                        examples=[
                            './demo_examples/bedroom_01.png',
                            './demo_examples/kitti_02.png',
                            './demo_examples/kitti_03.png',
                            './demo_examples/re10k_05.jpg',
                            './demo_examples/re10k_06.jpg',
                            './demo_examples/christ_church_cathedral.png',
                            './demo_examples/radcliffe.png',
                            './demo_examples/blenheim_palace_bedroom.png',
                            './demo_examples/blenheim_palace_living.png',
                            './demo_examples/blenheim_palace.JPG',
                        ],
                        inputs=[input_image],
                        cache_examples=False,
                        label="Examples",
                        examples_per_page=20,
                    )

                with gr.Row():
                    processed_image = gr.Image(label="Processed Image", interactive=False)

            with gr.Column(scale=2):
                with gr.Row():
                    with gr.Tab("Reconstruction"):
                        output_model = gr.Model3D(
                            height=640,
                            label="Output Model",
                            interactive=False
                        )

        gr.Markdown(
        """
            ## Comments:
            1. If you run the demo online, the first example you upload should take about 25 seconds (with preprocessing, saving and overhead), the following take about 14s (due to the .ply visualisation).
            2. The 3D viewer shows a .ply mesh extracted from a mix of 3D Gaussians. This is only an approximations and artefacts might show.
            3. Known limitations include:
            - a black dot appearing on the model from some viewpoints
            - while the multiple gaussians fill in resonable pixels to the invisible parts, the visual quality is still blurry.
            4. It achieves state-of-the-art results when trained and tested on RealEstate10k., and is **much** cheaper to train and run.
            5. When transferred to unseen datasets like NYU it outperforms competitors by a large margin.
            6. More impressively, when transferred to KITTI, Flash3D achieves better PSNR than methods trained specifically on that dataset. 
            ## How does it work?
            Given a single image I as input, Flash3D first estimates the metric depth D using a frozen off-the-shelf network.
            Then, a ResNet50-like encoder–decoder network predicts a set of shape and appearance parameters P of K layers of Gaussians for every pixel u,
            allowing unobserved and occluded surfaces to be modelled.
            From these predicted components, the depth can be obtained by summing the predicted (positive) offsets δi with the predicted monocular depth D,
            allowing the mean vector for every layer of Gaussians to be computed.
            This strategy ensures that the layers are depth-ordered, encouraging the network to model occluded surfaces.
            For more results see the [project page](https://www.robots.ox.ac.uk/~vgg/research/flash3d/).
            """
        )

        submit.click(fn=check_input_image, inputs=[input_image]).success(
            fn=preprocess,
            inputs=[input_image, dynamic_size, padding],
            outputs=[processed_image],
        ).success(
            fn=reconstruct_and_export,
            inputs=[processed_image],
            outputs=[output_model],
        )

    demo.queue(max_size=1)
    demo.launch()


if __name__ == "__main__":
    main()