Python torch.cross方法代碼示例

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def qrot(q, v):
    """
    Rotate vector(s) v about the rotation described by quaternion(s) q.
    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
    where * denotes any number of dimensions.
    Returns a tensor of shape (*, 3).
    """
    assert q.shape[-1] == 4
    assert v.shape[-1] == 3
    assert q.shape[:-1] == v.shape[:-1]
    
    original_shape = list(v.shape)
    q = q.view(-1, 4)
    v = v.view(-1, 3)
    
    qvec = q[:, 1:]
    uv = torch.cross(qvec, v, dim=1)
    uuv = torch.cross(qvec, uv, dim=1)
    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape) 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def cross_prod(x, y):
    z = torch.cross(x.view(-1, 3), y.view(-1, 3), dim=1).view_as(x)
    return z 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def forward(self, z, pos, batch=None):
        """"""
        edge_index = radius_graph(pos, r=self.cutoff, batch=batch)

        i, j, idx_i, idx_j, idx_k, idx_kj, idx_ji = self.triplets(
            edge_index, num_nodes=z.size(0))

        # Calculate distances.
        dist = (pos[i] - pos[j]).pow(2).sum(dim=-1).sqrt()

        # Calculate angles.
        pos_i = pos[idx_i]
        pos_ji, pos_ki = pos[idx_j] - pos_i, pos[idx_k] - pos_i
        a = (pos_ji * pos_ki).sum(dim=-1)
        b = torch.cross(pos_ji, pos_ki).norm(dim=-1)
        angle = torch.atan2(b, a)

        rbf = self.rbf(dist)
        sbf = self.sbf(dist, angle, idx_kj)

        # Embedding block.
        x = self.emb(z, rbf, i, j)
        P = self.output_blocks[0](x, rbf, i, num_nodes=pos.size(0))

        # Interaction blocks.
        for interaction_block, output_block in zip(self.interaction_blocks,
                                                   self.output_blocks[1:]):
            x = interaction_block(x, rbf, sbf, idx_kj, idx_ji)
            P += output_block(x, rbf, i)

        return P.sum(dim=0) if batch is None else scatter(P, batch, dim=0) 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def get_angle(v1: Tensor, v2: Tensor) -> Tensor:
    return torch.atan2(
        torch.cross(v1, v2, dim=1).norm(p=2, dim=1), (v1 * v2).sum(dim=1)) 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def compute_normal_map(x_img, y_img, z, intrinsics):
    cam_coords = lift(x_img, y_img, z, intrinsics)
    cam_coords = util.lin2img(cam_coords)

    shift_left = cam_coords[:, :, 2:, :]
    shift_right = cam_coords[:, :, :-2, :]

    shift_up = cam_coords[:, :, :, 2:]
    shift_down = cam_coords[:, :, :, :-2]

    diff_hor = F.normalize(shift_right - shift_left, dim=1)[:, :, :, 1:-1]
    diff_ver = F.normalize(shift_up - shift_down, dim=1)[:, :, 1:-1, :]

    cross = torch.cross(diff_hor, diff_ver, dim=1)
    return cross 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def get_normal_map(opt,index,vertices,faces):
	face_vertices = vertices[faces.long()]
	v1,v2,v3 = torch.unbind(face_vertices,dim=1)
	normal = F.normalize(torch.cross(v2-v1,v3-v2),dim=1)
	# face normals towards camera
	normal[normal[:,2]<0] *= -1
	normal_color = (normal+1)/2
	normal_color = torch.cat([torch.zeros(1,3,device=opt.device),normal_color],dim=0)
	normal_color[0] = 0
	normal_map = normal_color[index.long()+1].permute(2,0,1)
	return normal_map 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def get_angles(a, b):
    '''
    calculate the angle between vector a and b
    :param a: Bx3xMxK tensor
    :param b: Bx3xMxK tensor
    :return: Bx1xMxK tensor
    '''
    axb = torch.cross(a, b, dim=1)  # Bx3xMxK
    a_1x3 = a.permute(0, 2, 3, 1).contiguous().unsqueeze(3)  # BxMxKx3 -> BxMxKx1x3
    b_3x1 = b.permute(0, 2, 3, 1).contiguous().unsqueeze(4)  # BxMxKx3 -> BxMxKx3x1
    ab = torch.matmul(a_1x3, b_3x1).squeeze(3).squeeze(3)  # BxMxKx1x1

    angle = torch.atan2(torch.norm(axb, dim=1, keepdim=False), ab).unsqueeze(1)
    return angle 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def qrot(q, v):
    """
    Rotate vector(s) v about the rotation described by quaternion(s) q.
    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
    where * denotes any number of dimensions.
    Returns a tensor of shape (*, 3).
    """
    assert q.shape[-1] == 4
    assert v.shape[-1] == 3
    assert q.shape[:-1] == v.shape[:-1]

    qvec = q[..., 1:]
    uv = torch.cross(qvec, v, dim=len(q.shape) - 1)
    uuv = torch.cross(qvec, uv, dim=len(q.shape) - 1)
    return v + 2 * (q[..., :1] * uv + uuv) 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def vertex_normals(vertices, faces):
    """
    :param vertices: [batch size, number of vertices, 3]
    :param faces: [batch size, number of faces, 3]
    :return: [batch size, number of vertices, 3]
    """
    assert (vertices.ndimension() == 3)
    assert (faces.ndimension() == 3)
    assert (vertices.shape[0] == faces.shape[0])
    assert (vertices.shape[2] == 3)
    assert (faces.shape[2] == 3)

    bs, nv = vertices.shape[:2]
    bs, nf = faces.shape[:2]
    device = vertices.device
    normals = torch.zeros(bs * nv, 3).to(device)

    faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None] # expanded faces
    vertices_faces = vertices.reshape((bs * nv, 3))[faces.long()]

    faces = faces.view(-1, 3)
    vertices_faces = vertices_faces.view(-1, 3, 3)

    normals.index_add_(0, faces[:, 1].long(), 
                       torch.cross(vertices_faces[:, 2] - vertices_faces[:, 1], vertices_faces[:, 0] - vertices_faces[:, 1]))
    normals.index_add_(0, faces[:, 2].long(), 
                       torch.cross(vertices_faces[:, 0] - vertices_faces[:, 2], vertices_faces[:, 1] - vertices_faces[:, 2]))
    normals.index_add_(0, faces[:, 0].long(),
                       torch.cross(vertices_faces[:, 1] - vertices_faces[:, 0], vertices_faces[:, 2] - vertices_faces[:, 0]))

    normals = F.normalize(normals, eps=1e-6, dim=1)
    normals = normals.reshape((bs, nv, 3))
    # pytorch only supports long and byte tensors for indexing
    return normals 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def surface_normals(self):
        if self._surface_normals_update:
            v10 = self.face_vertices[:, :, 0] - self.face_vertices[:, :, 1]
            v12 = self.face_vertices[:, :, 2] - self.face_vertices[:, :, 1]
            self._surface_normals = F.normalize(torch.cross(v12, v10), p=2, dim=2, eps=1e-6)
            self._surface_normals_update = False
        return self._surface_normals 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def rot_from_2_vectors(v1, v2):
        """ Returns a Rotation matrix between vectors 'v1' and 'v2'    """
        v1 = v1/torch.norm(v1)
        v2 = v2/torch.norm(v2)
        v = torch.cross(v1, v2)
        cosang = v1.matmul(v2)
        sinang = torch.norm(v)
        Rot = TORCHIEKF.Id3 + TORCHIEKF.skew(v) + \
              TORCHIEKF.skew(v).mm(TORCHIEKF.skew(v))*(1-cosang)/(sinang**2)
        return Rot 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def s2s2_to_SO3(v1, v2=None):
    '''Normalize 2 3-vectors. Project second to orthogonal component.
    Take cross product for third. Stack to form SO matrix.'''
    if v2 is None:
        assert v1.shape[-1] == 6
        v2 = v1[...,3:]
        v1 = v1[...,0:3]
    u1 = v1
    e1 = u1 / u1.norm(p=2, dim=-1, keepdim=True).clamp(min=1E-5)
    u2 = v2 - (e1 * v2).sum(-1, keepdim=True) * e1
    e2 = u2 / u2.norm(p=2, dim=-1, keepdim=True).clamp(min=1E-5)
    e3 = torch.cross(e1, e2)
    return torch.stack([e1, e2, e3], 1) 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def calculate_normals(points , policy = "upright"):
    if policy is "upright":
        points_temp = F.pad(points, (0, 1, 0, 0), mode="replicate")
        dx = points_temp[:, :, :, :-1] - points_temp[:, :, :, 1:]  # NCHW
        points_temp = F.pad(points, (0, 0, 0, 1), mode="replicate")
        dy = points_temp[:, :, :-1, :] - points_temp[:, :, 1:, :]  # NCHW
        normals = torch.cross(dy,dx)
        #mask = (points[:,2,:,:] == 0).float()
        #normals /= torch.sqrt(torch.sum(normals*normals, 1) + mask)        
        #return normals
        return torch.nn.functional.normalize(normals) 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def nomal_loss(pred, targetN,params,depthI,depthJ):
    depthI = depthI.permute(0, 2, 3, 1)
    depthJ = depthJ.permute(0, 2, 3, 1)

    predN_1 = torch.zeros_like(targetN)
    predN_2 = torch.zeros_like(targetN)

    f = params[:, :, :, 0]
    cx = params[:, :, :, 1]
    cy = params[:, :, :, 2]

    z1 = depthJ - pred
    z1 = torch.squeeze(z1)
    depthJ = torch.squeeze(depthJ)
    predN_1[:, :, :, 0] = ((MatJ - cx) * z1 + depthJ) * 1.0 / f
    predN_1[:, :, :, 1] = (MatI - cy) * z1 * 1.0 / f
    predN_1[:, :, :, 2] = z1

    z2 = depthI - pred
    z2 = torch.squeeze(z2)
    depthI = torch.squeeze(depthI)
    predN_2[:, :, :, 0] = (MatJ - cx) * z2  * 1.0 / f
    predN_2[:, :, :, 1] = ((MatI - cy) * z2 + depthI) * 1.0 / f
    predN_2[:, :, :, 2] = z2

    predN = torch.cross(predN_1, predN_2)
    pred_n = F.normalize(predN)
    pred_n = pred_n.contiguous().view(-1, 3)
    target_n = targetN.contiguous().view(-1, 3)

    loss_function = nn.CosineEmbeddingLoss()
    loss = loss_function(pred_n, target_n, Variable(torch.Tensor(pred_n.size(0)).cuda().fill_(1.0)))
    return loss 

# 需要導入模塊: import torch [as 別名]
# 或者: from torch import cross [as 別名]
def compute_face_normals(self, vertices):
        triangles = vertices[self.faces_torch, :]
        u = triangles[::, 1] - triangles[::, 0]
        v = triangles[::, 2] - triangles[::, 0]
        if self.clockwise:
            n = -torch.cross(u, v)
        else:
            n = torch.cross(u, v)
        l2 = (n ** 2).sum(dim=1)
        norm = l2.sqrt()
        nn = n / norm[:, None]
        return nn