text: import torch
  import torch.nn as nn
  import numpy as np
  from torch.autograd import Variable
  import cv2
  import math
  from torch.autograd.gradcheck import gradcheck
  from torchvision import transforms, utils
  import matplotlib.pyplot as plt
  
  
  def polar_transformer(U, out_size, device, log=True, radius_factor=0.707):
      """Polar Transformer Layer
  
      Based on https://github.com/tensorflow/models/blob/master/transformer/spatial_transformer.py.
      _repeat(), _interpolate() are exactly the same;
      the polar transform implementation is in _transform()
  
      Args:
          U, theta, out_size, name: same as spatial_transformer.py
          log (bool): log-polar if True; else linear polar
          radius_factor (float): 2maxR / Width
      """
      def _repeat(x, n_repeats):
          rep = torch.ones(n_repeats)
          rep.unsqueeze(0)
          x = torch.reshape(x, (-1, 1))
          x = x * rep
          return torch.reshape(x, [-1])
  
      def _interpolate(im, x, y, out_size):  # im [B,H,W,C]
          # constants
          x = x.to(device)
          y = y.to(device)
          num_batch = im.shape[0]
          height = im.shape[1]
          width = im.shape[2]
          channels = im.shape[3]
          height_f = height
          width_f = width
  
          x = x.double()
          y = y.double()
          out_height = out_size[0]
          out_width = out_size[1]
          zero = torch.zeros([])
          max_y = im.shape[1] - 1
          max_x = im.shape[2] - 1
  
          # do sampling
          x0 = torch.floor(x).long()
          x1 = x0 + 1
          y0 = torch.floor(y).long()
          y1 = y0 + 1
  
          x0 = torch.clamp(x0, zero, max_x)
          x1 = torch.clamp(x1, zero, max_x)
          y0 = torch.clamp(y0, zero, max_y)
          y1 = torch.clamp(y1, zero, max_y)
          dim2 = width
          dim1 = width * height
  
          base = _repeat(
              torch.range(0, num_batch - 1, dtype=int) * dim1,
              out_height * out_width)
          base = base.long()
          base = base.to(device)
          base_y0 = base + y0 * dim2
          base_y1 = base + y1 * dim2
          idx_a = base_y0 + x0
          idx_b = base_y1 + x0
          idx_c = base_y0 + x1
          idx_d = base_y1 + x1
  
          # use indices to lookup pixels in the flat image and restore
          # channels dim
          im_flat = torch.reshape(im, [-1, channels])
          im_flat = im_flat.clone().float().to(device)
  
          Ia = im_flat.gather(0, idx_a.unsqueeze(1))
          Ib = im_flat.gather(0, idx_b.unsqueeze(1))
          Ic = im_flat.gather(0, idx_c.unsqueeze(1))
          Id = im_flat.gather(0, idx_d.unsqueeze(1))
  
          # Ia = im_flat[idx_a].to(device)
          # Ib = im_flat[idx_b].to(device)
          # Ic = im_flat[idx_c].to(device)
          # Id = im_flat[idx_d].to(device)
  
          # and finally calculate interpolated values
          x0_f = x0.double()
          x1_f = x1.double()
          y0_f = y0.double()
          y1_f = y1.double()
          # print(((x1_f-x) * (y1_f-y)).shape)
          # print("-------------")
          wa = ((x1_f - x) * (y1_f - y)).unsqueeze(1)
          wb = ((x1_f - x) * (y - y0_f)).unsqueeze(1)
          wc = ((x - x0_f) * (y1_f - y)).unsqueeze(1)
          wd = ((x - x0_f) * (y - y0_f)).unsqueeze(1)
  
          # output = Ia + Ib + Ic + Id
          output = wa * Ia + wb * Ib + wc * Ic + wd * Id
          return output
  
      def _meshgrid(height, width):
          x_t = torch.ones([height, 1]) * torch.linspace(
              0.0, 1.0 * width - 1, width).unsqueeze(1).permute(1, 0)
          y_t = torch.linspace(0.0, 1.0, height).unsqueeze(1) * torch.ones(
              [1, width])
  
          x_t_flat = torch.reshape(x_t, (1, -1))
          y_t_flat = torch.reshape(y_t, (1, -1))
          grid = torch.cat((x_t_flat, y_t_flat), 0)
  
          return grid
  
      def _transform(input_dim, out_size):
          # radius_factor = torch.sqrt(torch.tensor(2.))/2.
          num_batch = input_dim.shape[0]  # input [B,H,W,C]
          num_channels = input_dim.shape[3]
  
          out_height = out_size[0]
          out_width = out_size[1]
          grid = _meshgrid(out_height, out_width)  # (2, WxH)
          grid = grid.unsqueeze(0)
          grid = torch.reshape(grid, [-1])
          grid = grid.repeat(num_batch)
          grid = torch.reshape(grid, [num_batch, 2, -1])  # (B,2,WxH)
  
          ## here we do the polar/log-polar transform
          W = torch.tensor(input_dim.shape[1], dtype=torch.double)
          # print(" W shape", np.shape(W))
          # print(" W", W)
  
          # W = input_dim.shape[1].float()
          maxR = W * radius_factor
  
          # if radius is from 1 to W/2; log R is from 0 to log(W/2)
          # we map the -1 to +1 grid to log R
          # then remap to 0 to 1
          EXCESS_CONST = 1.1
  
          logbase = torch.exp(torch.log(W * EXCESS_CONST / 2) /
                              W)  #10. ** (torch.log10(maxR) / W)
          #torch.exp(torch.log(W*EXCESS_CONST/2) / W) #
          # get radius in pix
          if log:
              # min=1, max=maxR
              r_s = torch.pow(logbase, grid[:, 0, :])
          else:
              # min=1, max=maxR
              r_s = 1 + (grid[:, 0, :] + 1) / 2 * (maxR - 1)
  
          # y is from -1 to 1; theta is from 0 to 2pi
          theta = np.linspace(0., np.pi, input_dim.shape[1],
                              endpoint=False) * -1.0
          t_s = torch.from_numpy(theta).unsqueeze(1) * torch.ones([1, out_width])
          t_s = torch.reshape(t_s, (1, -1))
  
          # use + theta[:, 0] to deal with origin
          x_s = r_s * torch.cos(t_s) + (W / 2)
          y_s = r_s * torch.sin(t_s) + (W / 2)
  
          x_s_flat = torch.reshape(x_s, [-1])
          y_s_flat = torch.reshape(y_s, [-1])
  
          input_transformed = _interpolate(input_dim, x_s_flat, y_s_flat,
                                           out_size)
          output = torch.reshape(
              input_transformed,
              [num_batch, out_height, out_width, num_channels]).to(device)
          return output, logbase
  
      output, logbase = _transform(U, out_size)
      return [output, logbase]
  
  