185 lines
6.8 KiB
Python
185 lines
6.8 KiB
Python
import numpy as np
|
|
import torch
|
|
import time
|
|
import nvdiffrast.torch as dr
|
|
from PIL import Image
|
|
|
|
from ..crm.util.utils import get_tri
|
|
|
|
class CRMSampler:
|
|
@classmethod
|
|
def process_pixel_img(
|
|
cls,
|
|
pixel_img: Image.Image,
|
|
bg_color = "#7F7F7F"
|
|
):
|
|
# expand image to 1:1 squared and add fixed background color
|
|
width, height = pixel_img.size
|
|
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(pixel_img, paste_position)
|
|
|
|
# given an RGBA image fixed background color, alpha channel is used as mask to add background color
|
|
background = Image.new("RGBA", new_image.size, bg_color)
|
|
return Image.alpha_composite(background, new_image).convert("RGB")
|
|
|
|
@classmethod
|
|
def stage1_sample(
|
|
cls,
|
|
stage1_sampler,
|
|
pixel_img: Image.Image,
|
|
prompt="3D assets",
|
|
neg_texts="uniform low no texture ugly, boring, bad anatomy, blurry, pixelated, obscure, unnatural colors, poor lighting, dull, and unclear.",
|
|
seed=0,
|
|
scale=5,
|
|
step=50,
|
|
ddim_eta=0.0
|
|
):
|
|
"""
|
|
The first stage was condition on single pixel image, gererate multi-view pixel image, based on the v2pp config
|
|
"""
|
|
|
|
stage1_sampler.seed = seed
|
|
|
|
uc = stage1_sampler.model.get_learned_conditioning([neg_texts]).to(stage1_sampler.device)
|
|
stage1_images = stage1_sampler.i2i(
|
|
stage1_sampler.model,
|
|
stage1_sampler.size,
|
|
prompt,
|
|
uc=uc,
|
|
sampler=stage1_sampler.sampler,
|
|
ip=pixel_img,
|
|
step=step,
|
|
scale=scale,
|
|
batch_size=stage1_sampler.batch_size,
|
|
ddim_eta=ddim_eta,
|
|
dtype=stage1_sampler.dtype,
|
|
device=stage1_sampler.device,
|
|
camera=stage1_sampler.camera,
|
|
num_frames=stage1_sampler.num_frames,
|
|
pixel_control=(stage1_sampler.mode == "pixel"),
|
|
transform=stage1_sampler.image_transform,
|
|
offset_noise=stage1_sampler.offset_noise,
|
|
)
|
|
|
|
stage1_images = stage1_images[:-1] # remove reference view
|
|
return stage1_images # (N, H, W, 3) in [0, 1]
|
|
|
|
@classmethod
|
|
def stage2_sample(
|
|
cls,
|
|
stage2_sampler,
|
|
pixel_img: Image.Image,
|
|
stage1_images,
|
|
prompt="3D assets",
|
|
neg_texts="uniform low no texture ugly, boring, bad anatomy, blurry, pixelated, obscure, unnatural colors, poor lighting, dull, and unclear.",
|
|
seed=0,
|
|
scale=5,
|
|
step=50
|
|
):
|
|
"""
|
|
The second stage was condition on multiview pixel image generated by the first stage, generate the final image, based on the stage2-test config
|
|
"""
|
|
|
|
stage2_sampler.seed = seed
|
|
|
|
# Convert torch image to PIL.Image
|
|
stage1_images = list((255.0 * stage1_images.cpu().numpy()).astype(np.uint8))
|
|
stage1_images = [Image.fromarray(img) for img in stage1_images]
|
|
|
|
uc = stage2_sampler.model.get_learned_conditioning([neg_texts]).to(stage2_sampler.device)
|
|
stage2_images = stage2_sampler.i2iStage2(
|
|
stage2_sampler.model,
|
|
stage2_sampler.size,
|
|
prompt,
|
|
uc=uc,
|
|
sampler=stage2_sampler.sampler,
|
|
pixel_images=stage1_images,
|
|
ip=pixel_img,
|
|
step=step,
|
|
scale=scale,
|
|
batch_size=stage2_sampler.batch_size,
|
|
ddim_eta=0.0,
|
|
dtype=stage2_sampler.dtype,
|
|
device=stage2_sampler.device,
|
|
camera=stage2_sampler.camera,
|
|
num_frames=stage2_sampler.num_frames,
|
|
pixel_control=(stage2_sampler.mode == "pixel"),
|
|
transform=stage2_sampler.image_transform,
|
|
offset_noise=stage2_sampler.offset_noise,
|
|
)
|
|
|
|
return stage2_images # (N, H, W, 3) in [0, 1]
|
|
|
|
@classmethod
|
|
def generate3d(cls, crm_model, rgb, ccm, device):
|
|
|
|
color_tri = torch.from_numpy(rgb)
|
|
xyz_tri = torch.from_numpy(ccm[:,:,(2,1,0)])
|
|
color = color_tri.permute(2,0,1)
|
|
xyz = xyz_tri.permute(2,0,1)
|
|
|
|
|
|
def get_imgs(color):
|
|
# color : [C, H, W*6]
|
|
color_list = []
|
|
color_list.append(color[:,:,256*5:256*(1+5)])
|
|
for i in range(0,5):
|
|
color_list.append(color[:,:,256*i:256*(1+i)])
|
|
return torch.stack(color_list, dim=0)# [6, C, H, W]
|
|
|
|
triplane_color = get_imgs(color).permute(0,2,3,1).unsqueeze(0).to(device)# [1, 6, H, W, C]
|
|
|
|
color = get_imgs(color)
|
|
xyz = get_imgs(xyz)
|
|
|
|
color = get_tri(color, dim=0, blender= True, scale = 1).unsqueeze(0)
|
|
xyz = get_tri(xyz, dim=0, blender= True, scale = 1, fix= True).unsqueeze(0)
|
|
|
|
triplane = torch.cat([color,xyz],dim=1).to(device)
|
|
# 3D visualize
|
|
crm_model.eval()
|
|
glctx = dr.RasterizeCudaContext()
|
|
|
|
if crm_model.denoising == True:
|
|
tnew = 20
|
|
tnew = torch.randint(tnew, tnew+1, [triplane.shape[0]], dtype=torch.long, device=triplane.device)
|
|
noise_new = torch.randn_like(triplane) *0.5+0.5
|
|
triplane = crm_model.scheduler.add_noise(triplane, noise_new, tnew)
|
|
start_time = time.time()
|
|
with torch.no_grad():
|
|
triplane_feature2 = crm_model.unet2(triplane,tnew)
|
|
end_time = time.time()
|
|
elapsed_time = end_time - start_time
|
|
print(f"Running unet takes {elapsed_time}s")
|
|
else:
|
|
triplane_feature2 = crm_model.unet2(triplane)
|
|
|
|
|
|
with torch.no_grad():
|
|
data_config = {
|
|
'resolution': [1024, 1024],
|
|
"triview_color": triplane_color.to(device),
|
|
}
|
|
|
|
verts, faces = crm_model.decode(data_config, triplane_feature2)
|
|
|
|
data_config['verts'] = verts[0]
|
|
data_config['faces'] = faces
|
|
|
|
|
|
from kiui.mesh_utils import clean_mesh
|
|
verts, faces = clean_mesh(data_config['verts'].squeeze().cpu().numpy().astype(np.float32), data_config['faces'].squeeze().cpu().numpy().astype(np.int32), repair = False, remesh=False, remesh_size=0.005)
|
|
data_config['verts'] = torch.from_numpy(verts).cuda().contiguous()
|
|
data_config['faces'] = torch.from_numpy(faces).cuda().contiguous()
|
|
|
|
start_time = time.time()
|
|
|
|
mesh = crm_model.get_mesh_wt_uv(glctx, data_config, device, res=(1024,1024), tri_fea_2=triplane_feature2)
|
|
|
|
end_time = time.time()
|
|
elapsed_time = end_time - start_time
|
|
print(f"Extracting mesh & texture takes {elapsed_time}s")
|
|
|
|
return mesh |