memef4rmer/edit_anything

1import os2import numpy as np3import torch4import torch.nn as nn5 6from PIL import Image7 8from my.utils import tqdm9 10from pytorch3d.structures import Pointclouds11from pytorch3d.renderer.cameras import PerspectiveCameras   12from pytorch3d.renderer import (13    PointsRasterizer,14    AlphaCompositor,15    look_at_view_transform,16)17 18import torch.nn.functional as F19 20from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config21from point_e.diffusion.sampler import PointCloudSampler22from point_e.models.download import load_checkpoint23from point_e.models.configs import MODEL_CONFIGS, model_from_config24 25 26class PointsRenderer(nn.Module):27    """28    Modified version of Pytorch3D PointsRenderer29    """30 31    def __init__(self, rasterizer, compositor) -> None:32        super().__init__()33        self.rasterizer = rasterizer34        self.compositor = compositor35 36    def to(self, device):37        # Manually move to device rasterizer as the cameras38        # within the class are not of type nn.Module39        self.rasterizer = self.rasterizer.to(device)40        self.compositor = self.compositor.to(device)41        return self42 43    def forward(self, point_clouds, **kwargs) -> torch.Tensor:44        fragments = self.rasterizer(point_clouds, **kwargs)45        46        # import pdb; pdb.set_trace()47        48        depth_map = fragments[1][0,...,:1]49        50        # Construct weights based on the distance of a point to the true point.51        # However, this could be done differently: e.g. predicted as opposed52        # to a function of the weights.53        r = self.rasterizer.raster_settings.radius54 55        dists2 = fragments.dists.permute(0, 3, 1, 2)56        weights = 1 - dists2 / (r * r)57        images = self.compositor(58            fragments.idx.long().permute(0, 3, 1, 2),59            weights,60            point_clouds.features_packed().permute(1, 0),61            **kwargs,62        )63 64        # permute so image comes at the end65        images = images.permute(0, 2, 3, 1)66 67        return images, depth_map68 69 70def render_depth_from_cloud(points, angles, raster_settings, device,calibration_value=0):71    72    radius = 2.373    74    horizontal = angles[0]+calibration_value75    elevation = angles[1]76    FoV = angles[2]77 78 79    camera = py3d_camera(radius, elevation, horizontal, FoV, device)80    81    point_loc = torch.tensor(points.coords).to(device)82    colors = torch.tensor(np.stack([points.channels["R"], points.channels["G"], points.channels["B"]], axis=-1)).to(device)83 84    matching_rotation = torch.tensor([[[1.0, 0.0, 0.0],85                                    [0.0, 0.0, 1.0],86                                    [0.0, -1.0, 0.0]]]).to(device)87    88    rot_points = (matching_rotation @ point_loc[...,None]).squeeze()      89    90    point_cloud = Pointclouds(points=[rot_points], features=[colors])91    92    _, raw_depth_map = pointcloud_renderer(point_cloud, camera, raster_settings, device)    93 94    disparity = camera.focal_length[0,0] / (raw_depth_map + 1e-9)95    96    max_disp = torch.max(disparity) 97    min_disp = torch.min(disparity[disparity > 0])98    99    norm_disparity = (disparity - min_disp) / (max_disp - min_disp)100    101    mask = norm_disparity > 0102    norm_disparity = norm_disparity * mask103    104    depth_map = F.interpolate(norm_disparity.permute(2,0,1)[None,...],size=512,mode='bilinear')[0]105    depth_map = depth_map.repeat(3,1,1)106    107    return depth_map108 109 110def py3d_camera(radius, elevation, horizontal, FoV, device, img_size=800):111    112    fov_rad = torch.deg2rad(torch.tensor(FoV))113    focal = 1 / torch.tan(fov_rad / 2) * (2. / 2)114    115    focal_length = torch.tensor([[focal,focal]]).float()116    image_size = torch.tensor([[img_size,img_size]]).double()117 118    119    R, T = look_at_view_transform(dist=radius, elev=elevation, azim=horizontal, degrees=True)120 121    122    camera = PerspectiveCameras(123            R=R,124            T=T,125            focal_length=focal_length,126            image_size=image_size,127            device=device,128        )129 130    return camera131 132def pointcloud_renderer(point_cloud, camera, raster_settings, device):133 134    camera = camera.to(device)135 136    rasterizer = PointsRasterizer(cameras=camera, raster_settings=raster_settings)137    renderer = PointsRenderer(138        rasterizer=rasterizer,139        compositor=AlphaCompositor()140    ).to(device)141    142    image = renderer(point_cloud)143    144    return image145    146def point_e(device,exp_dir):147    print('creating base model...')148    base_name = 'base1B' # use base300M or base1B for better results149    base_model = model_from_config(MODEL_CONFIGS[base_name], device)150    base_model.eval()151    base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name])152 153    print('creating upsample model...')154    upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device)155    upsampler_model.eval()156    upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample'])157 158    print('downloading base checkpoint...')159    base_model.load_state_dict(load_checkpoint(base_name, device))160 161    print('downloading upsampler checkpoint...')162    upsampler_model.load_state_dict(load_checkpoint('upsample', device))163 164    sampler = PointCloudSampler(165        device=device,166        models=[base_model, upsampler_model],167        diffusions=[base_diffusion, upsampler_diffusion],168        num_points=[1024, 4096 - 1024],169        aux_channels=['R', 'G', 'B'],170        guidance_scale=[3.0, 3.0],171    )172 173    img = Image.open(os.path.join(exp_dir,'initial_image','instance0.png'))174 175    samples = None176    for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(images=[img]))):177        samples = x178        179    pc = sampler.output_to_point_clouds(samples)[0]180    181    return pc182 183 184def point_e_gradio(img,device):185    print('creating base model...')186    base_name = 'base1B' # use base300M or base1B for better results187    base_model = model_from_config(MODEL_CONFIGS[base_name], device)188    base_model.eval()189    base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name])190 191    print('creating upsample model...')192    upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device)193    upsampler_model.eval()194    upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample'])195 196    print('downloading base checkpoint...')197    base_model.load_state_dict(load_checkpoint(base_name, device))198 199    print('downloading upsampler checkpoint...')200    upsampler_model.load_state_dict(load_checkpoint('upsample', device))201 202    sampler = PointCloudSampler(203        device=device,204        models=[base_model, upsampler_model],205        diffusions=[base_diffusion, upsampler_diffusion],206        num_points=[1024, 4096 - 1024],207        aux_channels=['R', 'G', 'B'],208        guidance_scale=[3.0, 3.0],209    )210 211 212    samples = None213    for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(images=[img]))):214        samples = x215        216    pc = sampler.output_to_point_clouds(samples)[0]217    218    return pc