app.py

#!/usr/bin/env python

from __future__ import annotations

import os
import datetime

import gradio as gr
import spaces


@spaces.GPU(duration=60 * 3)
def run_on_gpu(input_point_cloud: gr.utils.NamedString,
               gen_resolution_global: int,
               padding_factor: float,
               gen_subsample_manifold_iter: int,
               gen_refine_iter: int):
    print('Started inference at {}'.format(datetime.datetime.now()))
    print('Inputs:', input_point_cloud, gen_resolution_global, padding_factor,
          gen_subsample_manifold_iter, gen_refine_iter)
    print('Types:', type(input_point_cloud), type(gen_resolution_global), type(padding_factor),
          type(gen_subsample_manifold_iter), type(gen_refine_iter))

    import os
    import sys
    sys.path.append(os.path.abspath('ppsurf'))
    # from ppsurf.pps import cli_main
    import subprocess

    in_file = '{}'.format(input_point_cloud.name)

    # append 'rec' to the input file name
    splitext_result = os.path.splitext(in_file)
    out_file = splitext_result[0] + '_rec' + splitext_result[1]
    out_dir = os.path.dirname(out_file)
    model_path = 'models/ppsurf_50nn/version_0/checkpoints/last.ckpt'

    args = [
        'pps.py', 'predict',
        '-c', 'ppsurf/configs/poco.yaml',
        '-c', 'ppsurf/configs/ppsurf.yaml',
        '-c', 'ppsurf/configs/ppsurf_50nn.yaml',
        '--ckpt_path', model_path,
        '--data.init_args.in_file', in_file,
        '--model.init_args.results_dir', out_dir,
        '--trainer.logger', 'False',
        '--trainer.devices', '1',
        '--model.init_args.gen_resolution_global', str(gen_resolution_global),
        '--data.init_args.padding_factor', str(padding_factor),
        '--model.init_args.gen_subsample_manifold_iter', str(gen_subsample_manifold_iter),
        '--model.init_args.gen_refine_iter', str(gen_refine_iter),
        ]

    sys.argv = args
    subprocess.run(['python', 'ppsurf/pps.py'] + args[1:])
    # cli_main()
    print('Finished inference at {}'.format(datetime.datetime.now()))

    result_3d_model = out_file
    output_file = out_file
    progress_text = 'done'

    return result_3d_model, output_file, progress_text


def main():
    description = '''# [PPSurf](https://github.com/cg-tuwien/ppsurf)
    
    Supported file formats: 
    - PLY, STL, OBJ and other mesh files, 
    - XYZ as whitespace-separated text file, 
    - NPY and NPZ (key='arr_0'), 
    - LAS and LAZ (version 1.0-1.4), COPC and CRS.
    Best results for 50k-250k points. 
    
    This method is meant for scans of single and few objects. 
    Quality for scenes and landscapes will be lower.
    
    Inference takes about 2 minutes.
    '''

    # can't render many input types directly in Gradio Model3D
    # so we need to convert to supported format
    # Gradio can't draw point clouds anyway, so we skip this for now
    # def convert_to_ply(input_point_cloud_upload: gr.utils.NamedString):
    #
    #     # add absolute path to import dirs
    #     import sys
    #     import os
    #     sys.path.append(os.path.abspath('ppsurf'))
    #
    #     # import os
    #     # os.chdir('ppsurf')
    #
    #     print('Inputs:', input_point_cloud_upload, type(input_point_cloud_upload))
    #     input_shape: str = input_point_cloud_upload.name
    #     if not input_shape.endswith('.ply'):
    #         # load file
    #         from ppsurf.source.occupancy_data_module import OccupancyDataModule
    #         pts_np = OccupancyDataModule.load_pts(input_shape)
    #
    #         # convert to ply
    #         import trimesh
    #         mesh = trimesh.Trimesh(vertices=pts_np[:, :3])
    #         input_shape = input_shape + '.ply'
    #         mesh.export(input_shape)
    #
    #     print('ls:\n', subprocess.run(['ls', os.path.dirname(input_shape)]))
    #
    #     # show in viewer
    #     print(type(input_tabs))
    #     # print(type(input_point_cloud_viewer))
    #     # input_tabs.selected = 'pc_viewer'
    #     # input_point_cloud_viewer.value = input_shape

    if (SPACE_ID := os.getenv('SPACE_ID')) is not None:
        description += (f'\n<p>For faster inference without waiting in queue, '
                        f'you may duplicate the space and upgrade to GPU in settings. '
                        f'<a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true">'
                        f'<img style="display: inline; margin-top: 0em; margin-bottom: 0em" '
                        f'src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>')

    with gr.Blocks(css='style.css') as demo:
        gr.Markdown(description)
        with gr.Row():
            with gr.Column():
                # with gr.Tabs() as input_tabs:
                #     with gr.TabItem(label='Input Point Cloud Upload', id='pc_upload'):
                input_point_cloud_upload = gr.File(show_label=False, file_count='single')
                        # input_point_cloud_upload.upload(
                        #     fn=convert_to_ply,
                        #     inputs=[
                        #         input_point_cloud_upload,
                        #     ],
                        #     outputs=[
                        #         # input_point_cloud_viewer,  # not available here
                        #     ])
                    # with gr.TabItem(label='Input Point Cloud Viewer', id='pc_viewer'):
                    #     input_point_cloud_viewer = gr.Model3D(show_label=False)
                gen_resolution_global = gr.Slider(
                    label='Grid Resolution (larger for more details)',
                    minimum=17, maximum=513, value=129, step=2)
                padding_factor = gr.Slider(
                    label='Padding Factor (larger if object is cut off at boundaries)',
                    minimum=0, maximum=1.0, value=0.05, step=0.05)
                gen_subsample_manifold_iter = gr.Slider(
                    label='Subsample Manifold Iterations (larger for larger point clouds)',
                    minimum=3, maximum=30, value=10, step=1)
                gen_refine_iter = gr.Slider(
                    label='Edge Refinement Iterations (larger for more details)',
                    minimum=3, maximum=30, value=10, step=1)
            with gr.Column():
                progress_text = gr.Text(label='Progress')
                with gr.Tabs():
                    with gr.TabItem(label='Reconstructed 3D model'):
                        result_3d_model = gr.Model3D(show_label=False)
                    with gr.TabItem(label='Output mesh file'):
                        output_file = gr.File(show_label=False)
        # with gr.Row():
        #     examples = [
        #         ['shapes/dragon1.obj', 'a photo of a dragon', 0, 7.5],
        #         ['shapes/dragon2.obj', 'a photo of a dragon', 0, 7.5],
        #         ['shapes/eagle.obj', 'a photo of an eagle', 0, 7.5],
        #         ['shapes/napoleon.obj', 'a photo of Napoleon Bonaparte', 3, 7.5],
        #         ['shapes/nascar.obj', 'A next gen nascar', 2, 10],
        #     ]
        #     gr.Examples(examples=examples,
        #                 inputs=[
        #                     input_point_cloud_viewer,
        #                     text,
        #                     seed,
        #                     guidance_scale,
        #                 ],
        #                 outputs=[
        #                     result_3d_model,
        #                     output_file,
        #                 ],
        #                 cache_examples=False)

        with gr.Row():
            run_button = gr.Button('Reconstruct with PPSurf')

        run_button.click(fn=run_on_gpu,
                         inputs=[
                             input_point_cloud_upload,
                             gen_resolution_global,
                             padding_factor,
                             gen_subsample_manifold_iter,
                             gen_refine_iter,
                         ],
                         outputs=[
                             result_3d_model,
                             output_file,
                             progress_text,
                         ])

    demo.queue(max_size=5)
    demo.launch(debug=True)


if __name__ == '__main__':
    main()