Untitled

def parse_isect_ids(isect_ids_np: np.ndarray, tile_n_bits: int = 22, image_n_bits: int = 10):
    """
    Parse 64-bit intersection keys into image_ids, tile_ids, and float depths.
    Layout: [image_id (image_n_bits)] | [tile_id (tile_n_bits)] | [depth (32 bits)]
    """
    depths_f, image_ids, tile_ids = [], [], []
    for val in isect_ids_np:
    
        image_id = (val >> (32 + tile_n_bits)) & ((1 << image_n_bits) - 1)
        tile_id = (val >> 32) & ((1 << tile_n_bits) - 1)
        # Extract the lower 32 bits, interpret as raw bytes, then as float32 using numpy
        depths = np.frombuffer(np.uint32(val & 0xFFFFFFFF).tobytes(), dtype=np.float32)[0]
        depths_f.append(depths)
        image_ids.append(image_id)
        tile_ids.append(tile_id)
    return depths_f, image_ids, tile_ids

def compute_ray_loss(info):
    radius = info["radii"]

    # image id (10 bits) | tile id (22 bits) | depth (32 bits)
    isect_ids = info["isect_ids"]
    flatten_ids = info["flatten_ids"]
    isect_ids_np = isect_ids.cpu().numpy().astype(np.uint64)

    depths_f, image_ids, tile_ids = parse_isect_ids(isect_ids_np, tile_n_bits=22, image_n_bits=10)
    offsets = {}

    for i, tile in enumerate(tile_ids):
        if tile not in offsets:
            offsets[tile] = i

    gaussains_per_tile = {}
    for tile, start_idx in offsets.items():
        while start_idx < len(tile_ids) and tile_ids[start_idx] == tile:
            if tile not in gaussains_per_tile:
                gaussains_per_tile[tile] = []            
            gaussains_per_tile[tile].append(
                (depths_f[start_idx], info["opacities"].cpu().numpy()[0][flatten_ids[start_idx]])
            )
            start_idx += 1


    # Squeeze out the first (batch) dimension and mask out zero radii
    radius_arr = np.array(radius.cpu().numpy()).squeeze(0)
    mask = np.logical_and(radius_arr[:, 0] != 0, radius_arr[:, 1] != 0)
    radius_arr = radius_arr[mask]


    loss = torch.tensor(0.0, device="cuda")
    for tile_idx_1d, gaussains in gaussains_per_tile.items():
        depths = torch.tensor([x[0] for x in gaussains])

        peak_positions = [depths[0], depths[-1]]
        opacities = torch.tensor([x[1] for x in gaussains])
        
        for p in peak_positions:
            weight = torch.exp(-0.5 * ((depths - p) / 0.1) ** 2)
            weight = weight / torch.sum(weight)  # normalize weight to sum=1
            loss += torch.sum(opacities * weight)

    return loss