app.py

import numpy as np

import gradio as gr


import os

import tempfile

import shutil

# from gradio_inter.predict_from_file import predict_from_file
from gradio_inter.create_bash_file import create_bash_file

from sample.reconstruct_data_taco import reconstruct_from_file

def create_temp_file(path: str) -> str:
    temp_dir = tempfile.gettempdir()
    temp_folder = os.path.join(temp_dir, "denoising")
    os.makedirs(temp_folder, exist_ok=True)
    # Clean up directory
    # for i in os.listdir(temp_folder):
    #     print("Removing", i)
    #     os.remove(os.path.join(temp_folder, i))

    temp_path = os.path.join(temp_folder, path.split("/")[-1])
    shutil.copy2(path, temp_path)
    return temp_path

# from gradio_inter.predict import predict_from_data
# from gradio_inter.predi

def transpose(matrix):
    return matrix.T


def predict(file_path: str):
    temp_file_path = create_temp_file(file_path)
    # predict_from_file
    temp_bash_file = create_bash_file(temp_file_path)
    
    os.system(f"bash {temp_bash_file}")
    
    res_file_path = "/tmp/denoising/save/predicted_infos_seed_0_tag_20231104_017_jts_spatial_t_100__st_0.npy"
    
    saved_path = reconstruct_from_file(temp_file_path)
    
    return saved_path

def create_demo():

    USAGE = """# GeneOH Diffusion: Towards Generalizable Hand-Object Interaction Denoising via Denoising Diffusion
    **[Project](https://meowuu7.github.io/GeneOH-Diffusion/) | [Paper](https://openreview.net/pdf?id=FvK2noilxT) | [Github](https://github.com/Meowuu7/GeneOH-Diffusion)**
    ## Input data format
    Currently, the demo accepts a `.pkl` file containing an hand-object sequence organized as the following format:
    ```python
        {
            "hand_pose": numpy.ndarray(seq_length, 48), # MANO pose at each frame
            "hand_trans": numpy.ndarray(seq_length, 3), # hand global translation at each frmae
            "hand_shape": numpy.ndarray(10), # MANO shape coefficients
            "hand_verts": numpy.ndarray(seq_length, 778, 3), # MANO hand vertices
            "hand_faces": numpy.ndarray(1538, 3), # MANO hand faces
            "obj_verts": numpy.ndarray(seq_length, num_obj_verts, 3), # object vertices at each frame
            "obj_faces": numpy.ndarray(num_obj_faces, 3), # object faces
            "obj_pose": numpy.ndarray(seq_length, 4, 4), # object pose at each frame
        }
    ```
    We provide an example [here](https://drive.google.com/file/d/17oqKMhQNpRqSdApyuuCmTrPkrFl0Cqp6/view?usp=sharing). **The demo is under developing and will support more data formats in the future.**
    
    
    ## To run the demo, 
    1. Upload a `pickle` file to the left box by draging your file or clicking the box to open the file explorer.
    2. Clik the `Submit` button to run the demo.
    3. The denoised sequence will be output as a `.npy` file and can be downloaded from the right box. 
    
    Since the model runs on CPU currently, the speed is not very fast. For instance, it takes abount 1200s to process the [example](https://drive.google.com/file/d/17oqKMhQNpRqSdApyuuCmTrPkrFl0Cqp6/view?usp=sharing) mentioned above which contains 288 frames. Please be patient and wait for the result. 
    
    To run the model faster, please visit our [github repo](https://github.com/Meowuu7/GeneOH-Diffusion), follow the instructions and run the model on your own server or local machine.
    
    ## Output data format 
    The output is a `.npy` file containing the denoised sequence organized as the following format:
    ```python
        {
            "predicted_info": {
                "targets": numpy.ndarray(seq_length, num_mano_joints, 3), # input MANO joints
                "outputs": numpy.ndarray(seq_length, num_mano_joints, 3), # denoised MANO joints
                "obj_verts": numpy.ndarray(seq_length, num_obj_verts, 3), # object vertices at each frame
                "obj_faces": numpy.ndarray(num_obj_faces, 3), # object faces
                ... # others
            }
            "bf_ct_verts": numpy.ndarray(seq_length, 778, 3), # denoised MANO vertices
            "bf_ct_rot_var": numpy.ndarray(seq_length, 3), # denoised MANO global rotation coefficients
            "bf_ct_theta_var": numpy.ndarray(seq_length, 45), # denoised MANO global pose coefficients
            "bf_ct_beta_var": numpy.ndarray(1, 10), # denoised MANO shape coefficients
            "bf_ct_transl_var": numpy.ndarray(seq_length, 3), # denoised hand global translation
        }
    ```
    The corresponding output file of the [example](https://drive.google.com/file/d/17oqKMhQNpRqSdApyuuCmTrPkrFl0Cqp6/view?usp=sharing) mentioned above can be downloaded [here](https://drive.google.com/file/d/1Ah-qwV6LXlOyaBBe0qQRu1lN-BpKt2Y3/view?usp=sharing).
    """


    with gr.Blocks() as demo:

        gr.Markdown(USAGE)

        # # demo = 
        # gr.Interface(
        #     predict,
        #     # gr.Dataframe(type="numpy", datatype="number", row_count=5, col_count=3),
        #     gr.File(type="filepath"),
        #     gr.File(type="filepath"),
        #     cache_examples=False
        # )


        input_file = gr.File(type="filepath")
        output_file = gr.File(type="filepath")
        
        gr.Interface(
            predict,
            # gr.Dataframe(type="numpy", datatype="number", row_count=5, col_count=3),
            input_file,
            output_file,
            cache_examples=False
        )
        
        inputs = input_file
        outputs = output_file
        gr.Examples(
            examples=[os.path.join(os.path.dirname(__file__), "./gradio_inter/20231104_017.pkl"), os.path.join(os.path.dirname(__file__), "./gradio_inter/20231104_010.pkl")],
            inputs=inputs,
            fn=predict,
            outputs=outputs,
        )

    return demo
    
if __name__ == "__main__":
    demo = create_demo()
    demo.launch()