init commit

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+gradio_splatting/frontend/node_modules/@esbuild/linux-x64/bin/esbuild filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,83 @@
+FROM nvidia/cuda:11.3.1-devel-ubuntu20.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+ENV PYTHONUNBUFFERED=1
+
+ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6"
+ENV TCNN_CUDA_ARCHITECTURES=86;80;75;70;61;60
+ENV FORCE_CUDA=1
+
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=${CUDA_HOME}/bin:/home/${USER_NAME}/.local/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
+ENV LIBRARY_PATH=${CUDA_HOME}/lib64/stubs:${LIBRARY_PATH}
+
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    build-essential \
+    curl \
+    git \
+    libegl1-mesa-dev \
+    libgl1-mesa-dev \
+    libgles2-mesa-dev \
+    libglib2.0-0 \
+    libsm6 \
+    libxext6 \
+    libxrender1 \
+    python-is-python3 \
+    python3-dev \
+    python3-pip \
+    wget \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+# Switch to the "user" user
+USER user
+# Set home to the user's home directory
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH \
+    PYTHONPATH=$HOME/app \
+	PYTHONUNBUFFERED=1 \
+	GRADIO_ALLOW_FLAGGING=never \
+	GRADIO_NUM_PORTS=1 \
+	GRADIO_SERVER_NAME=0.0.0.0 \
+	GRADIO_THEME=huggingface \
+	SYSTEM=spaces
+
+RUN pip install --upgrade pip setuptools ninja
+RUN pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
+
+RUN python -c "import torch; print(torch.version.cuda)"
+COPY requirements.txt /tmp
+RUN cd /tmp && pip install -r requirements.txt
+
+# install pointnet2_ops from snowflake
+RUN git clone https://github.com/AllenXiangX/SnowflakeNet.git /home/user/SnowflakeNet
+WORKDIR /home/user/SnowflakeNet/models/pointnet2_ops_lib
+RUN python setup.py install --user
+
+# install pytorch3d
+RUN git clone -b v0.7.3 https://github.com/facebookresearch/pytorch3d.git /home/user/pytorch3d-0.7.3
+WORKDIR /home/user/pytorch3d-0.7.3
+RUN  python setup.py install --user
+
+# install torch-scatter
+RUN git clone https://github.com/rusty1s/pytorch_scatter.git /home/user/pytorch_scatter
+WORKDIR /home/user/pytorch_scatter
+RUN  python setup.py install --user
+
+# install diff-gaussian-rasterization
+RUN git clone --recursive https://github.com/graphdeco-inria/diff-gaussian-rasterization.git /home/user/diff-gaussian-rasterization
+WORKDIR /home/user/diff-gaussian-rasterization
+RUN python setup.py install --user
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+RUN git clone https://github.com/dylanebert/gradio-splatting.git gradio_splatting
+
+CMD ["python", "app.py"]

README.md

@@ -0,0 +1,13 @@
+---
+title: TriplaneGaussian
+emoji: 👀
+colorFrom: blue
+colorTo: yellow
+sdk: docker
+# sdk: gradio
+# sdk_version: 4.13.0
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,177 @@
+import gradio as gr
+import argparse
+import os
+import glob
+import torch
+from PIL import Image
+from copy import deepcopy
+import sys
+import tempfile
+from huggingface_hub import snapshot_download
+
+LOCAL_CODE = os.environ.get("LOCAL_CODE", "1") == "1"
+AUTH = ("admin", os.environ["PASSWD"]) if "PASSWD" in os.environ else None
+
+code_dir = snapshot_download("zouzx/TriplaneGaussian", local_dir="./code", token=os.environ["HF_TOKEN"]) if not LOCAL_CODE else "./code"
+
+sys.path.append(code_dir)
+
+from utils import image_preprocess, pred_bbox, sam_init, sam_out_nosave, todevice
+from gradio_splatting.backend.gradio_model3dgs import Model3DGS
+import tgs
+from tgs.utils.config import ExperimentConfig, load_config
+from tgs.systems.infer import TGS
+
+SAM_CKPT_PATH = "code/checkpoints/sam_vit_h_4b8939.pth"
+MODEL_CKPT_PATH = "code/checkpoints/tgs_lvis_100v_rel.ckpt"
+CONFIG = "code/configs/single-rel.yaml"
+EXP_ROOT_DIR = "./outputs-gradio"
+
+gpu = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
+device = "cuda:{}".format(gpu) if torch.cuda.is_available() else "cpu"
+
+print("device: ", device)
+
+# load SAM checkpoint
+sam_predictor = sam_init(SAM_CKPT_PATH, gpu)
+print("load sam ckpt done.")
+
+# init system
+base_cfg: ExperimentConfig
+base_cfg = load_config(CONFIG, cli_args=[], n_gpus=1)
+base_cfg.system.weights = MODEL_CKPT_PATH
+system = TGS(cfg=base_cfg.system).to(device)
+print("load model ckpt done.")
+
+HEADER = """
+# Triplane Meets Gaussian Splatting: Fast and Generalizable Single-View 3D Reconstruction with Transformers
+
+<div>
+<a style="display: inline-block;" href="https://arxiv.org/abs/2312.09147"><img src="https://img.shields.io/badge/arxiv-2312.09147-B31B1B.svg"></a>
+</div>
+
+TGS enables fast reconstruction from single-view image in a few seconds based on a hybrid Triplane-Gaussian 3D representation.
+
+This model is trained on Objaverse-LVIS (~40K synthetic objects) only. And note that we normalize the input camera pose to a pre-set viewpoint during training stage following LRM, rather than directly using camera pose of input camera as implemented in our original paper.
+"""
+
+def preprocess(image_path, save_path=None, lower_contrast=False):
+    input_raw = Image.open(image_path)
+
+    input_raw.thumbnail([512, 512], Image.Resampling.LANCZOS)
+    image_sam = sam_out_nosave(
+        sam_predictor, input_raw.convert("RGB"), pred_bbox(input_raw)
+    )
+
+    if save_path is None:
+        save_path, ext = os.path.splitext(image_path)
+        save_path = save_path + "_rgba.png"
+    image_preprocess(image_sam, save_path, lower_contrast=lower_contrast, rescale=True)
+
+    return save_path
+
+def init_trial_dir():
+    if not os.path.exists(EXP_ROOT_DIR):
+        os.makedirs(EXP_ROOT_DIR, exist_ok=True)
+    trial_dir = tempfile.TemporaryDirectory(dir=EXP_ROOT_DIR).name
+    system.set_save_dir(trial_dir)
+    return trial_dir
+
+@torch.no_grad()
+def infer(image_path: str,
+          cam_dist: float,
+          fovy_deg: float,
+          only_3dgs: bool = False):
+    data_cfg = deepcopy(base_cfg.data)
+    data_cfg.only_3dgs = only_3dgs
+    data_cfg.cond_fovy_deg = fovy_deg
+    data_cfg.cond_camera_distance = cam_dist
+    data_cfg.image_list = [image_path]
+    dm = tgs.find(base_cfg.data_cls)(data_cfg)
+
+    dm.setup()
+    for batch_idx, batch in enumerate(dm.test_dataloader()):
+        batch = todevice(batch, device)
+        system.test_step(batch, batch_idx, save_3dgs=only_3dgs)
+    if not only_3dgs:
+        system.on_test_epoch_end()
+
+def run(image_path: str,
+        cam_dist: float,
+        fov_degree: float):
+    infer(image_path, cam_dist, fov_degree, only_3dgs=True)
+    save_path = system.get_save_dir()
+    gs = glob.glob(os.path.join(save_path, "*.ply"))[0]
+    return gs
+
+def run_video(image_path: str,
+            cam_dist: float,
+            fov_degree: float):
+    infer(image_path, cam_dist, fov_degree)
+    save_path = system.get_save_dir()
+    video = glob.glob(os.path.join(save_path, "*.mp4"))[0]
+    return video
+
+def launch(port):
+    with gr.Blocks(
+        title="TGS - Demo",
+        theme=gr.themes.Monochrome()
+    ) as demo:
+        with gr.Row(variant='panel'):
+            gr.Markdown(HEADER)
+    
+        with gr.Row(variant='panel'):
+            with gr.Column(scale=1):
+                input_image = gr.Image(value=None, width=512, height=512, type="filepath", label="Input Image")
+                fov_deg_slider = gr.Slider(20, 80, value=40, step=1, label="Camera Fov Degree")
+                camera_dist_slider = gr.Slider(1.0, 4.0, value=1.6, step=0.1, label="Camera Distance")
+                img_run_btn = gr.Button("Reconstruction")
+
+                gr.Examples(
+                    examples=[
+                        "example_images/green_parrot.webp",
+                        "example_images/rusty_gameboy.webp",
+                        "example_images/a_pikachu_with_smily_face.webp",
+                        "example_images/an_otter_wearing_sunglasses.webp",
+                        "example_images/lumberjack_axe.webp",
+                        "example_images/medieval_shield.webp"
+                    ],
+                    inputs=[input_image],
+                    cache_examples=False,
+                    label="Examples",
+                    examples_per_page=40
+                )
+            
+            with gr.Column(scale=1):
+                with gr.Row(variant='panel'):
+                    seg_image = gr.Image(value=None, type="filepath", height=256, width=256, image_mode="RGBA", label="Segmented Image", interactive=False)
+                    output_video = gr.Video(value=None, label="Video", height=256, autoplay=True)
+                output_3dgs = Model3DGS(value=None, label="3DGS")
+        
+        img_run_btn.click(
+            fn=preprocess,
+            inputs=[input_image],
+            outputs=[seg_image],
+            concurrency_limit=1,
+        ).success(
+            fn=init_trial_dir,
+            concurrency_limit=1,
+        ).success(fn=run,
+                inputs=[seg_image, camera_dist_slider, fov_deg_slider],
+                outputs=[output_3dgs],
+                concurrency_limit=1
+        ).success(fn=run_video,
+                inputs=[seg_image, camera_dist_slider, fov_deg_slider],
+                outputs=[output_video],
+                concurrency_limit=1)
+
+        launch_args = {"server_port": port}
+        demo.queue(max_size=10)
+        demo.launch(auth=AUTH, **launch_args)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    args, extra = parser.parse_known_args()
+    parser.add_argument("--port", type=int, default=7860)
+    args = parser.parse_args()
+    launch(args.port)

requirements.txt

@@ -0,0 +1,14 @@
+lightning==2.0.7
+pytorch-lightning==2.0.2
+plyfile
+OmegaConf
+matplotlib
+einops
+gradio
+diffusers==0.19.3
+transformers==4.34.1
+rembg
+segment_anything
+jaxtyping
+imageio
+imageio-ffmpeg

utils.py

@@ -0,0 +1,135 @@
+import os
+import time
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+from rembg import remove
+from segment_anything import SamPredictor, sam_model_registry
+import urllib.request
+from tqdm import tqdm
+
+
+def sam_init(sam_checkpoint, device_id=0):
+    # sam_checkpoint = os.path.join(os.path.dirname(__file__), "./sam_vit_h_4b8939.pth")
+    model_type = "vit_h"
+
+    device = "cuda:{}".format(device_id) if torch.cuda.is_available() else "cpu"
+
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=device)
+    predictor = SamPredictor(sam)
+    return predictor
+
+
+def sam_out_nosave(predictor, input_image, *bbox_sliders):
+    bbox = np.array(bbox_sliders)
+    image = np.asarray(input_image)
+
+    start_time = time.time()
+    predictor.set_image(image)
+
+    masks_bbox, scores_bbox, logits_bbox = predictor.predict(
+        box=bbox, multimask_output=True
+    )
+
+    out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
+    out_image[:, :, :3] = image
+    out_image_bbox = out_image.copy()
+    out_image_bbox[:, :, 3] = (
+        masks_bbox[-1].astype(np.uint8) * 255
+    )  # np.argmax(scores_bbox)
+    torch.cuda.empty_cache()
+    return Image.fromarray(out_image_bbox, mode="RGBA")
+
+
+# contrast correction, rescale and recenter
+def image_preprocess(input_image, save_path, lower_contrast=True, rescale=True):
+    image_arr = np.array(input_image)
+    in_w, in_h = image_arr.shape[:2]
+
+    if lower_contrast:
+        alpha = 0.8  # Contrast control (1.0-3.0)
+        beta = 0  # Brightness control (0-100)
+        # Apply the contrast adjustment
+        image_arr = cv2.convertScaleAbs(image_arr, alpha=alpha, beta=beta)
+        image_arr[image_arr[..., -1] > 200, -1] = 255
+
+    ret, mask = cv2.threshold(
+        np.array(input_image.split()[-1]), 0, 255, cv2.THRESH_BINARY
+    )
+    x, y, w, h = cv2.boundingRect(mask)
+    max_size = max(w, h)
+    ratio = 0.75
+    if rescale:
+        side_len = int(max_size / ratio)
+    else:
+        side_len = in_w
+    padded_image = np.zeros((side_len, side_len, 4), dtype=np.uint8)
+    center = side_len // 2
+    padded_image[
+        center - h // 2 : center - h // 2 + h, center - w // 2 : center - w // 2 + w
+    ] = image_arr[y : y + h, x : x + w]
+    rgba = Image.fromarray(padded_image).resize((256, 256), Image.LANCZOS)
+    rgba.save(save_path)
+
+    # rgba_arr = np.array(rgba) / 255.0
+    # rgb = rgba_arr[...,:3] * rgba_arr[...,-1:] + (1 - rgba_arr[...,-1:])
+    # return Image.fromarray((rgb * 255).astype(np.uint8))
+
+
+def pred_bbox(image):
+    image_nobg = remove(image.convert("RGBA"), alpha_matting=True)
+    alpha = np.asarray(image_nobg)[:, :, -1]
+    x_nonzero = np.nonzero(alpha.sum(axis=0))
+    y_nonzero = np.nonzero(alpha.sum(axis=1))
+    x_min = int(x_nonzero[0].min())
+    y_min = int(y_nonzero[0].min())
+    x_max = int(x_nonzero[0].max())
+    y_max = int(y_nonzero[0].max())
+    return x_min, y_min, x_max, y_max
+
+# convert a function into recursive style to handle nested dict/list/tuple variables
+def make_recursive_func(func):
+    def wrapper(vars, *args, **kwargs):
+        if isinstance(vars, list):
+            return [wrapper(x, *args, **kwargs) for x in vars]
+        elif isinstance(vars, tuple):
+            return tuple([wrapper(x, *args, **kwargs) for x in vars])
+        elif isinstance(vars, dict):
+            return {k: wrapper(v, *args, **kwargs) for k, v in vars.items()}
+        else:
+            return func(vars, *args, **kwargs)
+
+    return wrapper
+
+@make_recursive_func
+def todevice(vars, device="cuda"):
+    if isinstance(vars, torch.Tensor):
+        return vars.to(device)
+    elif isinstance(vars, str):
+        return vars
+    elif isinstance(vars, bool):
+        return vars
+    elif isinstance(vars, float):
+        return vars
+    elif isinstance(vars, int):
+        return vars
+    else:
+        raise NotImplementedError("invalid input type {} for tensor2numpy".format(type(vars)))
+
+def download_checkpoint(url, save_path):
+    try:
+        with urllib.request.urlopen(url) as response, open(save_path, 'wb') as file:
+            file_size = int(response.info().get('Content-Length', -1))
+            chunk_size = 8192
+            num_chunks = file_size // chunk_size if file_size > chunk_size else 1
+
+            with tqdm(total=file_size, unit='B', unit_scale=True, desc='Downloading', ncols=100) as pbar:
+                for chunk in iter(lambda: response.read(chunk_size), b''):
+                    file.write(chunk)
+                    pbar.update(len(chunk))
+        
+        print(f"Checkpoint downloaded and saved to: {save_path}")
+    except Exception as e:
+        print(f"Error downloading checkpoint: {e}")