app.py

import os
import torch
import subprocess
import gradio as gr
from tqdm import tqdm
import post_process
import shutil
from sugar.gaussian_splatting import convert
from sugar.gaussian_splatting import train as train_gs
from sugar import train as train_sugar


def process_images(images):
    print("Preparing directories...")
    # Specify the target folder where you want to copy the images
    target_folder = "data/custom/input/"

    # Remove and Create the target folder if it doesn't exist
    if os.path.exists(target_folder):
        os.removedirs(target_folder)
    os.makedirs(target_folder, exist_ok=True)

    # Copy each image to the target folder
    file_paths = []
    for image in tqdm(images):
        image_path = image.name  # Get the original image path
        new_path = os.path.join(target_folder, os.path.basename(image_path))  # Construct the new path
        shutil.move(image_path, new_path)  # Move the image to the target folder
        file_paths.append(new_path)

    return file_paths


def get_3d(file_paths):
    # Your 3D model reconstruction logic using the file paths goes here
    # For demonstration, I'll just return the file paths
    print(file_paths)
    file_paths = process_images(file_paths)
    print(file_paths)

    # 3d reconstruction
    torch.cuda.empty_cache()

    # subprocess.run(['python', 'sugar/gaussian_splatting/convert.py', '-s', 'data/custom/'])
    convert.main()
    torch.cuda.empty_cache()

    # subprocess.run(['python', 'sugar/gaussian_splatting/train.py', '-s', 'data/custom/', '--iterations', '7000', '-m', 'output/custom/'])
    train_gs.main()
    torch.cuda.empty_cache()

    # subprocess.run(['python', 'sugar/train.py', '-c', 'output/custom/', '-r', 'density', '-t', 'True', '--export_ply', 'True'])
    train_sugar.main()
    torch.cuda.empty_cache()

    out_dir = "output/final_mesh/custom/"
    os.makedirs(out_dir, exist_ok=True)
    post_process.post_process_textured_mesh("output/refined_mesh/custom/sugarfine_3Dgs7000_densityestim02_sdfnorm02_level03_decim1000000_normalconsistency01_gaussperface6.obj", "output/final_mesh/custom/custom.obj", crop=True, remesh=False, repair=True, viz=False, save=True)
    torch.cuda.empty_cache()

    return f"output/refined_mesh/custom/sugarfine_3Dgs7000_densityestim02_sdfnorm02_level03_decim1000000_normalconsistency01_gaussperface6.obj", f"output/final_mesh/custom/custom.glb"


demo = gr.Interface(
    fn=get_3d,  # Function that will be called when the button is clicked
    inputs=gr.Files(file_types=["image", "video"], file_count="multiple"),
    outputs=[
        gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model (Recon.)"),
        gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model (Final)"),
    ],
    title="Alteredverse::3D Model Reconstruction",
    description="Upload a video or set of images to reconstruct a 3D model.",
)


if __name__ == "__main__":
    demo.queue()
    # demo.launch(share=True)
    demo.launch()