| 182 | |
| 183 | |
| 184 | def gen_sineembed_for_position(pos_tensor): |
| 185 | # n_query, bs, _ = pos_tensor.size() |
| 186 | # sineembed_tensor = torch.zeros(n_query, bs, 256) |
| 187 | scale = 2 * math.pi |
| 188 | dim_t = torch.arange(128, dtype=torch.float32, device=pos_tensor.device) |
| 189 | dim_t = 10000**(2 * (dim_t // 2) / 128) |
| 190 | x_embed = pos_tensor[:, :, 0] * scale |
| 191 | y_embed = pos_tensor[:, :, 1] * scale |
| 192 | pos_x = x_embed[:, :, None] / dim_t |
| 193 | pos_y = y_embed[:, :, None] / dim_t |
| 194 | pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), |
| 195 | dim=3).flatten(2) |
| 196 | pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), |
| 197 | dim=3).flatten(2) |
| 198 | if pos_tensor.size(-1) == 2: |
| 199 | pos = torch.cat((pos_y, pos_x), dim=2) |
| 200 | elif pos_tensor.size(-1) == 4: |
| 201 | w_embed = pos_tensor[:, :, 2] * scale |
| 202 | pos_w = w_embed[:, :, None] / dim_t |
| 203 | pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), |
| 204 | dim=3).flatten(2) |
| 205 | |
| 206 | h_embed = pos_tensor[:, :, 3] * scale |
| 207 | pos_h = h_embed[:, :, None] / dim_t |
| 208 | pos_h = torch.stack((pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), |
| 209 | dim=3).flatten(2) |
| 210 | |
| 211 | pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2) |
| 212 | else: |
| 213 | raise ValueError('Unknown pos_tensor shape(-1):{}'.format( |
| 214 | pos_tensor.size(-1))) |
| 215 | return pos |
| 216 | |
| 217 | |
| 218 | def oks_overlaps(kpt_preds, kpt_gts, kpt_valids, kpt_areas, sigmas): |