Spaces:
Build error
Build error
File size: 4,908 Bytes
fd5e0f7 |
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 |
import cv2
import numpy as np
from copy import deepcopy
from ops.utils import dpt2xyz,transform_points
class Connect_Tool():
def __init__(self) -> None:
pass
def _align_scale_shift_numpy(self, pred: np.array, target: np.array):
mask = (target > 0) & (pred < 199)
target_mask = target[mask]
pred_mask = pred[mask]
if np.sum(mask) > 10:
scale, shift = np.polyfit(pred_mask, target_mask, deg=1)
if scale < 0:
scale = np.median(target[mask]) / (np.median(pred[mask]) + 1e-8)
shift = 0
else:
scale = 1
shift = 0
return scale,shift
def __call__(self, render_dpt, inpaint_dpt, inpaint_msk):
if np.sum(inpaint_msk > 0.5) < 1.: return render_dpt
# get areas need to be aligned
render_dpt_valid = render_dpt[~inpaint_msk]
inpaint_dpt_valid = inpaint_dpt[~inpaint_msk]
# rectify
scale,shift = self._align_scale_shift_numpy(inpaint_dpt_valid,render_dpt_valid)
inpaint_dpt = inpaint_dpt*scale + shift
return inpaint_dpt
class Smooth_Connect_Tool():
def __init__(self) -> None:
self.coarse_align = Connect_Tool()
def _coarse_alignment(self, render_dpt, ipaint_dpt, ipaint_msk):
# determine the scale and shift of inpaint_dpt to coarsely align it to render_dpt
inpaint_dpt = self.coarse_align(render_dpt,ipaint_dpt,ipaint_msk)
return inpaint_dpt
def _refine_movements(self, render_dpt, ipaint_dpt, ipaint_msk):
'''
Follow https://arxiv.org/pdf/2311.13384
'''
# Determine the adjustment of un-inpainted area
ipaint_msk = ipaint_msk>.5
H, W = ipaint_msk.shape[0:2]
U = np.arange(W)[None,:].repeat(H,axis=0)
V = np.arange(H)[:,None].repeat(W,axis=1)
# on kept areas
keep_render_dpt = render_dpt[~ipaint_msk]
keep_ipaint_dpt = ipaint_dpt[~ipaint_msk]
keep_adjust_dpt = keep_render_dpt - keep_ipaint_dpt
# iterative refinement
complete_adjust = np.zeros_like(ipaint_dpt)
for i in range(100):
complete_adjust[~ipaint_msk] = keep_adjust_dpt
complete_adjust = cv2.blur(complete_adjust,(15,15))
# complete_adjust[~ipaint_msk] = keep_adjust_dpt
ipaint_dpt = ipaint_dpt + complete_adjust
return ipaint_dpt
def _affine_dpt_to_GS(self, render_dpt, inpaint_dpt, inpaint_msk):
if np.sum(inpaint_msk > 0.5) < 1.: return render_dpt
inpaint_dpt = self._coarse_alignment(render_dpt,inpaint_dpt,inpaint_msk)
inpaint_dpt = self._refine_movements(render_dpt,inpaint_dpt,inpaint_msk)
return inpaint_dpt
def _scale_dpt_to_GS(self, render_dpt, inpaint_dpt, inpaint_msk):
if np.sum(inpaint_msk > 0.5) < 1.: return render_dpt
inpaint_dpt = self._refine_movements(render_dpt,inpaint_dpt,inpaint_msk)
return inpaint_dpt
class Occlusion_Removal():
def __init__(self) -> None:
pass
def __call__(self,scene,frame):
# first get xyz of the newly added frame
xyz = dpt2xyz(frame.dpt,frame.intrinsic)
# we only check newly added areas
xyz = xyz[frame.inpaint]
# move these xyzs to world coor system
inv_extrinsic = np.linalg.inv(frame.extrinsic)
xyz = transform_points(xyz,inv_extrinsic)
# we will add which pixels to the gaussian scene
msk = np.ones_like(xyz[...,0])
# project the xyzs to already built frames
for former_frame in scene.frames:
# xyz in camera frustrum
xyz_camera = transform_points(deepcopy(xyz),former_frame.extrinsic)
# uvz in camera frustrum
uvz_camera = np.einsum(f'ab,pb->pa',former_frame.intrinsic,xyz_camera)
# uv and d in camra frustrum
uv,d = uvz_camera[...,:2]/uvz_camera[...,-1:], uvz_camera[...,-1]
# in-frusturm pixels
valid_msk = (uv[...,0]>0) & (uv[...,0]<former_frame.W) & (uv[...,1]>0) & (uv[...,1]<former_frame.H) & (d>1e-2)
valid_idx = np.where(valid_msk)[0]
uv,d = uv[valid_idx].astype(np.uint32),d[valid_idx]
# make comparsion: compare_d < d is ok -- compare_d - d < 0(or a small number)
compare_d = former_frame.dpt[uv[:,1],uv[:,0]]
remove_msk = (compare_d-d)>(d+compare_d)/2./15.
# else to unvalid pixels
invalid_idx = valid_idx[remove_msk]
msk[invalid_idx] = 0.
# USE indexes rather than [][]
inpaint_idx_v,inpaint_idx_u = np.where(frame.inpaint)
inpaint_idx_v = inpaint_idx_v[msk<.5]
inpaint_idx_u = inpaint_idx_u[msk<.5]
frame.inpaint[inpaint_idx_v,inpaint_idx_u] = False
return frame |