Spaces:
Sleeping
Sleeping
File size: 4,825 Bytes
2e968f0 |
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 |
import logging
import os
import time
import numpy as np
from PIL import Image, ImageOps
import numpy as np
import torch
import xatlas
from PIL import Image
from tsr.system import TSR
from tsr.utils import save_video
from tsr.bake_texture import bake_texture
class Timer:
def __init__(self):
self.items = {}
self.time_scale = 1000.0 # ms
self.time_unit = "ms"
def start(self, name: str) -> None:
if torch.cuda.is_available():
torch.cuda.synchronize()
self.items[name] = time.time()
logging.info(f"{name} ...")
def end(self, name: str) -> float:
if name not in self.items:
return
if torch.cuda.is_available():
torch.cuda.synchronize()
start_time = self.items.pop(name)
delta = time.time() - start_time
t = delta * self.time_scale
logging.info(f"{name} finished in {t:.2f}{self.time_unit}.")
def initialize_model(pretrained_model_name_or_path="stabilityai/TripoSR",
chunk_size=8192,
device="cuda:0" if torch.cuda.is_available() else "cpu"):
timer.start("Initializing model")
model = TSR.from_pretrained(
pretrained_model_name_or_path,
config_name="config.yaml",
weight_name="model.ckpt",
)
model.renderer.set_chunk_size(chunk_size)
model.to(device)
timer.end("Initializing model")
return model
def remove_background(image_path, output_path, background_value=127, new_size=(425, 425)):
# Open the image
image = Image.open(image_path).convert("RGBA")
# Split the image into its respective channels
r, g, b, alpha = image.split()
# Convert the alpha channel to binary mask where transparency is 0 and opaque is 255
alpha = ImageOps.invert(alpha)
# Replace the transparent areas with the specified background value
background = Image.new("L", image.size, color=background_value)
image_rgb = Image.composite(background, r, alpha), Image.composite(background, g, alpha), Image.composite(background, b, alpha)
# Merge the channels back into an image
image = Image.merge("RGB", image_rgb)
# Resize the image to the desired size
image = image.resize(new_size, Image.LANCZOS)
# Save the processed image
# image.save(output_path)
return image
def process_image(image_path, output_dir, no_remove_bg, foreground_ratio):
timer.start("Processing image")
if no_remove_bg:
rembg_session = None
image = np.array(Image.open(image_path).convert("RGB"))
else:
image = remove_background(image_path ,output_dir)
# Save the processed image
os.makedirs(output_dir, exist_ok=True)
image.save(os.path.join(output_dir, "processed_input.png"))
timer.end("Processing image")
return image
def run_model(model, image, output_dir, device, render, mc_resolution, model_save_format, bake_texture_flag, texture_resolution):
logging.info("Running model...")
timer.start("Running model")
with torch.no_grad():
scene_codes = model([image], device=device)
timer.end("Running model")
out_video_path = None
if render:
timer.start("Rendering")
render_images = model.render(scene_codes, n_views=30, return_type="pil")
for ri, render_image in enumerate(render_images[0]):
render_image.save(os.path.join(output_dir, f"render_{ri:03d}.png"))
out_video_path = os.path.join(output_dir, "render.mp4")
save_video(
render_images[0], out_video_path, fps=30
)
timer.end("Rendering")
timer.start("Extracting mesh")
meshes = model.extract_mesh(scene_codes, not bake_texture_flag, resolution=mc_resolution)
timer.end("Extracting mesh")
out_mesh_path = os.path.join(output_dir, f"mesh.{model_save_format}")
if bake_texture_flag:
out_texture_path = os.path.join(output_dir, "texture.png")
timer.start("Baking texture")
bake_output = bake_texture(meshes[0], model, scene_codes[0], texture_resolution)
timer.end("Baking texture")
timer.start("Exporting mesh and texture")
xatlas.export(out_mesh_path, meshes[0].vertices[bake_output["vmapping"]], bake_output["indices"], bake_output["uvs"], meshes[0].vertex_normals[bake_output["vmapping"]])
Image.fromarray((bake_output["colors"] * 255.0).astype(np.uint8)).transpose(Image.FLIP_TOP_BOTTOM).save(out_texture_path)
timer.end("Exporting mesh and texture")
else:
timer.start("Exporting mesh")
meshes[0].export(out_mesh_path)
timer.end("Exporting mesh")
return out_mesh_path ,out_video_path
logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO)
timer = Timer()
|