spikezoo 0.1.1__py3-none-any.whl → 0.2__py3-none-any.whl

spikezoo/archs/stir/package_core/build/lib/package_core/geometry.py

@@ -0,0 +1,458 @@
+"""
+A collection of geometric transformation operations
+@author: Zhaoyang Lv 
+@Date: March, 2019
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import torch
+import torch.nn as nn
+from torch import sin, cos, atan2, acos
+
+_NEXT_AXIS = [1, 2, 0, 1]
+
+# map axes strings to/from tuples of inner axis, parity, repetition, frame
+_AXES2TUPLE = {
+    'sxyz': (0, 0, 0, 0), 'sxyx': (0, 0, 1, 0), 'sxzy': (0, 1, 0, 0),
+    'sxzx': (0, 1, 1, 0), 'syzx': (1, 0, 0, 0), 'syzy': (1, 0, 1, 0),
+    'syxz': (1, 1, 0, 0), 'syxy': (1, 1, 1, 0), 'szxy': (2, 0, 0, 0),
+    'szxz': (2, 0, 1, 0), 'szyx': (2, 1, 0, 0), 'szyz': (2, 1, 1, 0),
+    'rzyx': (0, 0, 0, 1), 'rxyx': (0, 0, 1, 1), 'ryzx': (0, 1, 0, 1),
+    'rxzx': (0, 1, 1, 1), 'rxzy': (1, 0, 0, 1), 'ryzy': (1, 0, 1, 1),
+    'rzxy': (1, 1, 0, 1), 'ryxy': (1, 1, 1, 1), 'ryxz': (2, 0, 0, 1),
+    'rzxz': (2, 0, 1, 1), 'rxyz': (2, 1, 0, 1), 'rzyz': (2, 1, 1, 1)}
+
+_TUPLE2AXES = dict((v, k) for k, v in _AXES2TUPLE.items())
+
+def meshgrid(H, W, B=None, is_cuda=False):
+    """ torch version of numpy meshgrid function
+    :input
+    :param height
+    :param width
+    :param batch size
+    :param initialize a cuda tensor if true
+    -------
+    :return 
+    :param meshgrid in column
+    :param meshgrid in row
+    """
+    u = torch.arange(0, W)
+    v = torch.arange(0, H)
+
+    if is_cuda:
+        u, v = u.cuda(), v.cuda()
+
+    u = u.repeat(H, 1).view(1,H,W)
+    v = v.repeat(W, 1).t_().view(1,H,W)
+
+    if B is not None:
+        u, v = u.repeat(B,1,1,1), v.repeat(B,1,1,1)
+    return u, v
+
+def generate_xy_grid(B, H, W, K):
+    """ Generate a batch of image grid from image space to world space 
+        px = (u - cx) / fx
+        py = (y - cy) / fy
+        function tested in 'test_geometry.py'
+    :input
+    :param batch size
+    :param height
+    :param width
+    :param camera intrinsic array [fx,fy,cx,cy] 
+    ---------
+    :return 
+    :param 
+    :param 
+    """
+    fx, fy, cx, cy = K.split(1,dim=1)
+    uv_grid = meshgrid(H, W, B)
+    u_grid, v_grid = [uv.type_as(cx) for uv in uv_grid]
+    px = ((u_grid.view(B,-1) - cx) / fx).view(B,1,H,W)
+    py = ((v_grid.view(B,-1) - cy) / fy).view(B,1,H,W)
+    return px, py
+
+def batch_inverse_Rt(R, t):
+    """ The inverse of the R, t: [R' | -R't] 
+        function tested in 'test_geometry.py'
+    :input 
+    :param rotation Bx3x3
+    :param translation Bx3
+    ----------
+    :return 
+    :param rotation inverse Bx3x3
+    :param translation inverse Bx3
+    """
+    R_t = R.transpose(1,2)
+    t_inv = -torch.bmm(R_t, t.contiguous().view(-1, 3, 1))
+
+    return R_t, t_inv.view(-1,3)
+
+def batch_Rt_compose(d_R, d_t, R0, t0):
+    """ Compose operator of R, t: [d_R*R | d_R*t + d_t] 
+        We use left-mulitplication rule here. 
+        function tested in 'test_geometry.py'
+    
+    :input
+    :param rotation incremental Bx3x3
+    :param translation incremental Bx3
+    :param initial rotation Bx3x3
+    :param initial translation Bx3
+    ----------
+    :return 
+    :param composed rotation Bx3x3
+    :param composed translation Bx3
+    """
+    R1 = d_R.bmm(R0)
+    t1 = d_R.bmm(t0.view(-1,3,1)) + d_t.view(-1,3,1)
+    return R1, t1.view(-1,3)
+
+def batch_Rt_between(R0, t0, R1, t1): 
+    """ Between operator of R, t, transform of T_0=[R0, t0] to T_1=[R1, t1]
+        which is T_1 \compose T^{-1}_0 
+        function tested in 'test_geometry.py'
+    
+    :input 
+    :param rotation of source Bx3x3
+    :param translation of source Bx3
+    :param rotation of target Bx3x3
+    :param translation of target Bx3
+    ----------
+    :return 
+    :param incremental rotation Bx3x3
+    :param incremnetal translation Bx3
+    """
+    R0t = R0.transpose(1,2)
+    dR = R1.bmm(R0t)
+    dt = t1.view(-1,3) - dR.bmm(t0.view(-1,3,1)).view(-1,3)
+    return dR, dt
+
+def batch_skew(w):
+    """ Generate a batch of skew-symmetric matrices. 
+        function tested in 'test_geometry.py'
+    :input
+    :param skew symmetric matrix entry Bx3
+    ---------
+    :return 
+    :param the skew-symmetric matrix Bx3x3
+    """
+    B, D = w.size()
+    assert(D == 3)
+    o = torch.zeros(B).type_as(w)
+    w0, w1, w2 = w[:, 0], w[:, 1], w[:, 2]
+    return torch.stack((o, -w2, w1, w2, o, -w0, -w1, w0, o), 1).view(B, 3, 3)
+
+def batch_twist2Mat(twist):
+    """ The exponential map from so3 to SO3
+        Calculate the rotation matrix using Rodrigues' Rotation Formula
+        http://electroncastle.com/wp/?p=39 
+        or Ethan Eade's lie group note:
+        http://ethaneade.com/lie.pdf equation (13)-(15) 
+        @todo: may rename the interface to batch_so3expmap(twist)
+        functioned tested with cv2.Rodrigues implementation in 'test_geometry.py'
+    :input
+    :param twist/axis angle Bx3 \in \so3 space 
+    ----------
+    :return 
+    :param Rotation matrix Bx3x3 \in \SO3 space
+    """
+    B = twist.size()[0]
+    theta = twist.norm(p=2, dim=1).view(B, 1)
+    w_so3 = twist / theta.expand(B, 3)
+    W = batch_skew(w_so3)
+    return torch.eye(3).repeat(B,1,1).type_as(W) \
+        + W*sin(theta.view(B,1,1)) \
+        + W.bmm(W)*(1-cos(theta).view(B,1,1))
+
+def batch_mat2angle(R):
+    """ Calcuate the axis angles (twist) from a batch of rotation matrices
+        Ethan Eade's lie group note:
+        http://ethaneade.com/lie.pdf equation (17)
+        function tested in 'test_geometry.py'
+    :input
+    :param Rotation matrix Bx3x3 \in \SO3 space
+    --------
+    :return 
+    :param the axis angle B
+    """
+    R1 = [torch.trace(R[i]) for i in range(R.size()[0])]
+    R_trace = torch.stack(R1)
+    # clamp if the angle is too large (break small angle assumption)
+    # @todo: not sure whether it is absoluately necessary in training. 
+    angle = acos( ((R_trace - 1)/2).clamp(-1,1))
+    return angle
+
+def batch_mat2twist(R):
+    """ The log map from SO3 to so3
+        Calculate the twist vector from Rotation matrix 
+        Ethan Eade's lie group note:
+        http://ethaneade.com/lie.pdf equation (18)
+        @todo: may rename the interface to batch_so3logmap(R)
+        function tested in 'test_geometry.py'
+        @note: it currently does not consider extreme small values. 
+        If you use it as training loss, you may run into problems
+    :input
+    :param Rotation matrix Bx3x3 \in \SO3 space 
+    --------
+    :param the twist vector Bx3 \in \so3 space
+    """
+    B = R.size()[0]
+ 
+    R1 = [torch.trace(R[i]) for i in range(R.size()[0])]
+    tr = torch.stack(R1)
+    theta = acos( ((tr - 1)/2).clamp(-1,1) )
+
+    r11,r12,r13,r21,r22,r23,r31,r32,r33 = torch.split(R.view(B,-1),1,dim=1)
+    res = torch.cat([r32-r23, r13-r31, r21-r12],dim=1)  
+
+    magnitude = (0.5*theta/sin(theta))
+
+    return magnitude.view(B,1) * res
+
+def batch_warp_inverse_depth(p_x, p_y, p_invD, pose, K):
+    """ Compute the warping grid w.r.t. the SE3 transform given the inverse depth
+    :input
+    :param p_x the x coordinate map
+    :param p_y the y coordinate map
+    :param p_invD the inverse depth
+    :param pose the 3D transform in SE3
+    :param K the intrinsics
+    --------
+    :return 
+    :param projected u coordinate in image space Bx1xHxW
+    :param projected v coordinate in image space Bx1xHxW
+    :param projected inverse depth Bx1XHxW
+    """
+    [R, t] = pose
+    B, _, H, W = p_x.shape
+
+    I = torch.ones((B,1,H,W)).type_as(p_invD)
+    x_y_1 = torch.cat((p_x, p_y, I), dim=1)
+
+    warped = torch.bmm(R, x_y_1.view(B,3,H*W)) + \
+        t.view(B,3,1).expand(B,3,H*W) * p_invD.view(B, 1, H*W).expand(B,3,H*W)
+
+    x_, y_, s_ = torch.split(warped, 1, dim=1)
+    fx, fy, cx, cy = torch.split(K, 1, dim=1)
+
+    u_ = (x_ / s_).view(B,-1) * fx + cx
+    v_ = (y_ / s_).view(B,-1) * fy + cy
+
+    inv_z_ = p_invD / s_.view(B,1,H,W)
+
+    return u_.view(B,1,H,W), v_.view(B,1,H,W), inv_z_
+
+def batch_warp_affine(pu, pv, affine):
+    # A = affine[:,:,:2]
+    # t = affine[:,:, 2]
+    B,_,H,W = pu.shape
+    ones = torch.ones(pu.shape).type_as(pu)
+    uv = torch.cat((pu, pv, ones), dim=1)
+    uv = torch.bmm(affine, uv.view(B,3,-1)) #+ t.view(B,2,1)
+    return uv[:,0].view(B,1,H,W), uv[:,1].view(B,1,H,W)
+
+def check_occ(inv_z_buffer, inv_z_ref, u, v, thres=1e-1):
+    """ z-buffering check of occlusion 
+    :param inverse depth of target frame
+    :param inverse depth of reference frame
+    """
+    B, _, H, W = inv_z_buffer.shape
+
+    inv_z_warped = warp_features(inv_z_ref, u, v)
+    inlier = (inv_z_buffer > inv_z_warped - thres)
+
+    inviews = inlier & (u > 0) & (u < W) & \
+        (v > 0) & (v < H)
+
+    return 1-inviews
+
+def warp_features(F, u, v):
+    """
+    Warp the feature map (F) w.r.t. the grid (u, v)
+    """
+    B, C, H, W = F.shape
+
+    u_norm = u / ((W-1)/2) - 1
+    v_norm = v / ((H-1)/2) - 1
+    uv_grid = torch.cat((u_norm.view(B,H,W,1), v_norm.view(B,H,W,1)), dim=3)
+    F_warped = nn.functional.grid_sample(F, uv_grid,
+        mode='bilinear', padding_mode='border')
+    return F_warped
+
+def batch_transform_xyz(xyz_tensor, R, t, get_Jacobian=True):
+    '''
+    transform the point cloud w.r.t. the transformation matrix
+    :param xyz_tensor: B * 3 * H * W
+    :param R: rotation matrix B * 3 * 3
+    :param t: translation vector B * 3
+    '''
+    B, C, H, W = xyz_tensor.size()
+    t_tensor = t.contiguous().view(B,3,1).repeat(1,1,H*W)
+    p_tensor = xyz_tensor.contiguous().view(B, C, H*W)
+    # the transformation process is simply:
+    # p' = t + R*p
+    xyz_t_tensor = torch.baddbmm(t_tensor, R, p_tensor)
+
+    if get_Jacobian:
+        # return both the transformed tensor and its Jacobian matrix
+        J_r = R.bmm(batch_skew_symmetric_matrix(-1*p_tensor.permute(0,2,1)))
+        J_t = -1 * torch.eye(3).view(1,3,3).expand(B,3,3)
+        J = torch.cat((J_r, J_t), 1)
+        return xyz_t_tensor.view(B, C, H, W), J
+    else:
+        return xyz_t_tensor.view(B, C, H, W)
+
+def flow_from_rigid_transform(depth, extrinsic, intrinsic):
+    """
+    Get the optical flow induced by rigid transform [R,t] and depth
+    """
+    [R, t] = extrinsic
+    [fx, fy, cx, cy] = intrinsic
+
+def batch_project(xyz_tensor, K):
+    """ Project a point cloud into pixels (u,v) given intrinsic K
+    [u';v';w] = [K][x;y;z]
+    u = u' / w; v = v' / w
+    :param the xyz points 
+    :param calibration is a torch array composed of [fx, fy, cx, cy]
+    -------
+    :return u, v grid tensor in image coordinate
+    (tested through inverse project)
+    """
+    B, _, H, W = xyz_tensor.size()
+    batch_K = K.expand(H, W, B, 4).permute(2,3,0,1)
+
+    x, y, z = torch.split(xyz_tensor, 1, dim=1)
+    fx, fy, cx, cy = torch.split(batch_K, 1, dim=1)
+
+    u = fx*x / z + cx
+    v = fy*y / z + cy
+    return torch.cat((u,v), dim=1)
+
+def batch_inverse_project(depth, K):
+    """ Inverse project pixels (u,v) to a point cloud given intrinsic 
+    :param depth dim B*H*W
+    :param calibration is torch array composed of [fx, fy, cx, cy]
+    :param color (optional) dim B*3*H*W
+    -------
+    :return xyz tensor (batch of point cloud)
+    (tested through projection)
+    """
+    if depth.dim() == 3:
+        B, H, W = depth.size()
+    else: 
+        B, _, H, W = depth.size()
+
+    x, y = generate_xy_grid(B,H,W,K)
+    z = depth.view(B,1,H,W)
+    return torch.cat((x*z, y*z, z), dim=1)
+
+def batch_euler2mat(ai, aj, ak, axes='sxyz'):
+    """ A torch implementation euler2mat from transform3d:
+    https://github.com/matthew-brett/transforms3d/blob/master/transforms3d/euler.py
+    :param ai : First rotation angle (according to `axes`).
+    :param aj : Second rotation angle (according to `axes`).
+    :param ak : Third rotation angle (according to `axes`).
+    :param axes : Axis specification; one of 24 axis sequences as string or encoded tuple - e.g. ``sxyz`` (the default).
+    -------
+    :return rotation matrix, array-like shape (B, 3, 3)
+    Tested w.r.t. transforms3d.euler module
+    """
+    B = ai.size()[0]
+
+    try:
+        firstaxis, parity, repetition, frame = _AXES2TUPLE[axes.lower()]
+    except (AttributeError, KeyError):
+        _TUPLE2AXES[axes]  # validation
+        firstaxis, parity, repetition, frame = axes
+
+    i = firstaxis
+    j = _NEXT_AXIS[i+parity]
+    k = _NEXT_AXIS[i-parity+1]
+    order = [i, j, k]
+
+    if frame:
+        ai, ak = ak, ai
+    if parity:
+        ai, aj, ak = -ai, -aj, -ak
+
+    si, sj, sk = sin(ai), sin(aj), sin(ak)
+    ci, cj, ck = cos(ai), cos(aj), cos(ak)
+    cc, cs = ci*ck, ci*sk
+    sc, ss = si*ck, si*sk
+
+    # M = torch.zeros(B, 3, 3).cuda()
+    if repetition:
+        c_i = [cj, sj*si, sj*ci]
+        c_j = [sj*sk, -cj*ss+cc, -cj*cs-sc]
+        c_k = [-sj*ck, cj*sc+cs, cj*cc-ss]
+    else:
+        c_i = [cj*ck, sj*sc-cs, sj*cc+ss]
+        c_j = [cj*sk, sj*ss+cc, sj*cs-sc]
+        c_k = [-sj, cj*si, cj*ci]
+
+    def permute(X): # sort X w.r.t. the axis indices
+        return [ x for (y, x) in sorted(zip(order, X)) ]
+
+    c_i = permute(c_i)
+    c_j = permute(c_j)
+    c_k = permute(c_k)
+
+    r =[torch.stack(c_i, 1),
+        torch.stack(c_j, 1),
+        torch.stack(c_k, 1)]
+    r = permute(r)
+
+    return torch.stack(r, 1)
+
+def batch_mat2euler(M, axes='sxyz'): 
+    """ A torch implementation euler2mat from transform3d:
+    https://github.com/matthew-brett/transforms3d/blob/master/transforms3d/euler.py
+    :param array-like shape (3, 3) or (4, 4). Rotation matrix or affine.
+    :param  Axis specification; one of 24 axis sequences as string or encoded tuple - e.g. ``sxyz`` (the default).
+    --------
+    :returns 
+    :param ai : First rotation angle (according to `axes`).
+    :param aj : Second rotation angle (according to `axes`).
+    :param ak : Third rotation angle (according to `axes`).
+    """
+    try:
+        firstaxis, parity, repetition, frame = _AXES2TUPLE[axes.lower()]
+    except (AttributeError, KeyError):
+        _TUPLE2AXES[axes]  # validation
+        firstaxis, parity, repetition, frame = axes
+
+    i = firstaxis
+    j = _NEXT_AXIS[i+parity]
+    k = _NEXT_AXIS[i-parity+1]
+
+    if repetition:
+        sy = torch.sqrt(M[:, i, j]**2 + M[:, i, k]**2)
+        # A lazy way to cope with batch data. Can be more efficient
+        mask = ~(sy > 1e-8) 
+        ax = atan2( M[:, i, j],  M[:, i, k])
+        ay = atan2( sy,          M[:, i, i])
+        az = atan2( M[:, j, i], -M[:, k, i])
+        if mask.sum() > 0:
+            ax[mask] = atan2(-M[:, j, k][mask], M[:, j, j][mask])
+            ay[mask] = atan2( sy[mask],         M[:, i, i][mask])
+            az[mask] = 0.0
+    else:
+        cy = torch.sqrt(M[:, i, i]**2 + M[:, j, i]**2)
+        mask = ~(cy > 1e-8)
+        ax = atan2( M[:, k, j],  M[:, k, k])
+        ay = atan2(-M[:, k, i],  cy)
+        az = atan2( M[:, j, i],  M[:, i, i])
+        if mask.sum() > 0:
+            ax[mask] = atan2(-M[:, j, k][mask],  M[:, j, j][mask])
+            ay[mask] = atan2(-M[:, k, i][mask],  cy[mask])
+            az[mask] = 0.0
+
+    if parity:
+        ax, ay, az = -ax, -ay, -az
+    if frame:
+        ax, az = az, ax
+    return ax, ay, az

spikezoo/archs/stir/package_core/build/lib/package_core/image_proc.py

@@ -0,0 +1,183 @@
+import cv2
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+
+def white_balance(img):
+    img = (img*255.).astype(np.uint8)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
+    avg_a = np.average(img[:, :, 1])
+    avg_b = np.average(img[:, :, 2])
+    img[:, :, 1] = img[:, :, 1] - ((avg_a - 128) * (img[:, :, 0] / 255.0) * 1.1)
+    img[:, :, 2] = img[:, :, 2] - ((avg_b - 128) * (img[:, :, 0] / 255.0) * 1.1)
+    img = cv2.cvtColor(img, cv2.COLOR_LAB2BGR)
+    img = img.astype(np.float)/255.
+    return img
+
+def warp_image_flow(ref_image, flow):
+    [B, _, H, W] = ref_image.size()
+    
+    # mesh grid 
+    xx = torch.arange(0, W).view(1,-1).repeat(H,1)
+    yy = torch.arange(0, H).view(-1,1).repeat(1,W)
+    xx = xx.view(1,1,H,W).repeat(B,1,1,1)
+    yy = yy.view(1,1,H,W).repeat(B,1,1,1)
+    grid = torch.cat((xx,yy),1).float()
+
+    if ref_image.is_cuda:
+        grid = grid.cuda()
+
+    flow_f = flow + grid
+    flow_fx = flow_f[:, 0, :, :] 
+    flow_fy = flow_f[:, 1, :, :]
+
+    with torch.no_grad():
+        mask_x = ~((flow_fx < 0) | (flow_fx > (W - 1)))
+        mask_y = ~((flow_fy < 0) | (flow_fy > (H - 1)))
+        mask = mask_x & mask_y
+        mask = mask.unsqueeze(1)
+
+    flow_fx = flow_fx / float(W) * 2. - 1.
+    flow_fy = flow_fy / float(H) * 2. - 1.
+
+    flow_fxy = torch.stack([flow_fx, flow_fy], dim=-1)
+    img = torch.nn.functional.grid_sample(ref_image, flow_fxy, padding_mode='zeros') 
+    return img, mask
+
+def warp_image(depth_cur, T_cur2ref, K, img_ref, crop_tl_h=0, crop_tl_w=0):
+    B,_, H, W = depth_cur.size()
+    fx, fy, cx, cy=K[:,0,0], K[:,1,1], K[:,0,2], K[:,1,2]
+    fx = fx.unsqueeze(-1)
+    fy = fy.unsqueeze(-1)
+    cx = cx.unsqueeze(-1)
+    cy = cy.unsqueeze(-1)
+
+    x_ref = torch.arange(0, W, 1).float().cuda() + crop_tl_w
+    y_ref = torch.arange(0, H, 1).float().cuda() + crop_tl_h
+
+    x_ref = x_ref.unsqueeze(0).repeat(B,1)
+    y_ref = y_ref.unsqueeze(0).repeat(B,1)
+
+    x_ref = (x_ref - cx)/fx
+    y_ref = (y_ref - cy)/fy
+
+    x_ref = x_ref.unsqueeze(1)
+    y_ref = y_ref.unsqueeze(-1)
+
+    xx_ref = x_ref.repeat(1, H, 1).unsqueeze(1)
+    yy_ref = y_ref.repeat(1, 1, W).unsqueeze(1)
+    ones = torch.ones_like(xx_ref)
+
+    p3d_ref = torch.cat([xx_ref, yy_ref, ones], dim=1)*depth_cur
+    ones = torch.ones_like(depth_cur)
+
+    p4d_ref = torch.cat([p3d_ref, ones], dim=1)
+    p4d_ref = p4d_ref.view(B,4,-1)
+
+    p3d_cur = T_cur2ref.bmm(p4d_ref)
+    p3d_cur = p3d_cur/(p3d_cur[:,2,:].unsqueeze(dim=1) + 1e-8)
+    p2d_cur = K.clone().bmm(p3d_cur)[:,:2,:]
+
+    p2d_cur[:,0,:]=p2d_cur[:,0,:] - crop_tl_w
+    p2d_cur[:,1,:]=p2d_cur[:,1,:] - crop_tl_h
+
+    # normalize
+    p2d_cur[:,0,:] = 2.0 * (p2d_cur[:,0,:] - W * 0.5 + 0.5) / (W - 1.)
+    p2d_cur[:,1,:] = 2.0 * (p2d_cur[:,1,:] - H * 0.5 + 0.5) / (H - 1.)
+
+    p2d_cur = p2d_cur.permute(0, 2, 1)
+    p2d_cur = p2d_cur.view(B,H,W,2)
+
+    with torch.no_grad():
+        mask_x = ~((p2d_cur[:,:,:,0] < -1.) | (p2d_cur[:,:,:,0] > 1.))
+        mask_y = ~((p2d_cur[:,:,:,1] < -1.) | (p2d_cur[:,:,:,1] > 1.))
+        mask = mask_x & mask_y
+        mask = mask.unsqueeze(1)
+
+    syn_ref_image = torch.nn.functional.grid_sample(img_ref, p2d_cur, padding_mode='zeros')    
+    return syn_ref_image, mask.float()
+
+class Grid_gradient_central_diff():
+    def __init__(self, nc, padding=True, diagonal=False):
+        self.conv_x = nn.Conv2d(nc, nc, kernel_size=2, stride=1, bias=False)
+        self.conv_y = nn.Conv2d(nc, nc, kernel_size=2, stride=1, bias=False)
+        self.conv_xy = None
+        if diagonal:
+            self.conv_xy = nn.Conv2d(nc, nc, kernel_size=2, stride=1, bias=False)
+    
+        self.padding=None
+        if padding:
+            self.padding = nn.ReplicationPad2d([0,1,0,1])
+
+        fx = torch.zeros(nc, nc, 2, 2).float().cuda()
+        fy = torch.zeros(nc, nc, 2, 2).float().cuda()
+        if diagonal:
+            fxy = torch.zeros(nc, nc, 2, 2).float().cuda()
+        
+        fx_ = torch.tensor([[1,-1],[0,0]]).cuda()
+        fy_ = torch.tensor([[1,0],[-1,0]]).cuda()
+        if diagonal:
+            fxy_ = torch.tensor([[1,0],[0,-1]]).cuda()
+
+        for i in range(nc):
+            fx[i, i, :, :] = fx_
+            fy[i, i, :, :] = fy_
+            if diagonal:
+                fxy[i,i,:,:] = fxy_
+            
+        self.conv_x.weight = nn.Parameter(fx)
+        self.conv_y.weight = nn.Parameter(fy)
+        if diagonal:
+            self.conv_xy.weight = nn.Parameter(fxy)
+
+    def __call__(self, grid_2d):
+        _image = grid_2d
+        if self.padding is not None:
+            _image = self.padding(_image)
+        dx = self.conv_x(_image)
+        dy = self.conv_y(_image)
+
+        if self.conv_xy is not None:
+            dxy = self.conv_xy(_image)
+            return dx, dy, dxy
+        return dx, dy
+
+class RandomScaleCrop(object):
+    """Randomly zooms images up to 15% and crop them to keep same size as before."""
+    def __call__(self, images, intrinsics):
+        assert intrinsics is not None
+        output_intrinsics = intrinsics.clone()
+
+        _, _, in_h, in_w = images.size()
+        x_scaling, y_scaling = np.random.uniform(1, 1.15, 2)
+        scaled_h, scaled_w = int(in_h * y_scaling), int(in_w * x_scaling)
+
+        output_intrinsics[:, 0, 0] *= x_scaling
+        output_intrinsics[:, 0, 2] *= x_scaling
+        output_intrinsics[:, 1, 1] *= y_scaling
+        output_intrinsics[:, 1, 2] *= y_scaling
+        scaled_images = F.interpolate(images, size=(scaled_h, scaled_w), mode='bilinear')
+       
+        offset_y = np.random.randint(scaled_h - in_h + 1)
+        offset_x = np.random.randint(scaled_w - in_w + 1)
+        cropped_images = scaled_images[:, :, offset_y:offset_y + in_h, offset_x:offset_x + in_w]
+
+        output_intrinsics[:, 0, 2] -= offset_x
+        output_intrinsics[:, 1, 2] -= offset_y
+
+        return cropped_images, output_intrinsics
+
+if __name__=='__main__':
+    from skimage import io
+    import cv2
+    image = io.imread('/home/peidong/leonhard/project/infk/cvg/liup/mydata/KITTI/odometry/resized/832x256/test/09/image_2/001545.jpg')
+    image = torch.from_numpy(image.transpose(2,0,1)).float()/255.
+    image = image.unsqueeze(0)
+
+    intrinsics = torch.eye(3).float().unsqueeze(0)
+    cropped_images, intrinsics = RandomScaleCrop()(image, intrinsics)
+
+    cv2.imshow('orig', image.numpy().transpose(0,2,3,1)[0])
+    cv2.imshow('crop', cropped_images.numpy().transpose(0,2,3,1)[0])
+    cv2.waitKey(0)

spikezoo/archs/stir/package_core/build/lib/package_core/linalg.py

@@ -0,0 +1,40 @@
+import torch
+
+def transformation_matrix_multiply(T1, T2):
+    if T1.dim() == 2:
+        T1=torch.unsqueeze(T1, 0)
+
+    if T2.dim() == 2:
+        T2=torch.unsqueeze(T2, 0)
+
+    R1 = T1[:, :, :3]
+    t1 = T1[:, :, 3:4]
+
+    R2 = T2[:, :, :3]
+    t2 = T2[:, :, 3:4]
+
+    R = torch.bmm(R1, R2)
+    t = torch.bmm(R1, t2) + t1
+
+    return torch.cat([R, t], dim=2)
+
+def transformation_matrix_inverse(T):
+    if T.dim() == 2:
+        T = T.unsqueeze(dim=0)
+    R = T[:, :, :3]
+    t = T[:, :, 3:4]
+
+    R_inv = R.transpose(2, 1)
+    t_inv = torch.bmm(R_inv, t)
+    t_inv = -1. * t_inv
+    return torch.cat([R_inv, t_inv], dim=2)
+
+def skew_matrix(phi):
+    Phi=torch.zeros([phi.shape[0], 3, 3], dtype=phi.dtype, device=phi.device)
+    Phi[:, 0, 1] = -phi[:, 2]
+    Phi[:, 1, 0] = phi[:, 2]
+    Phi[:, 0, 2] = phi[:, 1]
+    Phi[:, 2, 0] = -phi[:, 1]
+    Phi[:, 1, 2] = -phi[:, 0]
+    Phi[:, 2, 1] = phi[:, 0]
+    return Phi