Function initialize_first_timestep

sem_gauss.py:143–178

(dataset, num_frames, scene_radius_depth_ratio, mean_sq_dist_method, seg_net)

        return torch.optim.Adam(param_groups, lr=0.0, eps=1e-15)

def initialize_first_timestep(dataset, num_frames, scene_radius_depth_ratio, mean_sq_dist_method, seg_net):
    # Get RGB-D Data & semantic & Camera Parameters
    color, depth, semantic, intrinsics, pose = dataset[0]


    # get semantic feature
    rgb_input = color.permute(2, 0, 1).unsqueeze(0).cuda()/255.0  # [1,3,H,W]
    seg_net.set_mode_get_feature()
    sem_feature = seg_net.cnn(rgb_input)    # [1,16,H,W]    #TODO: sem_feature和color都在gpu上(cuda:0)
    sem_feature = sem_feature.squeeze(0).detach() # [16,H,W]

    # Process RGB-D Data
    depth = depth.permute(2, 0, 1) # (H, W, C) -> (C, H, W)
    color = color.permute(2, 0, 1) / 255

    # Process Camera Parameters
    intrinsics = intrinsics[:3, :3]
    w2c = torch.linalg.inv(pose)

    # Setup Camera
    cam = setup_camera(color.shape[2], color.shape[1], intrinsics.cpu().numpy(), w2c.detach().cpu().numpy())

    # Get Initial Point Cloud (PyTorch CUDA Tensor)
    mask = (depth > 0) # Mask out invalid depth values
    mask = mask.reshape(-1)
    init_pt_cld, mean3_sq_dist = get_pointcloud(color, depth, sem_feature, intrinsics, w2c,
                                                mask=mask, compute_mean_sq_dist=True, 
                                                mean_sq_dist_method=mean_sq_dist_method)

    # Initialize Parameters
    params, variables = initialize_params(init_pt_cld, num_frames, mean3_sq_dist)

    # Initialize an estimate of scene radius for Gaussian-Splatting Densification
    variables['scene_radius'] = torch.max(depth)/scene_radius_depth_ratio

    return params, variables, intrinsics, w2c, cam

def get_loss(params, curr_data, variables, iter_time_idx, loss_weights, use_sil_for_loss,
             sil_thres, use_l1,ignore_outlier_depth_loss, seg_net, tracking=False,