|
import os |
|
import numpy as np |
|
import torch |
|
import torch.nn as nn |
|
|
|
from PIL import Image |
|
|
|
from my.utils import tqdm |
|
|
|
from pytorch3d.structures import Pointclouds |
|
from pytorch3d.renderer.cameras import PerspectiveCameras |
|
from pytorch3d.renderer import ( |
|
PointsRasterizer, |
|
AlphaCompositor, |
|
look_at_view_transform, |
|
) |
|
|
|
import torch.nn.functional as F |
|
|
|
from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config |
|
from point_e.diffusion.sampler import PointCloudSampler |
|
from point_e.models.download import load_checkpoint |
|
from point_e.models.configs import MODEL_CONFIGS, model_from_config |
|
|
|
|
|
class PointsRenderer(nn.Module): |
|
""" |
|
Modified version of Pytorch3D PointsRenderer |
|
""" |
|
|
|
def __init__(self, rasterizer, compositor) -> None: |
|
super().__init__() |
|
self.rasterizer = rasterizer |
|
self.compositor = compositor |
|
|
|
def to(self, device): |
|
|
|
|
|
self.rasterizer = self.rasterizer.to(device) |
|
self.compositor = self.compositor.to(device) |
|
return self |
|
|
|
def forward(self, point_clouds, **kwargs) -> torch.Tensor: |
|
fragments = self.rasterizer(point_clouds, **kwargs) |
|
|
|
|
|
|
|
depth_map = fragments[1][0,...,:1] |
|
|
|
|
|
|
|
|
|
r = self.rasterizer.raster_settings.radius |
|
|
|
dists2 = fragments.dists.permute(0, 3, 1, 2) |
|
weights = 1 - dists2 / (r * r) |
|
images = self.compositor( |
|
fragments.idx.long().permute(0, 3, 1, 2), |
|
weights, |
|
point_clouds.features_packed().permute(1, 0), |
|
**kwargs, |
|
) |
|
|
|
|
|
images = images.permute(0, 2, 3, 1) |
|
|
|
return images, depth_map |
|
|
|
|
|
def render_depth_from_cloud(points, angles, raster_settings, device,calibration_value=0): |
|
|
|
radius = 2.3 |
|
|
|
horizontal = angles[0]+calibration_value |
|
elevation = angles[1] |
|
FoV = angles[2] |
|
|
|
|
|
camera = py3d_camera(radius, elevation, horizontal, FoV, device) |
|
|
|
point_loc = torch.tensor(points.coords).to(device) |
|
colors = torch.tensor(np.stack([points.channels["R"], points.channels["G"], points.channels["B"]], axis=-1)).to(device) |
|
|
|
matching_rotation = torch.tensor([[[1.0, 0.0, 0.0], |
|
[0.0, 0.0, 1.0], |
|
[0.0, -1.0, 0.0]]]).to(device) |
|
|
|
rot_points = (matching_rotation @ point_loc[...,None]).squeeze() |
|
|
|
point_cloud = Pointclouds(points=[rot_points], features=[colors]) |
|
|
|
_, raw_depth_map = pointcloud_renderer(point_cloud, camera, raster_settings, device) |
|
|
|
disparity = camera.focal_length[0,0] / (raw_depth_map + 1e-9) |
|
|
|
max_disp = torch.max(disparity) |
|
min_disp = torch.min(disparity[disparity > 0]) |
|
|
|
norm_disparity = (disparity - min_disp) / (max_disp - min_disp) |
|
|
|
mask = norm_disparity > 0 |
|
norm_disparity = norm_disparity * mask |
|
|
|
depth_map = F.interpolate(norm_disparity.permute(2,0,1)[None,...],size=512,mode='bilinear')[0] |
|
depth_map = depth_map.repeat(3,1,1) |
|
|
|
return depth_map |
|
|
|
|
|
def py3d_camera(radius, elevation, horizontal, FoV, device, img_size=800): |
|
|
|
fov_rad = torch.deg2rad(torch.tensor(FoV)) |
|
focal = 1 / torch.tan(fov_rad / 2) * (2. / 2) |
|
|
|
focal_length = torch.tensor([[focal,focal]]).float() |
|
image_size = torch.tensor([[img_size,img_size]]).double() |
|
|
|
|
|
R, T = look_at_view_transform(dist=radius, elev=elevation, azim=horizontal, degrees=True) |
|
|
|
|
|
camera = PerspectiveCameras( |
|
R=R, |
|
T=T, |
|
focal_length=focal_length, |
|
image_size=image_size, |
|
device=device, |
|
) |
|
|
|
return camera |
|
|
|
def pointcloud_renderer(point_cloud, camera, raster_settings, device): |
|
|
|
camera = camera.to(device) |
|
|
|
rasterizer = PointsRasterizer(cameras=camera, raster_settings=raster_settings) |
|
renderer = PointsRenderer( |
|
rasterizer=rasterizer, |
|
compositor=AlphaCompositor() |
|
).to(device) |
|
|
|
image = renderer(point_cloud) |
|
|
|
return image |
|
|
|
def point_e(device,exp_dir): |
|
print('creating base model...') |
|
base_name = 'base1B' |
|
base_model = model_from_config(MODEL_CONFIGS[base_name], device) |
|
base_model.eval() |
|
base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name]) |
|
|
|
print('creating upsample model...') |
|
upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device) |
|
upsampler_model.eval() |
|
upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample']) |
|
|
|
print('downloading base checkpoint...') |
|
base_model.load_state_dict(load_checkpoint(base_name, device)) |
|
|
|
print('downloading upsampler checkpoint...') |
|
upsampler_model.load_state_dict(load_checkpoint('upsample', device)) |
|
|
|
sampler = PointCloudSampler( |
|
device=device, |
|
models=[base_model, upsampler_model], |
|
diffusions=[base_diffusion, upsampler_diffusion], |
|
num_points=[1024, 4096 - 1024], |
|
aux_channels=['R', 'G', 'B'], |
|
guidance_scale=[3.0, 3.0], |
|
) |
|
|
|
img = Image.open(os.path.join(exp_dir,'initial_image','instance0.png')) |
|
|
|
samples = None |
|
for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(images=[img]))): |
|
samples = x |
|
|
|
pc = sampler.output_to_point_clouds(samples)[0] |
|
|
|
return pc |
|
|
|
|
|
def point_e_gradio(img,device): |
|
print('creating base model...') |
|
base_name = 'base1B' |
|
base_model = model_from_config(MODEL_CONFIGS[base_name], device) |
|
base_model.eval() |
|
base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name]) |
|
|
|
print('creating upsample model...') |
|
upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device) |
|
upsampler_model.eval() |
|
upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample']) |
|
|
|
print('downloading base checkpoint...') |
|
base_model.load_state_dict(load_checkpoint(base_name, device)) |
|
|
|
print('downloading upsampler checkpoint...') |
|
upsampler_model.load_state_dict(load_checkpoint('upsample', device)) |
|
|
|
sampler = PointCloudSampler( |
|
device=device, |
|
models=[base_model, upsampler_model], |
|
diffusions=[base_diffusion, upsampler_diffusion], |
|
num_points=[1024, 4096 - 1024], |
|
aux_channels=['R', 'G', 'B'], |
|
guidance_scale=[3.0, 3.0], |
|
) |
|
|
|
|
|
samples = None |
|
for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(images=[img]))): |
|
samples = x |
|
|
|
pc = sampler.output_to_point_clouds(samples)[0] |
|
|
|
return pc |