text: import torch
  import torch.optim as optim
  import numpy as np
  from PIL import Image
  import pano
  
  
  def vecang(vec1, vec2):
      vec1 = vec1 / np.sqrt((vec1 ** 2).sum())
      vec2 = vec2 / np.sqrt((vec2 ** 2).sum())
      return np.arccos(np.dot(vec1, vec2))
  
  
  def rotatevec(vec, theta):
      x = vec[0] * torch.cos(theta) - vec[1] * torch.sin(theta)
      y = vec[0] * torch.sin(theta) + vec[1] * torch.cos(theta)
      return torch.cat([x, y])
  
  
  def pts_linspace(pa, pb, pts=300):
      pa = pa.view(1, 2)
      pb = pb.view(1, 2)
      w = torch.arange(0, pts + 1, dtype=pa.dtype).view(-1, 1)
      return (pa * (pts - w) + pb * w) / pts
  
  
  def xyz2uv(xy, z=-1):
      c = torch.sqrt((xy ** 2).sum(1))
      u = torch.atan2(xy[:, 1], xy[:, 0]).view(-1, 1)
      v = torch.atan2(torch.zeros_like(c) + z, c).view(-1, 1)
      return torch.cat([u, v], dim=1)
  
  
  def uv2idx(uv, w, h):
      col = (uv[:, 0] / (2 * np.pi) + 0.5) * w - 0.5
      row = (uv[:, 1] / np.pi + 0.5) * h - 0.5
      return torch.cat([col.view(-1, 1), row.view(-1, 1)], dim=1)
  
  
  def wallidx(xy, w, h, z1, z2):
      col = (torch.atan2(xy[1], xy[0]) / (2 * np.pi) + 0.5) * w - 0.5
      c = torch.sqrt((xy ** 2).sum())
      row_s = (torch.atan2(torch.zeros_like(c) + z1, c) / np.pi + 0.5) * h - 0.5
      row_t = (torch.atan2(torch.zeros_like(c) + z2, c) / np.pi + 0.5) * h - 0.5
  
      pa = torch.cat([col.view(1), row_s.view(1)])
      pb = torch.cat([col.view(1), row_t.view(1)])
      return pts_linspace(pa, pb)
  
  
  def map_coordinates(input, coordinates):
      ''' PyTorch version of scipy.ndimage.interpolation.map_coordinates
      input: (H, W)
      coordinates: (2, ...)
      '''
      h = input.shape[0]
      w = input.shape[1]
  
      def _coordinates_pad_wrap(h, w, coordinates):
          coordinates[0] = coordinates[0] % h
          coordinates[1] = coordinates[1] % w
          return coordinates
  
      co_floor = torch.floor(coordinates).long()
      co_ceil = torch.ceil(coordinates).long()
      d1 = (coordinates[1] - co_floor[1].float())
      d2 = (coordinates[0] - co_floor[0].float())
      co_floor = _coordinates_pad_wrap(h, w, co_floor)
      co_ceil = _coordinates_pad_wrap(h, w, co_ceil)
      f00 = input[co_floor[0], co_floor[1]]
      f10 = input[co_floor[0], co_ceil[1]]
      f01 = input[co_ceil[0], co_floor[1]]
      f11 = input[co_ceil[0], co_ceil[1]]
      fx1 = f00 + d1 * (f10 - f00)
      fx2 = f01 + d1 * (f11 - f01)
      return fx1 + d2 * (fx2 - fx1)
  
  
  def pc2cor_id(pc, pc_vec, pc_theta, pc_height):
      ps = torch.stack([
          (pc + pc_vec),
          (pc + rotatevec(pc_vec, pc_theta)),
          (pc - pc_vec),
          (pc + rotatevec(pc_vec, pc_theta - np.pi))
      ])
  
      return torch.cat([
          uv2idx(xyz2uv(ps, z=-1), 1024, 512),
          uv2idx(xyz2uv(ps, z=pc_height), 1024, 512),
      ], dim=0)
  
  
  def project2sphere_score(pc, pc_vec, pc_theta, pc_height, scoreedg, scorecor, i_step=None):
  
      # Sample corner loss
      corid = pc2cor_id(pc, pc_vec, pc_theta, pc_height)
      corid_coordinates = torch.stack([corid[:, 1], corid[:, 0]])
      loss_cor = -map_coordinates(scorecor, corid_coordinates).mean()
  
      # Sample boundary loss
      p1 = pc + pc_vec
      p2 = pc + rotatevec(pc_vec, pc_theta)
      p3 = pc - pc_vec
      p4 = pc + rotatevec(pc_vec, pc_theta - np.pi)
  
      segs = [
          pts_linspace(p1, p2),
          pts_linspace(p2, p3),
          pts_linspace(p3, p4),
          pts_linspace(p4, p1),
      ]
  
      # wall-wall
  #    loss_wallwall = 0
  #    walls_idx = [
  #        wallidx(p1, 1024, 512, -1, pc_height),
  #        wallidx(p2, 1024, 512, -1, pc_height),
  #        wallidx(p3, 1024, 512, -1, pc_height),
  #        wallidx(p4, 1024, 512, -1, pc_height),
  #    ]
  #    for wall_idx in walls_idx:
  #        wall_coordinates = torch.stack([wall_idx[:, 1], wall_idx[:, 0]])
  #        loss_wallwall -= map_coordinates(scoreedg[..., 0], wall_coordinates).mean() / len(walls_idx)
  
      # ceil-wall
      loss_ceilwall = 0
      for seg in segs:
          ceil_uv = xyz2uv(seg, z=-1)
          ceil_idx = uv2idx(ceil_uv, 1024, 512)
          ceil_coordinates = torch.stack([ceil_idx[:, 1], ceil_idx[:, 0]])
          loss_ceilwall -= map_coordinates(scoreedg[..., 1], ceil_coordinates).mean() / len(segs)
  
      # floor-wall
      loss_floorwall = 0
      for seg in segs:
          floor_uv = xyz2uv(seg, z=pc_height)
          floor_idx = uv2idx(floor_uv, 1024, 512)
          floor_coordinates = torch.stack([floor_idx[:, 1], floor_idx[:, 0]])
          loss_floorwall -= map_coordinates(scoreedg[..., 2], floor_coordinates).mean() / len(segs)
  
      #losses = 1.0 * loss_cor + 0.1 * loss_wallwall + 0.5 * loss_ceilwall + 1.0 * loss_floorwall
      losses = 1.0 * loss_cor + 1.0 * loss_ceilwall + 1.0 * loss_floorwall
  
      if i_step is not None:
          with torch.no_grad():
              print('step %d: %.3f (cor %.3f, wall %.3f, ceil %.3f, floor %.3f)' % (
                  i_step, losses,
                  loss_cor, loss_wallwall,
                  loss_ceilwall, loss_floorwall))
  
      return losses
  
  
  def optimize_cor_id(cor_id, scoreedg, scorecor, num_iters=100, verbose=False):
      assert scoreedg.shape == (512, 1024, 3)
      assert scorecor.shape == (512, 1024)
  
      Z = -1
      ceil_cor_id = cor_id[0::2]
      floor_cor_id = cor_id[1::2]
      ceil_cor_id, ceil_cor_id_xy = pano.constraint_cor_id_same_z(ceil_cor_id, scorecor, Z)
      ceil_cor_id_xyz = np.hstack([ceil_cor_id_xy, np.zeros(4).reshape(-1, 1) + Z])
  
      #pc = (ceil_cor_id_xy[0] + ceil_cor_id_xy[2]) / 2
      pc = np.mean(ceil_cor_id_xy, axis=0)
      pc_vec = ceil_cor_id_xy[0] - pc
      pc_theta = vecang(pc_vec, ceil_cor_id_xy[1] - pc)
      pc_height = pano.fit_avg_z(floor_cor_id, ceil_cor_id_xy, scorecor)
  
      scoreedg = torch.FloatTensor(scoreedg)
      scorecor = torch.FloatTensor(scorecor)
      pc = torch.FloatTensor(pc)
      pc_vec = torch.FloatTensor(pc_vec)
      pc_theta = torch.FloatTensor([pc_theta])
      pc_height = torch.FloatTensor([pc_height])
      pc.requires_grad = True
      pc_vec.requires_grad = True
      pc_theta.requires_grad = True
      pc_height.requires_grad = True
  
      optimizer = optim.SGD([
          pc, pc_vec, pc_theta, pc_height
      ], lr=1e-3, momentum=0.9)
  
      best = {'score': 1e9}
  
      for i_step in range(num_iters):
          i = i_step if verbose else None
          optimizer.zero_grad()
          score = project2sphere_score(pc, pc_vec, pc_theta, pc_height, scoreedg, scorecor, i)
          if score.item() < best['score']:
              best['score'] = score.item()
              best['pc'] = pc.clone()
              best['pc_vec'] = pc_vec.clone()
              best['pc_theta'] = pc_theta.clone()
              best['pc_height'] = pc_height.clone()
          score.backward()
          optimizer.step()
  
      pc = best['pc']
      pc_vec = best['pc_vec']
      pc_theta = best['pc_theta']
      pc_height = best['pc_height']
      opt_cor_id = pc2cor_id(pc, pc_vec, pc_theta, pc_height).detach().numpy()
      opt_cor_id = np.stack([opt_cor_id[:4], opt_cor_id[4:]], axis=1).reshape(8, 2)
  
      return opt_cor_id
  