| import torch |
| import torch.nn.functional as F |
| from typing import Literal, Optional, Tuple |
|
|
| @torch.no_grad() |
| def select_tokens( |
| obj_masks: torch.Tensor, |
| grid_thw: Tuple[int,int,int], |
| *, |
| patch_size: int = 14, |
| spatial_merge_size: int = 2, |
| temporal_patch_size: int = 2, |
| coverage_thresh: float = 0.7, |
| time_reduce: Literal["mean","max","all"] = "max", |
| device: str | torch.device = "cpu", |
| retry_step: float = 0.1, |
| retry_times: int = 1, |
| ensure_at_least_one: bool = True, |
| dtype: torch.dtype = torch.float32, |
| ): |
| if obj_masks.dim() == 3: |
| obj_masks = obj_masks.unsqueeze(0) |
| O, N, H_rz, W_rz = obj_masks.shape |
| T, H, W = grid_thw |
| m, g = spatial_merge_size, temporal_patch_size |
| if N != T*g: |
| if N < T * g: |
| pad = T*g - N |
| last = obj_masks[:,-1:,:,:].repeat(1, pad, 1, 1) |
| obj_masks = torch.cat([obj_masks, last], dim=1) |
| N = T * g |
| else: |
| obj_masks = obj_masks[:, :T * g, :, :] |
| N = T * g |
| Hm, Wm = H // m, W // m |
| pix_h, pix_w = m * patch_size, m * patch_size |
| assert H_rz % pix_h == 0 and W_rz % pix_w == 0, "resized // (28×28)" |
|
|
| M = obj_masks.to(device=device, dtype=dtype).clamp(0, 1) |
|
|
| M_flat = M.view(O*N, 1, H_rz, W_rz) |
| cov_hw = F.avg_pool2d(M_flat, kernel_size=(pix_h, pix_w), stride=(pix_h, pix_w)) |
| cov_hw = cov_hw.view(O, N, Hm, Wm) |
|
|
| cov_hw = cov_hw.view(O, T, g, Hm, Wm) |
| if time_reduce == "mean": |
| cov_thw = cov_hw.mean(dim=2) |
| elif time_reduce == "max": |
| cov_thw = cov_hw.max(dim=2).values |
| elif time_reduce == "all": |
| cov_thw = cov_hw.min(dim=2).values |
| else: |
| raise ValueError("time_reduce ∈ {'mean','max','all'}") |
|
|
| per_obj_idx = [] |
| per_t = Hm * Wm |
| for o in range(O): |
| nz = torch.empty(0, 3, dtype=torch.long, device=device) |
| tried = 0 |
| thr = coverage_thresh |
| while tried <= retry_times: |
| thr_eff = max(0.0, float(thr)) |
| sel = (cov_thw[o] >= thr_eff) |
| nz = torch.nonzero(sel, as_tuple=False) |
| if nz.numel() > 0: |
| break |
| tried += 1 |
| thr -= retry_step |
| if nz.numel() == 0: |
| if ensure_at_least_one: |
| flat = cov_thw[o].reshape(-1) |
| arg = torch.argmax(flat) |
| t = arg // (Hm * Wm) |
| rem = arg % (Hm * Wm) |
| hp = rem // Wm |
| wp = rem % Wm |
| idx = (t * per_t + hp * Wm + wp).view(1) |
| per_obj_idx.append(idx.to(device=device, dtype=torch.long)) |
| else: |
| per_obj_idx.append(torch.empty(0, dtype=torch.long, device=device)) |
| else: |
| t = nz[:, 0] |
| hp = nz[:, 1] |
| wp = nz[:, 2] |
| idx = t * per_t + hp * Wm + wp |
| per_obj_idx.append(idx.to(device=device, dtype=torch.long)) |
|
|
| if len(per_obj_idx) == 0: |
| union_idx = torch.empty(0, dtype=torch.long, device=device) |
| else: |
| union_idx = torch.unique(torch.cat(per_obj_idx, dim=0)) if per_obj_idx[0].numel() else torch.empty(0, dtype=torch.long, device=device) |
|
|
| union_idx_cpu = union_idx.cpu() |
| per_obj_idx_cpu = [idx.cpu() for idx in per_obj_idx] |
| cov_thw_cpu = cov_thw.cpu() |
|
|
| del M, M_flat, cov_hw, cov_thw, per_obj_idx, union_idx |
| if O > 0: |
| del sel, nz |
|
|
| return union_idx_cpu, per_obj_idx_cpu, cov_thw_cpu |
|
|
