pose2theta.py

import math
import torch

def tfm2theta_2d(tfm, h, w):
    """Returns theta for affine_grid"""
    theta = tfm.new_zeros((tfm.size(0), 2, 3))
    theta[:,0,0] = tfm[:,0,0]
    theta[:,0,1] = tfm[:,0,1]*h/w
    theta[:,0,2] = tfm[:,0,2]*2/w + tfm[:,0,0] + tfm[:,0,1] - 1

    theta[:,1,0] = tfm[:,1,0]*w/h
    theta[:,1,1] = tfm[:,1,1]
    theta[:,1,2] = tfm[:,1,2]*2/h + tfm[:,1,0] + tfm[:,1,1] - 1

    return theta

def angle2rotmat_2d(theta):
    c, s = torch.cos(theta), torch.sin(theta)
    return torch.stack([torch.stack([c, -s], dim=-1),
                        torch.stack([s,  c], dim=-1)], dim=-1)

def pose2tfm_2d(pose): # x, y, rotation in deg
    x, y, r = pose.t()
    tfm = x.new_zeros((x.size(0), 3, 3))
    tfm[:, :2, :2] = angle2rotmat_2d(r/180.*math.pi)
    tfm[:,0,2] = x
    tfm[:,1,2] = y
    tfm[:,2,2] = 1
    return tfm

def pose2theta_2d(pose, h, w):
    return tfm2theta_2d(pose2tfm_2d(pose), h, w)