| 584 | return Li |
| 585 | |
| 586 | def apply_transformation(self, theta, points): |
| 587 | if theta.dim() == 2: |
| 588 | theta = theta.unsqueeze(2).unsqueeze(3) |
| 589 | # points should be in the [B,H,W,2] format, |
| 590 | # where points[:,:,:,0] are the X coords |
| 591 | # and points[:,:,:,1] are the Y coords |
| 592 | |
| 593 | # input are the corresponding control points P_i |
| 594 | batch_size = theta.size()[0] |
| 595 | # split theta into point coordinates |
| 596 | Q_X = theta[:, :self.N, :, :].squeeze(3) |
| 597 | Q_Y = theta[:, self.N:, :, :].squeeze(3) |
| 598 | |
| 599 | # get spatial dimensions of points |
| 600 | points_b = points.size()[0] |
| 601 | points_h = points.size()[1] |
| 602 | points_w = points.size()[2] |
| 603 | |
| 604 | # repeat pre-defined control points along spatial dimensions of points to be transformed |
| 605 | P_X = self.P_X.expand((1, points_h, points_w, 1, self.N)) |
| 606 | P_Y = self.P_Y.expand((1, points_h, points_w, 1, self.N)) |
| 607 | |
| 608 | # compute weigths for non-linear part |
| 609 | W_X = torch.bmm(self.Li[:, :self.N, :self.N].expand((batch_size, self.N, self.N)), Q_X) |
| 610 | W_Y = torch.bmm(self.Li[:, :self.N, :self.N].expand((batch_size, self.N, self.N)), Q_Y) |
| 611 | # reshape |
| 612 | # W_X,W,Y: size [B,H,W,1,N] |
| 613 | W_X = W_X.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(1, points_h, points_w, 1, 1) |
| 614 | W_Y = W_Y.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(1, points_h, points_w, 1, 1) |
| 615 | # compute weights for affine part |
| 616 | A_X = torch.bmm(self.Li[:, self.N:, :self.N].expand((batch_size, 3, self.N)), Q_X) |
| 617 | A_Y = torch.bmm(self.Li[:, self.N:, :self.N].expand((batch_size, 3, self.N)), Q_Y) |
| 618 | # reshape |
| 619 | # A_X,A,Y: size [B,H,W,1,3] |
| 620 | A_X = A_X.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(1, points_h, points_w, 1, 1) |
| 621 | A_Y = A_Y.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(1, points_h, points_w, 1, 1) |
| 622 | |
| 623 | # compute distance P_i - (grid_X,grid_Y) |
| 624 | # grid is expanded in point dim 4, but not in batch dim 0, as points P_X,P_Y are fixed for all batch |
| 625 | points_X_for_summation = points[:, :, :, 0].unsqueeze(3).unsqueeze(4).expand( |
| 626 | points[:, :, :, 0].size() + (1, self.N)) |
| 627 | points_Y_for_summation = points[:, :, :, 1].unsqueeze(3).unsqueeze(4).expand( |
| 628 | points[:, :, :, 1].size() + (1, self.N)) |
| 629 | |
| 630 | if points_b == 1: |
| 631 | delta_X = points_X_for_summation - P_X |
| 632 | delta_Y = points_Y_for_summation - P_Y |
| 633 | else: |
| 634 | # use expanded P_X,P_Y in batch dimension |
| 635 | delta_X = points_X_for_summation - P_X.expand_as(points_X_for_summation) |
| 636 | delta_Y = points_Y_for_summation - P_Y.expand_as(points_Y_for_summation) |
| 637 | |
| 638 | dist_squared = torch.pow(delta_X, 2) + torch.pow(delta_Y, 2) |
| 639 | # U: size [1,H,W,1,N] |
| 640 | dist_squared[dist_squared == 0] = 1 # avoid NaN in log computation |
| 641 | U = torch.mul(dist_squared, torch.log(dist_squared)) |
| 642 | |
| 643 | # expand grid in batch dimension if necessary |