# Not ready to use yet 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 import rembg 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 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, ) -> PIL.Image.Image: ensure_dis_onnx_model() # Create a temporary file to save the image with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as temp: # Save the PIL image to the temporary file image.save(temp.name) extracted_img, mask = dis_remove_background(DIS_ONNX_MODEL_PATH, temp.name) if isinstance(extracted_img, np.ndarray): if mask.dtype != np.uint8: mask = (np.clip(mask, 0, 1) * 255).astype(np.uint8) if mask.ndim == 3: mask = mask[..., 0] image = image.convert("RGBA") image_np = np.array(image) image_np[..., 3] = mask return Image.fromarray(image_np) 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 hex_to_rgb(hex_color: str) -> tuple[int, int, int]: """Converts a hex color string to an RGB tuple.""" hex_color = hex_color.lstrip('#') return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4)) def preprocess_image(image, background_choice, foreground_ratio, backgroud_color): """ Preprocesses the input image by optionally removing the background, resizing, and adding a new solid background. Returns an RGB PIL Image. """ if image.mode != 'RGBA': image = image.convert('RGBA') if background_choice == "Auto Remove background": image = remove_background(image) if image is None: raise gr.Error("Background removal failed. Please check the input image and ensure the model file exists and is valid.") # Resize the content of the image image = do_resize_content(image, foreground_ratio) # Add a solid background color rgb_background = hex_to_rgb(backgroud_color) image_with_bg = add_background(image, rgb_background) # Convert to RGB and return return image_with_bg.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", ) parser.add_argument("--device", type=str, default="cuda") 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(args.device) 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 = ''' 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()