<jupyter_start><jupyter_text>Train<jupyter_code>D = RayDataset(target_ims, rays, device)
train_loader = torch.utils.data.DataLoader(D, batch_size=5000, shuffle=True)
def train(epoch, opt_color, N_points):
    losses=[]
    for batch_idx, (rays, pixels) in enumerate(train_loader):
        rays, pixels = (rays[0].to(device),rays[1].to(device)), pixels.to(device)
        opt_color.zero_grad()
        
        pix_estims = VG.render_rays(rays, (N_points))
        loss = ((pix_estims-pixels)**2).sum()/rays[0].shape[0]
            
        #loss = ((pix_estims-pixels)**2).sum()/rays[0].shape[0] +0.0001*VG.total_variation()
        #+ 0.1*opacity_loss(VG.opacities)
        
        loss.backward()
        losses.append(loss.item())
        #VG.update_grads(1000)
        opt_color.step()
        VG.clamp()
        

        if batch_idx%10==0:
            print("Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( epoch,batch_idx,len(train_loader),100.0 * batch_idx / len(train_loader),loss.data.item(),))
    return losses
VG = VoxelGridSpherical(4)
VG.load('mic_carve.obj')
losses=[]
epochss = [2, 4, 8, 8]
N_pointss = [2**(i+4) for i in range(4)]
lrs = [1000/2**i for i in range(4)]
lrs2 = [30 for i in range(4)]

epochss = [4]
N_pointss = [64]
lrs = [1000]



for epochs, N_points, lr in zip(epochss,N_pointss, lrs):
    #VG.subdivide()
    opt_c= torch.optim.SGD(
                [VG.colors], 
                lr=lr/10
            )
    opt_o= torch.optim.SGD(
                [VG.opacities], 
                lr=lr
            )
   
    for epoch in tqdm(range(epochs)):
        #VG.save("some.obj")
        plt.imshow(VG.render_large_image_from_rays(disp_rays[0],(500,1.2)))
        plt.show()
        losses += train(epoch, opt_c, N_points)
        plt.plot(np.log(rolling_average(np.array(losses))))
        plt.show()


    
compute_psnr(VG, disp_rays_test, disp_ims_test, N_points=1000)<jupyter_output>100%|█████████████████████████████████████████| 200/200 [00:05<00:00, 39.88it/s]<jupyter_text>View images<jupyter_code>VG = VoxelGrid()

model = "drums"
VG.load(model+'.obj')


dataset= "../nerf_synthetic/"+model
test_focal, test_c2w, test_gt = get_data(dataset, "test")
red = 4
disp_ims_test, disp_rays_test = reduce_data(test_c2w, test_gt,test_focal, red)

image_ind=0


plt.imshow(disp_ims_test[image_ind])
plt.show()

new_im = VG.render_large_image_from_rays(disp_rays_test[image_ind],(32,1.2))

plt.imshow(new_im)
plt.show()<jupyter_output><empty_output><jupyter_text>ROTATION<jupyter_code>cust_c2ws = create_rotation_matrices(1.5,-20, n=60)
red_fac=2
ordir_rays=[]
for c2w in cust_c2ws:
    ray_np = get_rays_np(800,800, focal, c2w)
    oris = ray_np[0][::red_fac,::red_fac]
    direct = ray_np[1][::red_fac,::red_fac] # direction. optimal fac:3
    ordir_rays.append((oris, direct))
VG=VoxelGrid()
VG.load('lego.obj')
VG.subdivide()
VG.treshold_opacity(0.5)
imgs=[]
for image_ind in tqdm(range(len(cust_c2ws))):
    with torch.no_grad():
        new_im = VG.render_large_image_from_rays(ordir_rays[image_ind],(500, 1.2))
        plt.imshow(new_im)
        plt.show()
        break
        #imgs.append(np.uint8(255*new_im))
imageio.mimwrite('screenshots/ficus_best2.gif', imgs,  format='GIF', duration=0.04)
VG.save_pointcloud()
rays_or_dir = [get_rays_np(800,800, focal, c2w) for c2w in all_c2w]
camera_centers = get_cameras_centers(rays_or_dir)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot( camera_centers[:,0], camera_centers[:, 1], camera_centers[:,2], 'o')

n = 40
t = np.linspace(0,2*np.pi, n+1)[:-1]
cust_centers = np.zeros((n,3))
cust_centers[:,0] = np.cos(t)*3.8
cust_centers[:,1] = np.sin(t)*3.8
cust_centers[:,2] = 1

cust_c2ws = [create_rotation_transformation_matrix(cust_centers[i], t[i], -np.pi*0.1) for i in range(n)]

ax.plot( cust_centers[:,0], cust_centers[:,1], cust_centers[:,2], 'o')

print(cust_centers[0])


L=[]
k=0
i=0
while i<39:
    k+=1
    for _ in range(4):
        L.append(i)
        i+=k
L
L = range(15)

images=[]
for ind in tqdm([0,1,3,7,15,40]):
    VG.load('16a_'+str(ind)+'.obj')
    for image_ind in (range(len(cust_c2ws))):
        new_im = VG.render_large_image_from_rays(ordir_rays[image_ind], (500, 1.2))
        img = Image.fromarray(np.uint8(255*np.clip(new_im,0,1)))
        draw = ImageDraw.Draw(img)
        draw.text((10, 10),"Epoch:"+str(ind),(0,0,0))
        images.append(np.array(img))
        #plt.imshow(images[-1])
        #plt.show()
        #images.append(np.uint8(255*new_im))
        #plt.imsave('screenshots/rotate/'+str(ind)+'_'+str(image_ind)+'.png', new_im)

images2=[]
k=0
while k<len(images):
    for _ in range(30):
        images2.append(images[k])
        k+=1
    k+=60
plt.imshow(images2[61])
imageio.mimwrite('exports/movies_training_full.gif', images,  format='GIF', duration=0.04)<jupyter_output><empty_output><jupyter_text>INTERP<jupyter_code>offsets = np.array([[0,0,0], #base
                           [0,0,0], 
                           [1,0,0], 
                           [0,0,1],
                           [1,0,1],
                           [0,1,0],
                           [1,1,0],  
                           [0,1,1],
                           [1,1,1]])


class VoxelGridInterp(VoxelGrid):
    def in_bounds_indices(self, p):
        ''' input: Nx3 array
            output: index of valid arrays'''
        in_x = (p[...,0]>=0)*(p[...,0]<(self.size-1))
        in_y = (p[...,1]>=0)*(p[...,1]<(self.size-1))
        in_z = (p[...,2]>=0)*(p[...,2]<(self.size-1))
        return in_x*in_y*in_z
   
    def render_rays(self, ordir_tuple, N_points, inv_depth=1.2):
        with torch.no_grad():
            ori = ordir_tuple[0][:, None,:]
            # WARNING: Assuming constant distance
            distances = 10*torch.sqrt( (ordir_tuple[1]**2).sum(1, keepdim=True))/inv_depth/(N_points-1)
            scatter_points = torch.rand_like(distances)*distances + torch.linspace(0,10, N_points, device=device)[None, :]/inv_depth
            p = ori + scatter_points[:,:,None]*(ordir_tuple[1][:, None, :])

            offsets = torch.tensor([[0,0,0], 
                               [1,0,0], 
                               [0,0,1],
                               [1,0,1],
                               [0,1,0],
                               [1,1,0],  
                               [0,1,1],
                               [1,1,1]], device=device)
            offsets[:, 1] *= self.size
            offsets[:, 2] *= self.size**2
            offsets = offsets.sum(1)

        
            #indicies
            p = self.descartes_to_indices(p)   
            p0 = torch.floor(p)
            in_bounds = self.in_bounds_indices(p)

            mesh_coords = self.flatten_3d_indices(p0.long())
            mesh_coords[torch.logical_not(in_bounds)] = 0

            diff = (p-p0)
            ind_offsets = []

            mcs = []
            for i in range(8):
                mcs.append( mesh_coords + offsets[i])                
        # color interp
        ciii = [self.colors[mcs[i]]
                *(1-diff[...,0, None])+
                self.colors[mcs[i+1]]
                *diff[...,0, None] for i in range(0,8,2)]
        cii = [ciii[i]*(1-diff[...,1, None]) + ciii[i+2]*diff[...,1, None] for i in [0,1]]
        colors = cii[0]*(1-diff[...,2, None])+cii[1]*diff[...,2, None]
        
        # opacities interp
        oiii = [self.opacities[mcs[i]]
               *(1-diff[...,0]) +
                self.opacities[mcs[i+1]]
                *diff[...,0] for i in range(0,8,2)]
        oii = [oiii[i]*(1-diff[...,1]) + oiii[i+2]*diff[...,1] for i in [0,1]]
        opacities = (oii[0]*(1-diff[...,2])+oii[1]*diff[...,2])*distances*in_bounds.float() 

        cumsum_opacities = torch.cumsum(opacities, 1)

        transp_term = torch.exp(-cumsum_opacities)*(1-torch.exp(-opacities))
        return (colors*transp_term[..., None]).sum(1) + torch.exp(-cumsum_opacities[:, -1])[..., None]
    


VG = VoxelGrid(32,1)


with torch.no_grad():
    VG.opacities[:] = 2
    VG.colors[:] = 0
    VG.colors[:, 1] = 0.5
    for i in range(VG.size):
        for j in range(0, VG.size,2):
            for k in range(VG.size):
    
                VG.colors[i+VG.size*j+VG.size**2*k, 1] = 0
                VG.colors[i+VG.size*j+VG.size**2*k, 2] = 0.5
       
     

from time import time

#VG= VoxelGridInterp(16,1.4)

im1=VG.render_large_image_from_rays(disp_rays[10],(1000,1.2))
VG.subdivide()
im2=VG.render_large_image_from_rays(disp_rays[10],(1000,1.2))
plt.imshow(im1)
plt.show()
plt.imshow(im2)
plt.show()


plt.imshow(((im1-im2)**2).sum(2))
plt.colorbar()
plt.show()
#VG.alive<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1848127600554785, 'fraction_numerical': 0.04623208506703652, 'mean_word_length': 2.4723113964687, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '604092', 'license': 'MIT', 'n_tokens_mistral': 3707, 'n_tokens_neox': 3485, 'n_words': 494, 'path': 'notebooks/test-cuda-batch.ipynb', 'repo_name': 'nissmar/Radiance-Fields', 'size': '8652'}