strKernel = globals()[strFunction]

while True:

objMatch = re.search('(SIZE_)([0-4])(\()([^\)]*)(\))', strKernel)

if objMatch is None:

break

# end

intArg = int(objMatch.group(2))

strTensor = objMatch.group(4)

intSizes = objVariables[strTensor].size()

strKernel = strKernel.replace(objMatch.group(), str(intSizes[intArg]))

# end

while True:

objMatch = re.search('(OFFSET_)([0-4])(\()([^\)]+)(\))', strKernel)

if objMatch is None:

break

# end

intArgs = int(objMatch.group(2))

strArgs = objMatch.group(4).split(',')

strTensor = strArgs[0]

intStrides = objVariables[strTensor].stride()

strIndex = [ '((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str(intStrides[intArg]) + ')' for intArg in range(intArgs) ]

strKernel = strKernel.replace(objMatch.group(0), '(' + str.join('+', strIndex) + ')')

# end

while True:

objMatch = re.search('(VALUE_)([0-4])(\()([^\)]+)(\))', strKernel)

if objMatch is None:

break

# end

intArgs = int(objMatch.group(2))

strArgs = objMatch.group(4).split(',')

strTensor = strArgs[0]

intStrides = objVariables[strTensor].stride()

strIndex = [ '((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str(intStrides[intArg]) + ')' for intArg in range(intArgs) ]

strKernel = strKernel.replace(objMatch.group(0), strTensor + '[' + str.join('+', strIndex) + ']')

# end

@staticmethod

def forward(self, input, flow):

self.save_for_backward(input, flow)

intSamples = input.shape[0]

intInputDepth, intInputHeight, intInputWidth = input.shape[1], input.shape[2], input.shape[3]

intFlowDepth, intFlowHeight, intFlowWidth = flow.shape[1], flow.shape[2], flow.shape[3]

assert(intFlowDepth == 2)

assert(intInputHeight == intFlowHeight)

assert(intInputWidth == intFlowWidth)

assert(input.is_contiguous() == True)

assert(flow.is_contiguous() == True)

output = input.new_zeros([ intSamples, intInputDepth, intInputHeight, intInputWidth ])

if input.is_cuda == True:

n = output.nelement()

cupy_launch('kernel_Softsplat_updateOutput', cupy_kernel('kernel_Softsplat_updateOutput', {

'input': input,

'flow': flow,

'output': output

}))(

grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),

block=tuple([ 512, 1, 1 ]),

args=[ n, input.data_ptr(), flow.data_ptr(), output.data_ptr() ]

)

elif input.is_cuda == False:

raise NotImplementedError()

# end

return output

# end

@staticmethod

def backward(self, gradOutput):

input, flow = self.saved_tensors

intSamples = input.shape[0]

intInputDepth, intInputHeight, intInputWidth = input.shape[1], input.shape[2], input.shape[3]

intFlowDepth, intFlowHeight, intFlowWidth = flow.shape[1], flow.shape[2], flow.shape[3]

assert(intFlowDepth == 2)

assert(intInputHeight == intFlowHeight)

assert(intInputWidth == intFlowWidth)

assert(gradOutput.is_contiguous() == True)

gradInput = input.new_zeros([ intSamples, intInputDepth, intInputHeight, intInputWidth ]) if self.needs_input_grad[0] == True else None

gradFlow = input.new_zeros([ intSamples, intFlowDepth, intFlowHeight, intFlowWidth ]) if self.needs_input_grad[1] == True else None

if input.is_cuda == True:

if gradInput is not None:

n = gradInput.nelement()

cupy_launch('kernel_Softsplat_updateGradInput', cupy_kernel('kernel_Softsplat_updateGradInput', {

'input': input,

'flow': flow,

'gradOutput': gradOutput,

'gradInput': gradInput,

'gradFlow': gradFlow

}))(

grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),

block=tuple([ 512, 1, 1 ]),

args=[ n, input.data_ptr(), flow.data_ptr(), gradOutput.data_ptr(), gradInput.data_ptr(), None ]

)

# end

if gradFlow is not None:

n = gradFlow.nelement()

cupy_launch('kernel_Softsplat_updateGradFlow', cupy_kernel('kernel_Softsplat_updateGradFlow', {

'input': input,

'flow': flow,

'gradOutput': gradOutput,

'gradInput': gradInput,

'gradFlow': gradFlow

}))(

grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),

block=tuple([ 512, 1, 1 ]),

args=[ n, input.data_ptr(), flow.data_ptr(), gradOutput.data_ptr(), None, gradFlow.data_ptr() ]

)

# end

elif input.is_cuda == False:

raise NotImplementedError()

# end

return gradInput, gradFlow

assert(tenMetric is None or tenMetric.shape[1] == 1)

assert(strType in ['summation', 'average', 'linear', 'softmax'])

if strType == 'average':

tenInput = torch.cat([ tenInput, tenInput.new_ones(tenInput.shape[0], 1, tenInput.shape[2], tenInput.shape[3]) ], 1)

elif strType == 'linear':

tenInput = torch.cat([ tenInput * tenMetric, tenMetric ], 1)

elif strType == 'softmax':

tenInput = torch.cat([ tenInput * tenMetric.exp(), tenMetric.exp() ], 1)

# end

tenOutput = _FunctionSoftsplat.apply(tenInput, tenFlow)

if strType != 'summation':

tenOutput = tenOutput[:, :-1, :, :] / (tenOutput[:, -1:, :, :] + 0.0000001)

# end

def __init__(self, strType):

super(ModuleSoftsplat, self).__init__()

self.strType = strType

# end

def forward(self, tenInput, tenFlow, tenMetric):

return FunctionSoftsplat(tenInput, tenFlow, tenMetric, self.strType)

def __init__(self, *args):

self.vec = [];

if len(args) == 0:

# default is a null rotation

self.s = 1.0

self.v = matrix([0.0, 0.0, 0.0])

elif len(args) == 4:

self.s = args[0];

self.v = mat(args[1:4])

else:

print("error")

return None

def __repr__(self):

return "%f <%f, %f, %f>" % (self.s, self.v[0, 0], self.v[0, 1], self.v[0, 2])

def tr2q(self, t):

# TR2Q Convert homogeneous transform to a unit-quaternion

#

# Q = tr2q(T)

#

# Return a unit quaternion corresponding to the rotational part of the

# homogeneous transform T.

qs = sqrt(trace(t) + 1) / 2.0

kx = t[2, 1] - t[1, 2] # Oz - Ay

ky = t[0, 2] - t[2, 0] # Ax - Nz

kz = t[1, 0] - t[0, 1] # Ny - Ox

if (t[0, 0] >= t[1, 1]) and (t[0, 0] >= t[2, 2]):

kx1 = t[0, 0] - t[1, 1] - t[2, 2] + 1 # Nx - Oy - Az + 1

ky1 = t[1, 0] + t[0, 1] # Ny + Ox

kz1 = t[2, 0] + t[0, 2] # Nz + Ax

add = (kx >= 0)

elif (t[1, 1] >= t[2, 2]):

kx1 = t[1, 0] + t[0, 1] # Ny + Ox

ky1 = t[1, 1] - t[0, 0] - t[2, 2] + 1 # Oy - Nx - Az + 1

kz1 = t[2, 1] + t[1, 2] # Oz + Ay

add = (ky >= 0)

else:

kx1 = t[2, 0] + t[0, 2] # Nz + Ax

ky1 = t[2, 1] + t[1, 2] # Oz + Ay

kz1 = t[2, 2] - t[0, 0] - t[1, 1] + 1 # Az - Nx - Oy + 1

add = (kz >= 0)

if add:

kx = kx + kx1

ky = ky + ky1

kz = kz + kz1

else:

kx = kx - kx1

ky = ky - ky1

kz = kz - kz1

kv = matrix([kx, ky, kz])

nm = linalg.norm(kv)

if nm == 0:

self.s = 1.0

self.v = matrix([0.0, 0.0, 0.0])

else:

self.s = qs

self.v = (sqrt(1 - qs ** 2) / nm) * kv

############### OPERATORS #########################################

# PLUS Add two quaternion objects

#

# Invoked by the + operator

#

# q1+q2 standard quaternion addition

def __add__(self, q):

'''

if isinstance(q, quaternion):

qr = quaternion()

qr.s = 0

qr.s = self.s + q.s

qr.v = self.v + q.v

return qr

else:

raise ValueError

# MINUS Subtract two quaternion objects

#

# Invoked by the - operator

#

# q1-q2 standard quaternion subtraction

def __sub__(self, q):

'''

if isinstance(q, quaternion):

qr = quaternion()

qr.s = 0

qr.s = self.s - q.s

qr.v = self.v - q.v

return qr

else:

raise ValueError

# q * q or q * const

def __mul__(self, q2):

'''

qr = quaternion();

if isinstance(q2, quaternion):

# Multiply unit-quaternion by unit-quaternion

#

# QQ = qqmul(Q1, Q2)

# decompose into scalar and vector components

s1 = self.s;

v1 = self.v

s2 = q2.s;

v2 = q2.v

# form the product

qr.s = s1 * s2 - v1 * v2.T

qr.v = s1 * v2 + s2 * v1 + cross(v1, v2)

elif type(q2) is matrix:

# Multiply vector by unit-quaternion

#

# Rotate the vector V by the unit-quaternion Q.

if q2.shape == (1, 3) or q2.shape == (3, 1):

qr = self * quaternion(q2) * self.inv()

return qr.v;

else:

raise ValueError;

else:

qr.s = self.s * q2

qr.v = self.v * q2

return qr

def __rmul__(self, c):

'''

qr = quaternion()

qr.s = self.s * c

qr.v = self.v * c

return qr

def __imul__(self, x):

'''

if isinstance(x, quaternion):

s1 = self.s;

v1 = self.v

s2 = x.s

v2 = x.v

# form the product

self.s = s1 * s2 - v1 * v2.T

self.v = s1 * v2 + s2 * v1 + cross(v1, v2)

elif isscalar(x):

self.s *= x;

self.v *= x;

return self;

def __pow__(self, p):

'''

# check that exponent is an integer

if not isinstance(p, int):

raise ValueError

qr = quaternion()

q = quaternion(self);

# multiply by itself so many times

for i in range(0, abs(p)):

qr *= q

# if exponent was negative, invert it

if p < 0:

qr = qr.inv()

return qr

def double(self):

"""Return the quaternion as 4-element vector.

return concatenate((mat(self.s), self.v), 1) # Debug present

def unit_Q(self):

'''

# Still have some errors with linalg.norm(self.v)

qr = quaternion()

try:

nm = linalg.norm(self.v) / sqrt(1 - pow(self.s, 2))

qr.s = self.s

qr.v = self.v / nm

except:

qr.s = self.s

qr.v = self.v

q1 = Q1.double()

q2 = Q2.double()

theta = acos(q1 * q2.T)

q = []

count = 0

if isscalar(r):

if r < 0 or r > 1:

raise Exception('R out of range')

if theta == 0:

q = quaternion(Q1)

else:

Qq = np.copy((sin((1 - r) * theta) * q1 + sin(r * theta) * q2) / sin(theta))

q = quaternion(Qq[0, 0], Qq[0, 1], Qq[0, 2], Qq[0, 3])

else:

for R in r:

if theta == 0:

qq = Q1

else:

qq = quaternion((sin((1 - R) * theta) * q1 + sin(R * theta) * q2) / sin(theta))

q.append(qq)

M = np.asarray(M)

if cy_thresh is None:

try:

cy_thresh = np.finfo(M.dtype).eps * 4

except ValueError:

cy_thresh = _FLOAT_EPS_4

r11, r12, r13, r21, r22, r23, r31, r32, r33 = M.flat

# cy: sqrt((cos(y)*cos(z))**2 + (cos(x)*cos(y))**2)

cy = math.sqrt(r33*r33 + r23*r23)

if cy > cy_thresh: # cos(y) not close to zero, standard form

z = math.atan2(-r12, r11) # atan2(cos(y)*sin(z), cos(y)*cos(z))

y = math.atan2(r13, cy) # atan2(sin(y), cy)

x = math.atan2(-r23, r33) # atan2(cos(y)*sin(x), cos(x)*cos(y))

else: # cos(y) (close to) zero, so x -> 0.0 (see above)

# so r21 -> sin(z), r22 -> cos(z) and

z = math.atan2(r21, r22)

y = math.atan2(r13, cy) # atan2(sin(y), cy)

x = 0.0

'''

V_AB = B -A

C = A + T*(V_AB)

return np.array([

[1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,

2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],

2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]],

[2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],

1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,

2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]],

[2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],

2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],

Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat

K = np.array([

[Rxx - Ryy - Rzz, 0, 0, 0],

[Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],

[Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],

[Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0

eigvals, eigvecs = np.linalg.eigh(K)

qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]

if qvec[0] < 0:

qvec *= -1

track = {'lin': [], 'rot': []}

q_a = rotmat2qvec(A_rot)

q_b = rotmat2qvec(B_rot)

QA = quaternion(q_a[0], q_a[1], q_a[2], q_a[3])

QB = quaternion(q_b[0], q_b[1], q_b[2], q_b[3])

track['lin'] = linear_translation(A_trans, B_trans, T).tolist()

q = interpolate(QA, QB, T)

track['rot'] = [q.s] + (q.v).tolist()[0] # List of quaternion [w x y z]

raw_rgba_s = torch.cat([ret['raw_rgb'], ret['raw_alpha'].unsqueeze(-1)], dim=-1)

raw_rgba = raw_rgba_s[:, :, j, :].permute(2, 0, 1).unsqueeze(0).contiguous().cuda()

pts_ref = ret['pts_ref'][:, :, j, :3]

pts_ref_e_G = NDC2Euclidean(pts_ref, H, W, focal)

if fwd_flow:

pts_post = pts_ref + ret['raw_sf_ref2post'][:, :, j, :]

else:

pts_post = pts_ref + ret['raw_sf_ref2prev'][:, :, j, :]

pts_post_e_G = NDC2Euclidean(pts_post, H, W, focal)

pts_mid_e_G = (pts_post_e_G - pts_ref_e_G) * ratio + pts_ref_e_G

pts_mid_e_local = se3_transform_points(pts_mid_e_G,

R_w2t.unsqueeze(0).unsqueeze(0),

t_w2t.unsqueeze(0).unsqueeze(0))

pts_2d_mid = perspective_projection(pts_mid_e_local, H, W, focal)

xx, yy = torch.meshgrid(torch.linspace(0, W - 1, W), torch.linspace(0, H - 1, H))

xx = xx.t()

yy = yy.t()

pts_2d_original = torch.stack([xx, yy], -1)

flow_2d = pts_2d_mid - pts_2d_original

flow_2d = flow_2d.permute(2, 0, 1).unsqueeze(0).contiguous().cuda()

splat_raw_rgba_dy = FunctionSoftsplat(tenInput=raw_rgba,

tenFlow=flow_2d,

tenMetric=None,

strType='average')

# splatting for static nerf

pts_rig_e_local = se3_transform_points(pts_ref_e_G,

R_w2t.unsqueeze(0).unsqueeze(0),

t_w2t.unsqueeze(0).unsqueeze(0))

pts_2d_rig = perspective_projection(pts_rig_e_local, H, W, focal)

flow_2d_rig = pts_2d_rig - pts_2d_original

flow_2d_rig = flow_2d_rig.permute(2, 0, 1).unsqueeze(0).contiguous().cuda()

raw_rgba_rig = torch.cat([ret['raw_rgb_rigid'], ret['raw_alpha_rigid'].unsqueeze(-1)], dim=-1)

raw_rgba_rig = raw_rgba_rig[:, :, j, :].permute(2, 0, 1).unsqueeze(0).contiguous().cuda()

splat_raw_rgba_rig = FunctionSoftsplat(tenInput=raw_rgba_rig,

tenFlow=flow_2d_rig,

tenMetric=None,

strType='average')

splat_alpha_dy = splat_raw_rgba_dy[0, 3:4, :, :]

splat_rgb_dy = splat_raw_rgba_dy[0, 0:3, :, :]

splat_alpha_rig = splat_raw_rgba_rig[0, 3:4, :, :]

splat_rgb_rig = splat_raw_rgba_rig[0, 0:3, :, :]

hwf, chunk, render_kwargs,

gt_imgs=None, savedir=None,

render_factor=0, target_idx=10):

# import scipy.io

start_frame = target_idx - 6.

end_frame = target_idx + 6.

H, W, focal = hwf

if render_factor != 0:

# Render downsampled for speed

H = H // render_factor

W = W // render_factor

focal = focal / render_factor

t = time.time()

count = 0

save_img_dir = os.path.join(savedir, 'images')

os.makedirs(save_img_dir, exist_ok=True)

print("##########", target_idx, start_frame, end_frame)

for i, cur_time in enumerate(np.linspace(start_frame, end_frame, 200 + 1).tolist()):

flow_time = int(np.floor(cur_time))

ratio = cur_time - np.floor(cur_time)

print('cur_time ', i, cur_time, ratio)

t = time.time()

int_rot, int_trans = linear_pose_interp(render_poses[flow_time, :3, 3],

render_poses[flow_time, :3, :3],

render_poses[flow_time + 1, :3, 3],

render_poses[flow_time + 1, :3, :3],

ratio)

int_poses = np.concatenate((int_rot, int_trans[:, np.newaxis]), 1)

int_poses = np.concatenate([int_poses[:3, :4], np.array([0.0, 0.0, 0.0, 1.0])[np.newaxis, :]], axis=0)

int_poses = np.dot(int_poses, bt_poses[i])

render_pose = torch.Tensor(int_poses).to(device)

R_w2t = render_pose[:3, :3].transpose(0, 1)

t_w2t = -torch.matmul(R_w2t, render_pose[:3, 3:4])

num_img = gt_imgs.shape[0]

img_idx_embed_1 = (np.floor(cur_time)) / float(num_img) * 2. - 1.0

img_idx_embed_2 = (np.floor(cur_time) + 1) / float(num_img) * 2. - 1.0

print('img_idx_embed_1 ', cur_time, img_idx_embed_1)

ret1 = render_sm(img_idx_embed_1, 0,

num_img,

H, W, focal,

chunk=1024 * 16,

c2w=render_pose,

**render_kwargs)

ret2 = render_sm(img_idx_embed_2, 0,

num_img,

H, W, focal,

chunk=1024 * 16,

c2w=render_pose,

**render_kwargs)

T_i = torch.ones((1, H, W))

final_rgb = torch.zeros((3, H, W))

num_sample = ret1['raw_rgb'].shape[2]

# final_depth = torch.zeros((1, H, W))

z_vals = ret1['z_vals']

for j in range(0, num_sample):

splat_alpha_dy_1, splat_rgb_dy_1, \

splat_alpha_rig_1, splat_rgb_rig_1 = splat_rgb_img(ret1, ratio, R_w2t, t_w2t,

j, H, W, focal, True)

splat_alpha_dy_2, splat_rgb_dy_2, \

splat_alpha_rig_2, splat_rgb_rig_2 = splat_rgb_img(ret2, 1. - ratio, R_w2t, t_w2t,

j, H, W, focal, False)

final_rgb += T_i * (splat_alpha_dy_1 * splat_rgb_dy_1 + \

splat_alpha_rig_1 * splat_rgb_rig_1) * (1.0 - ratio)

final_rgb += T_i * (splat_alpha_dy_2 * splat_rgb_dy_2 + \

splat_alpha_rig_2 * splat_rgb_rig_2) * ratio

# splat_alpha = splat_alpha1 * (1. - ratio) + splat_alpha2 * ratio

# final_rgb += T_i * (splat_alpha1 * (1. - ratio) * splat_rgb1 + splat_alpha2 * ratio * splat_rgb2)

alpha_1_final = (1.0 - (1. - splat_alpha_dy_1) * (1. - splat_alpha_rig_1)) * (1. - ratio)

alpha_2_fianl = (1.0 - (1. - splat_alpha_dy_2) * (1. - splat_alpha_rig_2)) * ratio

alpha_final = alpha_1_final + alpha_2_fianl

# final_depth += T_i * (alpha_final) * z_vals[..., j]

T_i = T_i * (1.0 - alpha_final + 1e-10)

filename = os.path.join(savedir, 'slow-mo_%03d.jpg' % (i))

rgb8 = to8b(final_rgb.permute(1, 2, 0).cpu().numpy())

new_x = rgb8.shape[0] // 2 * 2

new_y = rgb8.shape[1] // 2 * 2

rgb8 = rgb8[:new_x, :new_y, :]

# final_depth = torch.clamp(final_depth/percentile(final_depth, 98), 0., 1.)

# depth8 = to8b(final_depth.permute(1, 2, 0).repeat(1, 1, 3).cpu().numpy())

filename = os.path.join(save_img_dir, '{:03d}.jpg'.format(i))

hwf, chunk, render_kwargs,

gt_imgs=None, savedir=None,

render_factor=0,

target_idx=5):

H, W, focal = hwf

if render_factor != 0:

# Render downsampled for speed

H = H // render_factor

W = W // render_factor

focal = focal / render_factor

H, W, focal = hwf

if render_factor != 0:

# Render downsampled for speed

H = H // render_factor

W = W // render_factor

focal = focal / render_factor

t = time.time()

count = 0

for i, cur_time in enumerate(np.linspace(target_idx - 8., target_idx + 8., 160 + 1).tolist()):

ratio = cur_time - np.floor(cur_time)

render_pose = ref_c2w[:3, :4] # render_poses[i % num_frame_per_cycle][:3,:4]

R_w2t = render_pose[:3, :3].transpose(0, 1)

t_w2t = -torch.matmul(R_w2t, render_pose[:3, 3:4])

num_img = gt_imgs.shape[0]

img_idx_embed_1 = (np.floor(cur_time)) / float(num_img) * 2. - 1.0

img_idx_embed_2 = (np.floor(cur_time) + 1) / float(num_img) * 2. - 1.0

print('render lock camera time ', i, cur_time, ratio, time.time() - t)

t = time.time()

ret1 = render_sm(img_idx_embed_1, 0,

num_img,

H, W, focal,

chunk=1024 * 16,

c2w=render_pose,

**render_kwargs)

ret2 = render_sm(img_idx_embed_2, 0,

num_img,

H, W, focal,

chunk=1024 * 16,

c2w=render_pose,

**render_kwargs)

T_i = torch.ones((1, H, W))

final_rgb = torch.zeros((3, H, W))

num_sample = ret1['raw_rgb'].shape[2]

for j in range(0, num_sample):

splat_alpha_dy_1, splat_rgb_dy_1, \

splat_alpha_rig_1, splat_rgb_rig_1 = splat_rgb_img(ret1, ratio, R_w2t,

t_w2t, j, H, W,

focal, True)

splat_alpha_dy_2, splat_rgb_dy_2, \

splat_alpha_rig_2, splat_rgb_rig_2 = splat_rgb_img(ret2, 1. - ratio, R_w2t,

t_w2t, j, H, W,

focal, False)

final_rgb += T_i * (splat_alpha_dy_1 * splat_rgb_dy_1 + \

splat_alpha_rig_1 * splat_rgb_rig_1) * (1.0 - ratio)

final_rgb += T_i * (splat_alpha_dy_2 * splat_rgb_dy_2 + \

splat_alpha_rig_2 * splat_rgb_rig_2) * ratio

alpha_1_final = (1.0 - (1. - splat_alpha_dy_1) * (1. - splat_alpha_rig_1)) * (1. - ratio)

alpha_2_fianl = (1.0 - (1. - splat_alpha_dy_2) * (1. - splat_alpha_rig_2)) * ratio

alpha_final = alpha_1_final + alpha_2_fianl

T_i = T_i * (1.0 - alpha_final + 1e-10)

filename = os.path.join(savedir, '%03d.jpg' % (i))

rgb8 = to8b(final_rgb.permute(1, 2, 0).cpu().numpy())

new_x = rgb8.shape[0] // 2 * 2

new_y = rgb8.shape[1] // 2 * 2

rgb8 = rgb8[:new_x, :new_y, :]

#start_y = (rgb8.shape[1] - 512) // 2