A piecewise linear function y = f(x), using xp and yp as keypoints. We implement f(x) in a differentiable way (i.e. applicable for autograd). The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
(x, xp, yp)
| 1102 | ############################################################# |
| 1103 | |
| 1104 | def interpolate_fn(x, xp, yp): |
| 1105 | """ |
| 1106 | A piecewise linear function y = f(x), using xp and yp as keypoints. |
| 1107 | We implement f(x) in a differentiable way (i.e. applicable for autograd). |
| 1108 | The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.) |
| 1109 | Args: |
| 1110 | x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver). |
| 1111 | xp: PyTorch tensor with shape [C, K], where K is the number of keypoints. |
| 1112 | yp: PyTorch tensor with shape [C, K]. |
| 1113 | Returns: |
| 1114 | The function values f(x), with shape [N, C]. |
| 1115 | """ |
| 1116 | N, K = x.shape[0], xp.shape[1] |
| 1117 | all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2) |
| 1118 | sorted_all_x, x_indices = torch.sort(all_x, dim=2) |
| 1119 | x_idx = torch.argmin(x_indices, dim=2) |
| 1120 | cand_start_idx = x_idx - 1 |
| 1121 | start_idx = torch.where( |
| 1122 | torch.eq(x_idx, 0), |
| 1123 | torch.tensor(1, device=x.device), |
| 1124 | torch.where( |
| 1125 | torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx, |
| 1126 | ), |
| 1127 | ) |
| 1128 | end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1) |
| 1129 | start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2) |
| 1130 | end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2) |
| 1131 | start_idx2 = torch.where( |
| 1132 | torch.eq(x_idx, 0), |
| 1133 | torch.tensor(0, device=x.device), |
| 1134 | torch.where( |
| 1135 | torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx, |
| 1136 | ), |
| 1137 | ) |
| 1138 | y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1) |
| 1139 | start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2) |
| 1140 | end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2) |
| 1141 | cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x) |
| 1142 | return cand |
| 1143 | |
| 1144 | |
| 1145 | def expand_dims(v, dims): |
no outgoing calls
no test coverage detected