File size: 10,218 Bytes
4932a8f f4e8cf6 5a9aacb 1e3c695 f4e8cf6 0ac4eff f4e8cf6 135bed3 5a9aacb 135bed3 f4e8cf6 bb9cdff b6b723a f4e8cf6 bb9cdff f4e8cf6 5a9aacb f4e8cf6 5a9aacb 1e3c695 fa6cc29 a3f04cd fa6cc29 0ac4eff f4e8cf6 bb9cdff f4e8cf6 867f45c 0ac4eff 135bed3 867f45c 1452c34 05a48c5 867f45c 1fa688f 867f45c 1fa688f bb9cdff 5407b39 1fa688f 9551df9 bb9cdff 9551df9 bfac008 bb9cdff f4e8cf6 e3cc233 50e2279 1eb08e4 552ea8a 1eb08e4 50e2279 732c53f 8e3405b 1eb08e4 732c53f 1eb08e4 50e2279 732c53f bb9cdff 1eb08e4 50e2279 1eb08e4 bb9cdff 50e2279 bb9cdff 1eb08e4 50e2279 bb9cdff 50e2279 81d5d11 50e2279 1eb08e4 50e2279 bb9cdff 50e2279 bb9cdff 50e2279 bb9cdff 50e2279 bb9cdff 50e2279 552ea8a 1fa688f |
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 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 |
import spaces
import argparse
import numpy as np
import gradio as gr
from omegaconf import OmegaConf
import torch
from PIL import Image
import PIL
from pipelines import TwoStagePipeline
from huggingface_hub import hf_hub_download
import os
from typing import Any
import json
import os
import json
import argparse
import requests
import tempfile
from model import CRM
from inference import generate3d
from dis_bg_remover import remove_background as dis_remove_background
# Configurable ONNX model path (can be set via environment variable)
DIS_ONNX_MODEL_PATH = os.environ.get("DIS_ONNX_MODEL_PATH", "isnet_dis.onnx")
DIS_ONNX_MODEL_URL = "https://huggingface.co/stoned0651/isnet_dis.onnx/resolve/main/isnet_dis.onnx"
pipeline = None
def expand_to_square(image, bg_color=(0, 0, 0, 0)):
# expand image to 1:1
width, height = image.size
if width == height:
return image
new_size = (max(width, height), max(width, height))
new_image = Image.new("RGBA", new_size, bg_color)
paste_position = ((new_size[0] - width) // 2, (new_size[1] - height) // 2)
new_image.paste(image, paste_position)
return new_image
def check_input_image(input_image):
if input_image is None:
raise gr.Error("No image uploaded!")
def ensure_dis_onnx_model():
if not os.path.exists(DIS_ONNX_MODEL_PATH):
try:
print(f"Model file not found at {DIS_ONNX_MODEL_PATH}. Downloading from {DIS_ONNX_MODEL_URL}...")
response = requests.get(DIS_ONNX_MODEL_URL, stream=True)
response.raise_for_status()
with open(DIS_ONNX_MODEL_PATH, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
print(f"Downloaded model to {DIS_ONNX_MODEL_PATH}")
except Exception as e:
raise gr.Error(
f"Failed to download DIS background remover model file: {e}\n"
f"Please manually download it from {DIS_ONNX_MODEL_URL} and place it in the project directory or set the DIS_ONNX_MODEL_PATH environment variable."
)
def remove_background(
image: PIL.Image.Image,
rembg_session: Any = None,
force: bool = False,
**rembg_kwargs,
) -> PIL.Image.Image:
ensure_dis_onnx_model()
with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as temp:
image.save(temp.name)
extracted_img, mask = dis_remove_background(DIS_ONNX_MODEL_PATH, temp.name)
# If extracted_img is a mask (single channel), use it as alpha for the original image
if isinstance(extracted_img, np.ndarray):
# If mask is float, convert to uint8
if mask.dtype != np.uint8:
mask = (np.clip(mask, 0, 1) * 255).astype(np.uint8)
# Ensure mask is 2D
if mask.ndim == 3:
mask = mask[..., 0]
# Convert original image to RGBA
image = image.convert("RGBA")
image_np = np.array(image)
image_np[..., 3] = mask
return Image.fromarray(image_np)
# If extracted_img is already a color image, just return it
return extracted_img
def do_resize_content(original_image: Image, scale_rate):
# resize image content wile retain the original image size
if scale_rate != 1:
# Calculate the new size after rescaling
new_size = tuple(int(dim * scale_rate) for dim in original_image.size)
# Resize the image while maintaining the aspect ratio
resized_image = original_image.resize(new_size)
# Create a new image with the original size and black background
padded_image = Image.new("RGBA", original_image.size, (0, 0, 0, 0))
paste_position = ((original_image.width - resized_image.width) // 2, (original_image.height - resized_image.height) // 2)
padded_image.paste(resized_image, paste_position)
return padded_image
else:
return original_image
def add_background(image, bg_color=(255, 255, 255)):
# given an RGBA image, alpha channel is used as mask to add background color
background = Image.new("RGBA", image.size, bg_color)
return Image.alpha_composite(background, image)
def preprocess_image(image, background_choice, foreground_ratio, backgroud_color):
"""
input image is a pil image in RGBA, return RGB image
"""
print(background_choice)
if background_choice == "Alpha as mask":
background = Image.new("RGBA", image.size, (0, 0, 0, 0))
image = Image.alpha_composite(background, image)
else:
image = remove_background(image, force=True)
if image is None:
raise gr.Error("Background removal failed. Please check the input image and ensure the model file exists and is valid.")
image = do_resize_content(image, foreground_ratio)
image = expand_to_square(image)
image = add_background(image, backgroud_color)
return image.convert("RGB")
@spaces.GPU
def gen_image(input_image, seed, scale, step):
global pipeline, model, args
pipeline.set_seed(seed)
rt_dict = pipeline(input_image, scale=scale, step=step)
stage1_images = rt_dict["stage1_images"]
stage2_images = rt_dict["stage2_images"]
np_imgs = np.concatenate(stage1_images, 1)
np_xyzs = np.concatenate(stage2_images, 1)
glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path#, obj_path
parser = argparse.ArgumentParser()
parser.add_argument(
"--stage1_config",
type=str,
default="configs/nf7_v3_SNR_rd_size_stroke.yaml",
help="config for stage1",
)
parser.add_argument(
"--stage2_config",
type=str,
default="configs/stage2-v2-snr.yaml",
help="config for stage2",
)
# Force CPU usage
parser.add_argument("--device", type=str, default="cpu")
args = parser.parse_args()
crm_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="CRM.pth")
specs = json.load(open("configs/specs_objaverse_total.json"))
model = CRM(specs)
model.load_state_dict(torch.load(crm_path, map_location="cpu"), strict=False)
model = model.to("cpu")
stage1_config = OmegaConf.load(args.stage1_config).config
stage2_config = OmegaConf.load(args.stage2_config).config
stage2_sampler_config = stage2_config.sampler
stage1_sampler_config = stage1_config.sampler
stage1_model_config = stage1_config.models
stage2_model_config = stage2_config.models
xyz_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="ccm-diffusion.pth")
pixel_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="pixel-diffusion.pth")
stage1_model_config.resume = pixel_path
stage2_model_config.resume = xyz_path
pipeline = TwoStagePipeline(
stage1_model_config,
stage2_model_config,
stage1_sampler_config,
stage2_sampler_config,
device=args.device,
dtype=torch.float32
)
_DESCRIPTION = '''
* Our [official implementation](https://github.com/thu-ml/CRM) uses UV texture instead of vertex color. It has better texture than this online demo.
* Project page of CRM: https://ml.cs.tsinghua.edu.cn/~zhengyi/CRM/
* If you find the output unsatisfying, try using different seeds:)
'''
with gr.Blocks() as demo:
gr.Markdown("# CRM: Single Image to 3D Textured Mesh with Convolutional Reconstruction Model")
gr.Markdown(_DESCRIPTION)
with gr.Row():
with gr.Column():
with gr.Row():
image_input = gr.Image(
label="Image input",
image_mode="RGBA",
sources="upload",
type="pil",
)
processed_image = gr.Image(label="Processed Image", interactive=False, type="pil", image_mode="RGB")
with gr.Row():
with gr.Column():
with gr.Row():
background_choice = gr.Radio([
"Alpha as mask",
"Auto Remove background"
], value="Auto Remove background",
label="backgroud choice")
# do_remove_background = gr.Checkbox(label=, value=True)
# force_remove = gr.Checkbox(label=, value=False)
back_groud_color = gr.ColorPicker(label="Background Color", value="#7F7F7F", interactive=False)
foreground_ratio = gr.Slider(
label="Foreground Ratio",
minimum=0.5,
maximum=1.0,
value=1.0,
step=0.05,
)
with gr.Column():
seed = gr.Number(value=1234, label="seed", precision=0)
guidance_scale = gr.Number(value=5.5, minimum=3, maximum=10, label="guidance_scale")
step = gr.Number(value=30, minimum=30, maximum=100, label="sample steps", precision=0)
text_button = gr.Button("Generate 3D shape")
gr.Examples(
examples=[os.path.join("examples", i) for i in os.listdir("examples")],
inputs=[image_input],
examples_per_page = 20,
)
with gr.Column():
image_output = gr.Image(interactive=False, label="Output RGB image")
xyz_ouput = gr.Image(interactive=False, label="Output CCM image")
output_model = gr.Model3D(
label="Output OBJ",
interactive=False,
)
gr.Markdown("Note: Ensure that the input image is correctly pre-processed into a grey background, otherwise the results will be unpredictable.")
inputs = [
processed_image,
seed,
guidance_scale,
step,
]
outputs = [
image_output,
xyz_ouput,
output_model,
# output_obj,
]
text_button.click(fn=check_input_image, inputs=[image_input]).success(
fn=preprocess_image,
inputs=[image_input, background_choice, foreground_ratio, back_groud_color],
outputs=[processed_image],
).success(
fn=gen_image,
inputs=inputs,
outputs=outputs,
)
demo.queue().launch()
|