Corrct Format for Query Points at Different Timesteps

[dongki-fai]

Hello, and thank you so much for this awesome work! 😊

I hope to ask a question regarding how to provide query points at different timesteps. In my dataset, some pedestrians are visible in the initial frame, but there are also some occluded pedestrian who are not visible initially and only appear at a later timestep. To track all pedestrians—including those visible in the initial frame and the one that appears later—I would like to supply appropriate query points using inference.py.

From what I observed, the query_point tensor [line] has a shape of (Batch size, 4). The first column seems to represent the timestep at which the query is provided (as described in the paper). I attempted to manually change this timestep from 0 to another value (e.g., 16), but the results were not too satisfactory. To make sure, I hope to ask whether there is any documentation or guidance on the correct way to specify query points at different timesteps to track all pedestrians?

Thank you so much!

[dadadadawjb]

Hi! I ran into the same issue and found a setup that works for me.

I think setting the first column in the query_point works. But if you are using inference.py via get_grid_queries() function to generate queries, the generated query_point is computed from the first frame by default. In addition to setting the first column in query_point to the frame index, you should also change the used depths, intrinsics and extrinsics to the query frame index.

Here is my code snippet:

def get_grid_queries(grid_size: int, depths: torch.Tensor, intrinsics: torch.Tensor, extrinsics: torch.Tensor, query_idx: int = 0):
    if len (depths.shape) == 3:
        return get_grid_queries(
            grid_size=grid_size,
            depths=depths.unsqueeze(0),
            intrinsics=intrinsics.unsqueeze(0),
            extrinsics=extrinsics.unsqueeze(0),
            query_idx=query_idx,
        ).squeeze(0)

    image_size = depths.shape[-2:]
    if grid_size == -1:
        xy = torch.stack(torch.meshgrid(
            torch.arange(image_size[1], dtype=torch.float32, device=intrinsics.device), 
            torch.arange(image_size[0], dtype=torch.float32, device=intrinsics.device), 
            indexing='xy'), dim=-1).view(1, -1, 2) # type: ignore
    else:
        xy = get_points_on_a_grid(grid_size, image_size).to(intrinsics.device) # type: ignore
    ji = torch.round(xy).to(torch.int32)
    d = depths[:, query_idx][torch.arange(depths.shape[0])[:, None], ji[..., 1], ji[..., 0]]

    assert d.shape[0] == 1, "batch size must be 1"
    mask = d[0] > 0
    d = d[:, mask]
    xy = xy[:, mask]
    ji = ji[:, mask]

    inv_intrinsics0 = torch.linalg.inv(intrinsics[0, query_idx])
    inv_extrinsics0 = torch.linalg.inv(extrinsics[0, query_idx])

    xy_homo = torch.cat([xy, torch.ones_like(xy[..., :1])], dim=-1)
    xy_homo = torch.einsum('ij,bnj->bni', inv_intrinsics0, xy_homo)
    local_coords = xy_homo * d[..., None]
    local_coords_homo = torch.cat([local_coords, torch.ones_like(local_coords[..., :1])], dim=-1)
    world_coords = torch.einsum('ij,bnj->bni', inv_extrinsics0, local_coords_homo)
    world_coords = world_coords[..., :3]

    queries = torch.cat([torch.ones_like(xy[:, :, :1]) * query_idx, world_coords], dim=-1).to(depths.device)  # type: ignore
    return queries