gffx 0.1.0__py3-none-any.whl

gffx/__init__.py

@@ -0,0 +1,5 @@
+from . import linalg
+from . import obj
+from . import context
+
+from . import ray

gffx/__main__.py

@@ -0,0 +1,172 @@
+import gffx
+import torch
+import argparse
+from pathlib import Path
+import matplotlib.pyplot as plt
+
+################################################################
+# ARGUMENT PARSER
+################################################################
+parser = argparse.ArgumentParser(description='gffx')
+parser.add_argument(
+    '--vertices',
+    type=Path,
+    help='Path to vertices PyTorch file (*.pt)',
+    required=True
+)
+parser.add_argument(
+    '--faces',
+    type=Path,
+    help='Path to faces PyTorch file (*.pt)',
+    required=True
+)
+parser.add_argument(
+    '--center_and_scale',
+    action  = 'store_true',
+    help    = 'Center and scale object',
+    default = False
+)
+
+parser.add_argument(
+    '--camera_width',
+    type=int,
+    help='Camera width',
+    default=512
+)
+parser.add_argument(
+    '--camera_height',
+    type=int,
+    help='Camera height',
+    default=512
+)
+parser.add_argument(
+    '--camera_pos',
+    type=list,
+    help='Camera position',
+    default=[0, 0, 1.2]
+)
+parser.add_argument(
+    '--camera_dir',
+    type=list,
+    help='Camera direction',
+    default=[0, 0, -1]
+)
+parser.add_argument(
+    '--camera_screen_distance',
+    type    = float,
+    help    = 'Camera screen distance',
+    default = 1.0
+)
+
+parser.add_argument(
+    '--ambient_color',
+    type=list,
+    help='Ambient color',
+    default=[0.5, 0.5, 0.5]
+)
+parser.add_argument(
+    '--diffuse_color',
+    type=list,
+    help='Diffuse color',
+    default=[0.5, 0.5, 0.5]
+)
+parser.add_argument(
+    '--specular_color',
+    type=list,
+    help='Specular color',
+    default=[0.5, 0.5, 0.5]
+)
+parser.add_argument(
+    '--specular_coefficient',
+    type=float,
+    help='Specular coefficient',
+    default=1
+)
+parser.add_argument(
+    '--light_intensity',
+    type=float,
+    help='Light intensity',
+    default=1
+)
+parser.add_argument(
+    '--ambient_intensity',
+    type=float,
+    help='Ambient intensity',
+    default=0.2
+)
+parser.add_argument(
+    '--light_pos',
+    type=list,
+    help='Light position',
+    default=[5, 5, 5]
+)
+parser.add_argument(
+    '--background_color',
+    type=list,
+    help='Background color',
+    default=[0, 0, 0]
+)
+parser.add_argument(
+    '--ray_chunk_size',
+    type=int,
+    help='Ray chunk size',
+    default=1024
+)
+
+args = parser.parse_args()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Setup Camera
+camera = gffx.ray.Camera(
+    width  = args.camera_width,
+    height = args.camera_height,
+    pos    = args.camera_pos,
+    dir    = args.camera_dir,
+    screen_distance = args.camera_screen_distance,
+
+    device = device
+)
+
+# Setup list of meshes
+vertices          = torch.load(args.vertices)[0].to(device)
+faces             = torch.load(args.faces)[0].to(device)
+
+if args.center_and_scale:
+    vertices_mean     = torch.mean(vertices, dim=0)
+    vertices_centered = vertices - vertices_mean
+    vertices_scale    = torch.max(torch.linalg.norm(vertices_centered, dim=-1))
+    vertices          = vertices / vertices_scale
+
+object_list = [
+    gffx.obj.mesh_from_vertices_and_faces(
+        vertices             = vertices,
+        faces                = faces,
+        init_translation     = [0, 0.0, 0],
+        init_rotation        = [0, 0, 0],
+        init_scale           = [1, 1, 1],
+        
+        ambient_color        = args.ambient_color,
+        diffuse_color        = args.diffuse_color,
+        specular_color       = args.specular_color,
+        specular_coefficient = args.specular_coefficient,
+        
+        device           = device
+    )
+]
+
+# Ray Trace Render
+images = gffx.ray.mesh_render(
+    meshes = object_list,
+    camera = camera,
+    light_intensity     = args.light_intensity,
+    ambient_intensity   = args.ambient_intensity,
+    light_pos           = args.light_pos,
+    background_color    = args.background_color,
+    ray_chunk_size      = args.ray_chunk_size,
+    
+    device = device
+)
+
+plt.imshow((images[0].cpu()).permute(1, 0, 2))
+plt.gca().invert_yaxis()
+plt.show()

gffx/context/__init__.py

@@ -0,0 +1,14 @@
+class InteractiveDev:
+    def __init__(
+        self,
+        enabled : bool = False
+    ):
+        self.enabled = enabled
+    
+    def __enter__(self):
+        breakpoint()
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        breakpoint()
+        pass

gffx/linalg/__init__.py

@@ -0,0 +1,158 @@
+import torch
+
+def cramer(
+    A : torch.Tensor,
+    b : torch.Tensor
+) -> torch.Tensor:
+    """
+        Solve Ax = b via Cramer's Rule
+        
+        Parameters
+        ----------
+        A : torch.Tensor
+            Coefficient matrices    -> dim(..., N, N)
+        b : torch.Tensor
+            Right-hand side vectors -> dim(..., N, 1)
+    """
+    # Get shape for reshape later
+    A_shape = A.shape
+    
+    # Flatten A and b
+    A = A.flatten(start_dim=0, end_dim=-3)
+    b = b.flatten(start_dim=0, end_dim=-3)
+    device = A.device
+    
+    # Asserts
+    B, N, _ = A.shape
+    assert N == A.shape[2], "Matrices A must be square"
+    assert B == b.shape[0], "Matrices A and vectors b must have the same batch size"
+    assert N == b.shape[1], "Matrices A and vectors b must have the same number of rows"
+    assert b.shape[2] == 1, "Vectors b must be a column vector"
+    
+    # Compute determinant of A
+    M = torch.linalg.det(A) # dim(B, 1)
+    
+    # Compute determinants of A_i
+    A_dets = torch.zeros((B, N), device=device)
+    for i in range(N):
+        A_i = A.clone()
+        A_i[:, :, i] = b.squeeze(-1)
+        A_dets[:, i] = torch.linalg.det(A_i)
+    
+    # Compute solution
+    x = A_dets / M.unsqueeze(-1)
+    
+    # Reshape x
+    x = torch.reshape(x, [*A_shape[:-2], N, 1]) # dim(..., N, 1)
+    
+    return x
+
+def euler_angles_to_matrix(
+    rot_vec : torch.Tensor
+):
+    """
+        Convert Euler angles to rotation matrix
+        
+        Parameters
+        ----------
+        rot_vec : torch.Tensor
+            Rotation vector -> dim(B, 3)
+    """
+    B = rot_vec.shape[0]
+    device = rot_vec.device
+    
+    # Rodrigues' rotation formula
+    theta = torch.linalg.norm(rot_vec, dim=-1, keepdim=True)[..., None]
+    K = torch.zeros((B, 3, 3), device=device)
+    K[:,0,1] = -rot_vec[:,2]
+    K[:,0,2] =  rot_vec[:,1]
+    K[:,1,0] =  rot_vec[:,2]
+    K[:,1,2] = -rot_vec[:,0]
+    K[:,2,0] = -rot_vec[:,1]
+    K[:,2,1] =  rot_vec[:,0]
+
+    R = (
+        torch.eye(3, device=device)[None,:].expand(B,-1,-1) 
+        + torch.sin(theta) * K 
+        + (1 - torch.cos(theta)) * (K @ K)
+    )
+    R = torch.cat([R, torch.zeros((B, 3, 1), device=device)], dim=-1)
+    R = torch.cat([R, torch.zeros((B, 1, 4), device=device)], dim=-2)
+    R[:,3,3] = 1
+    
+    return R
+
+def scale_matrix(
+    scale_vec : torch.Tensor
+):
+    """
+        Convert scale vector to scale matrix
+        
+        Parameters
+        ----------
+        scale_vec : torch.Tensor
+            Scale vector -> dim(B, 3)
+    """
+    B = scale_vec.shape[0]
+    device = scale_vec.device
+    
+     # Scale matrix
+    return torch.cat([
+        torch.cat([
+            torch.eye(3, device=device)[None,:].expand(B,-1,-1) * scale_vec[..., None], 
+            torch.zeros((B, 3, 1), device=device)
+        ], dim=-1),
+        torch.cat([
+            torch.zeros((B, 1, 3), device=device),
+            torch.ones((B, 1, 1), device=device)
+        ], dim=-1)
+    ], dim=-2)
+    
+def translation_matrix(
+    translation_vec : torch.Tensor
+):
+    """
+        Convert translation vector to translation matrix
+        
+        Parameters
+        ----------
+        translation_vec : torch.Tensor
+            Translation vector -> dim(B, 3)
+    """
+    B = translation_vec.shape[0]
+    device = translation_vec.device
+    
+     # Translation matrix
+    T = torch.eye(4, device=device)[None,:].expand(B,-1,-1).clone()
+    T[:,0:3,3] = translation_vec
+    
+    return T
+
+def transformation_matrix(
+    translation_vec : torch.Tensor,
+    rotation_vec    : torch.Tensor,
+    scale_vec       : torch.Tensor
+):
+    """
+        Convert translation, rotation, and scale vectors to transformation matrix
+        
+        Parameters
+        ----------
+        translation_vec : torch.Tensor
+            Translation vector -> dim(B, 3)
+        rotation_vec    : torch.Tensor
+            Rotation vector -> dim(B, 3)
+        scale_vec       : torch.Tensor
+            Scale vector -> dim(B, 3)
+    """
+    
+    # Rodrigues' rotation formula
+    R = euler_angles_to_matrix(rotation_vec)
+
+    # Scale Matrix
+    S = scale_matrix(scale_vec)
+
+    # Translation Matrix
+    T = translation_matrix(translation_vec)
+    
+    return T @ R @ S

gffx/obj/__init__.py

@@ -0,0 +1,367 @@
+import math
+import gffx
+import torch
+from typing import Optional
+
+class Transform:
+    def __init__(
+        self, 
+        translation : Optional[list | torch.Tensor] = None,
+        rotation    : Optional[list | torch.Tensor] = None,
+        scale       : Optional[list | torch.Tensor] = None,
+        device      : Optional[torch.device]        = None
+    ):
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Defaults
+        if not isinstance(translation, torch.Tensor) and not translation:
+            translation = [0., 0., 0.]
+        if isinstance(translation, list):
+            translation = torch.tensor(translation, device=device)[None,:]
+        
+        if not isinstance(rotation, torch.Tensor) and not rotation:
+            rotation = [0., 0., 0.]
+        if isinstance(rotation, list):
+            rotation = torch.tensor(rotation, device=device)[None,:]
+            
+        if not isinstance(scale, torch.Tensor) and not scale:
+            scale = [1., 1., 1.]
+        if isinstance(scale, list):
+            scale = torch.tensor(scale, device=device)[None,:]
+        
+        # 
+        self.translation = translation.float()
+        self.rotation    = rotation.float()
+        self.scale       = scale.float()
+
+class MeshObject:
+    def __init__(
+        self, 
+        vertices         : torch.Tensor,
+        faces            : torch.Tensor,
+        normals          : Optional[torch.Tensor] = None,
+        init_transform   : Optional[Transform]    = None,
+        init_translation : Optional[torch.Tensor] = None,
+        init_rotation    : Optional[torch.Tensor] = None,
+        init_scale       : Optional[torch.Tensor] = None,
+        
+        ambient_color    : list | torch.Tensor = [0.5, 0.5, 0.5],
+        diffuse_color    : list | torch.Tensor = [0.5, 0.5, 0.5],
+        specular_color   : list | torch.Tensor = [0.5, 0.5, 0.5],
+        specular_coefficient : float = 1.0
+    ):
+        # Defaults
+        if init_transform:
+            assert init_translation is None, "If init_transform is not None, init_translation must be None"
+            assert init_rotation    is None, "If init_transform is not None, init_rotation must be None"
+            assert init_scale       is None, "If init_transform is not None, init_scale must be None"
+            self.transform = init_transform
+        else:
+            self.transform = Transform(
+                translation = init_translation,
+                rotation    = init_rotation,
+                scale       = init_scale
+            )
+        
+        self.vertices = vertices
+        self.faces    = faces
+        self.normals  = normals
+        self.device   = vertices.device
+        
+        # [4NOW] Flat colors
+        self.ambient_color = ambient_color
+        if isinstance(self.ambient_color, list):
+            self.ambient_color = torch.tensor(self.ambient_color, device=self.device, dtype=torch.float32)
+            
+        self.diffuse_color = diffuse_color
+        if isinstance(self.diffuse_color, list):
+            self.diffuse_color = torch.tensor(self.diffuse_color, device=self.device, dtype=torch.float32)
+            
+        self.specular_color = specular_color
+        if isinstance(self.specular_color, list):
+            self.specular_color = torch.tensor(self.specular_color, device=self.device, dtype=torch.float32)
+        
+        self.specular_coefficient = specular_coefficient
+        
+    def get_transformed(self):
+        """
+            Apply transformation to vertices
+            
+            Returns:
+            -------
+            torch.Tensor: Transformed vertices
+            
+            if normals are not None, also returns transformed normals
+        """
+        M = gffx.linalg.transformation_matrix(
+            translation_vec = self.transform.translation[None,...],
+            rotation_vec    = self.transform.rotation[None,...],
+            scale_vec       = self.transform.scale[None,...]
+        ) # dim(1, 4, 4)
+        vertices_h = torch.cat([self.vertices, torch.ones((self.vertices.shape[0], 1), device=self.device)], dim=-1)
+        vertices_h = vertices_h @ M.transpose(-1, -2)
+        
+        # [Optional] Transform normals
+        if self.normals is not None:
+            M = M[:, 0:3, 0:3].inverse()
+            normals = self.normals @ M
+            return vertices_h[0,..., 0:3], normals[0]
+        
+        return vertices_h[0,..., 0:3], None
+    
+def compute_normals(vertices: torch.Tensor, faces: torch.Tensor) -> torch.Tensor:
+    """
+        Compute per-vertex normals for a batch of meshes.
+
+        Args:
+            vertices (torch.Tensor): Tensor of shape (B, V, 3) with vertex positions.
+            faces (torch.Tensor): Long tensor of shape (B, F, 3) with indices into vertices.
+
+        Returns:
+            torch.Tensor: Tensor of shape (B, V, 3) with normalized vertex normals.
+    """
+    B, V, _ = vertices.shape
+    _, F, _ = faces.shape
+    device = vertices.device
+
+    # Initialize normals accumulator
+    vertex_normals = torch.zeros_like(vertices)
+
+    # Gather vertex positions for each corner of every face
+    v0 = torch.gather(vertices, 1, faces[:, :, 0].unsqueeze(-1).expand(-1, -1, 3))
+    v1 = torch.gather(vertices, 1, faces[:, :, 1].unsqueeze(-1).expand(-1, -1, 3))
+    v2 = torch.gather(vertices, 1, faces[:, :, 2].unsqueeze(-1).expand(-1, -1, 3))
+
+    # Compute face normals (unnormalized) via cross product
+    face_normals = torch.cross(v1 - v0, v2 - v0, dim=2) # dim(B,F,3)
+    face_normals = face_normals / (face_normals.norm(dim=2, keepdim=True) + 1e-8)
+
+    # Accumulate face normals to each vertex using scatter_add_
+    for i in range(3):
+        idx = faces[:, :, i].unsqueeze(-1).expand(-1, -1, 3)
+        vertex_normals.scatter_add_(1, idx, face_normals)
+
+    # Normalize the accumulated vertex normals
+    norm = vertex_normals.norm(dim=2, keepdim=True).clamp(min=1e-8)
+    vertex_normals = vertex_normals / norm
+
+    return vertex_normals
+
+def mesh_from_vertices_and_faces(
+    vertices         : torch.Tensor,
+    faces            : torch.Tensor,
+    init_transform   : Optional[Transform]           = None,
+    init_translation : Optional[list | torch.Tensor] = None,
+    init_rotation    : Optional[list | torch.Tensor] = None,
+    init_scale       : Optional[list | torch.Tensor] = None,
+    
+    ambient_color  : list | torch.Tensor = [0.5, 0.5, 0.5],
+    diffuse_color  : list | torch.Tensor = [0.5, 0.5, 0.5],
+    specular_color : list | torch.Tensor = [0.5, 0.5, 0.5],
+    specular_coefficient : float = 1.0,
+    
+    device           : Optional[torch.device]        = None
+):
+    # 
+    if isinstance(init_translation, list):
+        init_translation = torch.tensor(init_translation, device=device)
+    if isinstance(init_rotation, list):
+        init_rotation = torch.tensor(init_rotation, device=device)
+    if isinstance(init_scale, list):
+        init_scale = torch.tensor(init_scale, device=device)
+    
+    # Compute Normals
+    normals = compute_normals(vertices[None,:], faces[None,:])[0]
+    
+    return MeshObject(
+        vertices         = vertices,
+        faces            = faces,
+        normals          = normals,
+        init_transform   = init_transform,
+        init_translation = init_translation,
+        init_rotation    = init_rotation,
+        init_scale       = init_scale,
+        
+        ambient_color  = ambient_color,
+        diffuse_color  = diffuse_color,
+        specular_color = specular_color,
+        specular_coefficient = specular_coefficient
+    )
+    
+def generate_cube_mesh(
+    init_transform   : Optional[Transform]           = None,
+    init_translation : Optional[list | torch.Tensor] = None,
+    init_rotation    : Optional[list | torch.Tensor] = None,
+    init_scale       : Optional[list | torch.Tensor] = None,
+    
+    ambient_color  : list | torch.Tensor = [0.5, 0.5, 0.5],
+    diffuse_color  : list | torch.Tensor = [0.5, 0.5, 0.5],
+    specular_color : list | torch.Tensor = [0.5, 0.5, 0.5],
+    
+    device           : Optional[torch.device]        = None
+):
+    # 
+    if isinstance(init_translation, list):
+        init_translation = torch.tensor(init_translation, device=device)
+    if isinstance(init_rotation, list):
+        init_rotation = torch.tensor(init_rotation, device=device)
+    if isinstance(init_scale, list):
+        init_scale = torch.tensor(init_scale, device=device)
+    
+    # Define vertices for a cube centered at the origin with side length 2
+    vertices = torch.tensor([
+        [-1, -1, -1],
+        [-1, -1,  1],
+        [-1,  1, -1],
+        [-1,  1,  1],
+        [ 1, -1, -1],
+        [ 1, -1,  1],
+        [ 1,  1, -1],
+        [ 1,  1,  1]
+    ], device=device, dtype=torch.float32)
+
+    # Define faces as two triangles per cube face (12 triangles total)
+    faces = torch.tensor([
+        [0, 1, 2], [1, 3, 2],  # Left face
+        [4, 6, 5], [5, 6, 7],  # Right face
+        [0, 4, 1], [1, 4, 5],  # Bottom face
+        [2, 3, 6], [3, 7, 6],  # Top face
+        [0, 2, 4], [2, 6, 4],  # Back face
+        [1, 5, 3], [3, 5, 7]   # Front face
+    ], device=device, dtype=torch.int64)
+    
+    return MeshObject(
+        vertices         = vertices, 
+        faces            = faces,
+        normals          = vertices / torch.linalg.norm(vertices, dim=-1, keepdim=True),
+        init_transform   = init_transform,
+        init_translation = init_translation,
+        init_rotation    = init_rotation,
+        init_scale       = init_scale,
+        
+        ambient_color  = ambient_color,
+        diffuse_color  = diffuse_color,
+        specular_color = specular_color
+    )
+    
+    
+def generate_icosphere_mesh(
+    num_subdivisions : int                           = 1,
+    init_transform   : Optional[Transform]           = None,
+    init_translation : Optional[list | torch.Tensor] = None,
+    init_rotation    : Optional[list | torch.Tensor] = None,
+    init_scale       : Optional[list | torch.Tensor] = None,
+    
+    ambient_color  : list | torch.Tensor = [0.5, 0.5, 0.5],
+    diffuse_color  : list | torch.Tensor = [0.5, 0.5, 0.5],
+    specular_color : list | torch.Tensor = [0.5, 0.5, 0.5],
+    
+    device           : Optional[torch.device]        = None
+):
+    #
+    if isinstance(init_translation, list):
+        init_translation = torch.tensor(init_translation, device=device)
+    if isinstance(init_rotation, list):
+        init_rotation = torch.tensor(init_rotation, device=device)
+    if isinstance(init_scale, list):
+        init_scale = torch.tensor(init_scale, device=device)
+    
+    t = (1.0 + math.sqrt(5.0)) / 2.0
+
+    verts = [
+        [-1,  t,  0],
+        [ 1,  t,  0],
+        [-1, -t,  0],
+        [ 1, -t,  0],
+        [ 0, -1,  t],
+        [ 0,  1,  t],
+        [ 0, -1, -t],
+        [ 0,  1, -t],
+        [ t,  0, -1],
+        [ t,  0,  1],
+        [-t,  0, -1],
+        [-t,  0,  1],
+    ]
+    
+    # Normalize vertices to unit length
+    vertices = []
+    for v in verts:
+        norm = math.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
+        vertices.append([v[0] / norm, v[1] / norm, v[2] / norm])
+    vertices = torch.tensor(vertices, device=device, dtype=torch.float32)
+
+    faces = torch.tensor([
+        [0, 11, 5],
+        [0, 5, 1],
+        [0, 1, 7],
+        [0, 7, 10],
+        [0, 10, 11],
+        [1, 5, 9],
+        [5, 11, 4],
+        [11, 10, 2],
+        [10, 7, 6],
+        [7, 1, 8],
+        [3, 9, 4],
+        [3, 4, 2],
+        [3, 2, 6],
+        [3, 6, 8],
+        [3, 8, 9],
+        [4, 9, 5],
+        [2, 4, 11],
+        [6, 2, 10],
+        [8, 6, 7],
+        [9, 8, 1],
+    ], device=device, dtype=torch.int64)
+
+    # Subdivide faces
+    for _ in range(num_subdivisions):
+        new_faces = []
+        midpoints = {}
+        for face in faces:
+            a, b, c = face.tolist()
+            ab = tuple(sorted((a, b)))
+            ac = tuple(sorted((a, c)))
+            bc = tuple(sorted((b, c)))
+
+            if ab not in midpoints:
+                midpoints[ab] = len(vertices)
+                vertices = torch.cat([vertices, (vertices[a:a+1] + vertices[b:b+1]) / 2.0], dim=0)
+            if ac not in midpoints:
+                midpoints[ac] = len(vertices)
+                vertices = torch.cat([vertices, (vertices[a:a+1] + vertices[c:c+1]) / 2.0], dim=0)
+            if bc not in midpoints:
+                midpoints[bc] = len(vertices)
+                vertices = torch.cat([vertices, (vertices[b:b+1] + vertices[c:c+1]) / 2.0], dim=0)
+
+            d = midpoints[ab]
+            e = midpoints[ac]
+            f = midpoints[bc]
+
+            new_faces.extend([
+                [a, d, e],
+                [b, d, f],
+                [c, e, f],
+                [d, e, f]
+            ])
+        faces = torch.tensor(new_faces, device=device, dtype=torch.int64)
+    
+    # Project vertices to unit length
+    vertices = torch.nn.functional.normalize(vertices, dim=-1)
+    vertices = torch.cat([vertices, torch.ones((vertices.shape[0], 1), device=device)], dim=-1)
+    vertices = vertices @ torch.diag(torch.tensor([1, 1, 1, 0], device=device)).to(vertices.dtype)
+    vertices = vertices[..., 0:3]
+    
+    return MeshObject(
+        vertices         = vertices,
+        faces            = faces,
+        normals          = vertices,
+        init_transform   = init_transform,
+        init_translation = init_translation,
+        init_rotation    = init_rotation,
+        init_scale       = init_scale,
+        
+        ambient_color  = ambient_color,
+        diffuse_color  = diffuse_color,
+        specular_color = specular_color
+    )

gffx/ray/__init__.py

@@ -0,0 +1,284 @@
+import gffx
+import torch
+from tqdm import tqdm
+
+from .camera import *
+
+def ray_sphere_intersection(ray_origins, ray_directions, sphere_pos, sphere_radius):
+    # Check discriminant (ray-sphere intersection)
+    e_min_c = ray_origins - sphere_pos[None, None, None, :]                    # dim(B, H, W, 3)
+    d_dot_e_min_c = torch.sum(ray_directions * e_min_c, dim=-1, keepdim=True)  # dim(B, H, W, 1)
+    d_dot_d = torch.sum(ray_directions * ray_directions, dim=-1, keepdim=True) # dim(B, H, W, 1)
+    discriminant = d_dot_e_min_c ** 2 - d_dot_d * (torch.sum(e_min_c * e_min_c, dim=-1, keepdim=True) - sphere_radius ** 2)
+
+    # Compute t
+    hit_mask = (discriminant >= 0).float()
+    t = torch.sqrt(discriminant * hit_mask) / d_dot_d
+    t_minus = torch.clip(-d_dot_e_min_c - t, 0)
+    t_plus  = torch.clip(-d_dot_e_min_c + t, 0)
+    t = torch.maximum(t_minus, t_plus) * hit_mask
+
+    return t
+
+def ray_triangle_intersection(
+    ray_origins, 
+    ray_directions, 
+    triangle_vertices, 
+    t0 = 0, 
+    t1 = 100
+):
+    """
+        Compute ray-triangle intersection using Cramer's Rule
+        
+        Terms
+        -----
+        B = batch size
+        N = number of triangles
+        
+        Parameters
+        ----------
+        ray_origins : torch.Tensor
+            Ray origins -> dim(B, 3)
+        ray_directions : torch.Tensor
+            Ray directions -> dim(B, 3)
+        triangle_vertices : torch.Tensor
+            Triangle vertices -> dim(N, 3, 3)
+        t0 : float
+            Minimum t value (default: 0)
+        t1 : float
+            Maximum t value (default: 100)
+            
+        Returns
+        -------
+        beta : torch.Tensor
+            beta values -> dim(B, N)
+        gamma : torch.Tensor
+            gamma values -> dim(B, N)
+        t : torch.Tensor
+            t values -> dim(B, N)
+        intersect : torch.Tensor
+            Intersection mask -> dim(B, N)
+        
+        Notes
+        -----
+        | x_a - x_b; x_a - x_c; x_d | | beta  |   | x_a - x_e |
+        | y_a - y_b; y_a - y_c; y_d | | gamma | = | y_a - y_e |
+        | z_a - z_b; z_a - z_c; z_d | | t     |   | z_a - z_e |
+
+        A x = b
+
+        Solve via Cramer's Rule:
+            beta  = det(A_1) / det(A)
+            gamma = det(A_2) / det(A)
+            t     = det(A_3) / det(A)
+        Where A_i is A with column i replaced by b
+    """
+    B = ray_origins.shape[0]
+    N = triangle_vertices.shape[0]
+    device = triangle_vertices.device
+    
+    # Setup
+    A = torch.zeros((B, N, 3, 3), device=device)
+    b = torch.zeros((B, N, 3, 1), device=device)
+    triangle_vertices = triangle_vertices[None, :, :, :]
+    ray_directions = ray_directions[:, None, :]
+    ray_origins = ray_origins[:, None, :]
+    
+    # 
+    A[...,0,0] = triangle_vertices[...,0,0] - triangle_vertices[...,1,0]
+    A[...,0,1] = triangle_vertices[...,0,0] - triangle_vertices[...,2,0]
+    A[...,0,2] = ray_directions[...,0]
+    A[...,1,0] = triangle_vertices[...,0,1] - triangle_vertices[...,1,1]
+    A[...,1,1] = triangle_vertices[...,0,1] - triangle_vertices[...,2,1]
+    A[...,1,2] = ray_directions[...,1]
+    A[...,2,0] = triangle_vertices[...,0,2] - triangle_vertices[...,1,2]
+    A[...,2,1] = triangle_vertices[...,0,2] - triangle_vertices[...,2,2]
+    A[...,2,2] = ray_directions[...,2]
+
+    b[...,0,0] = triangle_vertices[...,0,0] - ray_origins[...,0]
+    b[...,1,0] = triangle_vertices[...,0,1] - ray_origins[...,1]
+    b[...,2,0] = triangle_vertices[...,0,2] - ray_origins[...,2]
+
+    # Cramer's Rule
+    x     = gffx.linalg.cramer(A, b)
+    beta  = x[...,0,0]
+    gamma = x[...,1,0]
+    t     = x[...,2,0]
+
+    # Intersection Test
+    intersect = ((t > t0) & (t < t1)) & ((gamma >= 0) & (gamma <= 1)) & ((beta > 0) & (beta + gamma < 1))
+
+    return beta, gamma, t, intersect
+
+# [4NOW] 1 camera
+def mesh_render(
+    meshes            : list | gffx.obj.MeshObject,
+    camera            : Camera,
+    light_intensity   : float               = 1,
+    ambient_intensity : float               = 0.2,
+    light_pos         : list | torch.Tensor = [5, 5, 5],
+    background_color  : list | torch.Tensor = [0, 0, 0],
+    ray_chunk_size    : int                 = 4096,
+    
+    verbose           : bool                = False,
+    
+    device : Optional[torch.device] = None
+):
+    # [4NOW] 1 camera  
+    B = 1
+
+    if isinstance(meshes, gffx.obj.MeshObject):
+        meshes = [meshes]
+    
+    # [4NOW] Flat colors
+    if isinstance(background_color, list):
+        background_color = torch.tensor(background_color, device=device, dtype=torch.float32)
+
+    diffuse_colors = [obj.diffuse_color for obj in meshes] + [background_color]
+    diffuse_colors = torch.stack(diffuse_colors, dim=0).to(device)
+
+    specular_coefficients = [obj.specular_coefficient for obj in meshes] + [1]
+    specular_coefficients = torch.tensor(specular_coefficients, device=device, dtype=torch.float32)
+
+    specular_colors = [obj.specular_color for obj in meshes] + [background_color]
+    specular_colors = torch.stack(specular_colors, dim=0).to(device)
+
+    ambient_colors = [obj.ambient_color for obj in meshes] + [background_color]
+    ambient_colors = torch.stack(ambient_colors, dim=0).to(device)
+
+    # Light setup
+    if isinstance(light_pos, list):
+        light_pos  = torch.tensor(light_pos, device=device, dtype=torch.float32)
+
+    # 
+    object_hit = torch.full((B * camera.width * camera.height,), -1, device=device, dtype=torch.int64)
+    face_hit   = torch.full((B * camera.width * camera.height,), -1, device=device, dtype=torch.int64)
+    t_val      = torch.full((B * camera.width * camera.height,), float('inf'), device=device)
+    normals    = torch.zeros((B * camera.width * camera.height, 3), device=device)
+    hit_pos    = torch.zeros((B * camera.width * camera.height, 3), device=device)
+
+    # 
+    if verbose:
+        print('Ray Chunk Size:', ray_chunk_size)
+    ray_origins    = camera.ray_origins.view(B * camera.width * camera.height, 3)       # dim(B * W * H, 3)
+    ray_directions = camera.ray_directions.view(B * camera.width * camera.height, 3) # dim(B * W * H, 3)
+
+    # 
+    for obj_idx, obj in enumerate(meshes):
+        transformed_vertices, transformed_normals = obj.get_transformed()
+        
+        face_tri_vertices = transformed_vertices[obj.faces] # dim(F, 3, 3)
+        face_tri_normals  = transformed_normals[obj.faces]  # dim(F, 3, 3)
+        
+        # Chunking
+        pbar = tqdm(
+            iterable = range(0, B * camera.width * camera.height, ray_chunk_size),
+            desc     = f'Object {obj_idx}',
+        )
+        for i in pbar:
+            C = ray_chunk_size if i + ray_chunk_size < B * camera.width * camera.height else (B * camera.width * camera.height - i)
+            
+            # 
+            beta, gamma, t, intersect = gffx.ray.ray_triangle_intersection(
+                ray_origins       = ray_origins[i:i+C],    # dim(C, 3)
+                ray_directions    = ray_directions[i:i+C], # dim(C, 3)
+                triangle_vertices = face_tri_vertices,     # dim(F, 3, 3)
+                t0                = 0,
+                t1                = 100
+            )
+            
+            # Choose the smallest t which intersects
+            t_valid = torch.where(
+                intersect,
+                torch.where(t < 0, float('inf'), t),
+                float('inf')
+            )
+            face_idx_valid_min = t_valid.argmin(dim=-1, keepdim=True)
+            beta_valid_min = torch.gather(
+                input = beta,
+                dim   = -1,
+                index = face_idx_valid_min
+            )[:,0]
+            gamma_valid_min = torch.gather(
+                input = gamma,
+                dim   = -1,
+                index = face_idx_valid_min
+            )[:,0]
+            
+            t_valid_min = torch.gather(
+                input = t_valid,
+                dim   = -1,
+                index = face_idx_valid_min
+            )[:,0]
+            intersect_valid_min = torch.gather(
+                input = intersect,
+                dim   = -1,
+                index = face_idx_valid_min
+            )[:,0]
+            
+            # 
+            object_hit[i:i+C] = torch.where(
+                ((t_valid_min < t_val[i:i+C]) & intersect_valid_min),
+                obj_idx,
+                object_hit[i:i+C]
+            )
+            
+            #
+            face_hit[i:i+C] = torch.where(
+                ((t_valid_min < t_val[i:i+C]) & intersect_valid_min),
+                face_idx_valid_min[:,0],
+                face_hit[i:i+C]
+            )
+            
+            # Interpolate normals
+            interpolated_normals = (1 - beta - gamma)[...,None] * face_tri_normals[None, :, 0] + beta[...,None] * face_tri_normals[None, :, 1] + gamma[...,None] * face_tri_normals[None, :, 2]
+
+            interpolated_normals = torch.gather(
+                input = interpolated_normals,   # dim(C, F, 3)
+                dim   = 1,
+                index = face_idx_valid_min[..., None].expand(-1, -1, 3)
+            )[:,0,:]
+            normals[i:i+C] = torch.where(
+                ((t_valid_min < t_val[i:i+C]) & intersect_valid_min)[:,None],
+                interpolated_normals,
+                normals[i:i+C]
+            )
+            
+            # Interpolate hit positions
+            hit_pos[i:i+C] = torch.where(
+                ((t_valid_min < t_val[i:i+C]) & intersect_valid_min)[:,None],
+                ray_origins[i:i+C] + t_valid_min[:,None] * ray_directions[i:i+C],
+                hit_pos[i:i+C]
+            )
+            
+            #
+            t_val[i:i+C] = torch.where(
+                ((t_valid_min < t_val[i:i+C]) & intersect_valid_min),
+                t_valid_min,
+                t_val[i:i+C]
+            )
+        
+    # Reshape
+    object_hit = object_hit.view(B, camera.width, camera.height)
+    face_hit   = face_hit.view(B, camera.width, camera.height)
+    normals    = normals.view(B, camera.width, camera.height, 3)
+    hit_pos    = hit_pos.view(B, camera.width, camera.height, 3)
+
+    # Compute Phong Shading
+    light_pos = light_pos[None, None, None, :]
+    light_dir = light_pos - hit_pos
+    light_dir /= torch.linalg.norm(light_dir, dim=-1, keepdim=True)
+
+    view_dir  = camera.pos[None, None, None, :] - hit_pos
+    view_dir /= torch.linalg.norm(view_dir, dim=-1, keepdim=True)
+    bisector_vec  = light_dir + view_dir
+    bisector_vec /= torch.linalg.norm(bisector_vec, dim=-1, keepdim=True)
+
+    diffuse_weight  = torch.clamp(torch.sum(light_dir * normals, dim=-1, keepdim=True), min=0)
+    specular_weight = torch.clamp(torch.sum(bisector_vec * normals, dim=-1, keepdim=True), min=0) ** specular_coefficients[object_hit][...,None]
+
+    L  = diffuse_colors[object_hit]  * light_intensity * diffuse_weight
+    L += specular_colors[object_hit] * light_intensity * specular_weight
+    L += ambient_colors[object_hit]  * ambient_intensity
+
+    return L

gffx/ray/camera.py

@@ -0,0 +1,66 @@
+import torch
+from typing import Optional
+
+class Camera:
+    def __init__(
+        self, 
+        width    : int = 512,
+        height   : int = 512,
+        pos      : list | torch.Tensor = [0, 0, 5], 
+        dir      : list | torch.Tensor = [0, 0, -1],
+        world_up : list | torch.Tensor = [0, 1, 0],
+        ray_origin_scale_du : float = 0,
+        ray_origin_scale_dv : float = 0,
+        screen_distance : float     = 1.0,
+        device   : Optional[torch.device] = None
+    ):
+        self.width  = width
+        self.height = height
+        self.ray_origin_scale_du = ray_origin_scale_du
+        self.ray_origin_scale_dv = ray_origin_scale_dv
+        self.screen_distance = screen_distance
+        
+        if not device:
+            self.device = device
+            
+        self.pos = pos
+        if isinstance(self.pos, list):
+            self.pos = torch.tensor(self.pos, device=device, dtype=torch.float32)
+        
+        self.dir = dir
+        if isinstance(self.dir, list):
+            self.dir = torch.tensor(self.dir, device=device, dtype=torch.float32)
+        
+        self.world_up  = world_up
+        if isinstance(self.world_up, list):
+            self.world_up = torch.tensor(self.world_up, device=device, dtype=torch.float32)
+            
+        # Camera basis
+        self.u  = torch.cross(self.dir, self.world_up, dim=-1)
+        self.u /= torch.linalg.norm(self.u)
+        self.v  = torch.cross(self.u, self.dir, dim=-1)
+        self.w  = -self.dir
+        
+        # Screen
+        self.aspect_ratio = self.width / self.height
+        self.left         = -self.aspect_ratio
+        self.right        = self.aspect_ratio
+        self.bottom       = -1
+        self.top          = 1
+        
+        self.screen_grids = torch.stack(torch.meshgrid(
+            (self.right - self.left) * ((torch.arange(0, self.width, device=device) + 0.5) / self.width) + self.left,
+            (self.top - self.bottom) * ((torch.arange(0, self.height, device=device) + 0.5) / self.height) + self.bottom,
+            indexing='ij'
+        ), dim=-1).float().expand(1, -1, -1, -1)
+
+        self.ray_origins  = self.ray_origin_scale_du * self.screen_grids[..., 0:1] * self.u[None, None, None, :]
+        self.ray_origins += self.ray_origin_scale_dv * self.screen_grids[..., 1:2] * self.v[None, None, None, :]
+        self.ray_origins += self.pos[None,None,None,:]
+
+        # Compute viewing rays
+        self.ray_directions  = self.screen_grids[..., 0:1] * self.u[None, None, None, :]
+        self.ray_directions += self.screen_grids[..., 1:2] * self.v[None, None, None, :]
+        self.ray_directions += self.pos[None,None,None,:] - (self.w[None,None,None,:] * self.screen_distance)
+        self.ray_directions -= self.ray_origins
+        self.ray_directions /= torch.linalg.norm(self.ray_directions, dim=-1, keepdim=True)

gffx-0.1.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Danzel Serrano
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

gffx-0.1.0.dist-info/METADATA

@@ -0,0 +1,84 @@
+Metadata-Version: 2.3
+Name: gffx
+Version: 0.1.0
+Summary: A minimal graphics library in PyTorch
+License: MIT
+Author: Danzel Serrano
+Author-email: danzel.serrano@gmail.com
+Requires-Python: >=3.8,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Description-Content-Type: text/markdown
+
+### gffx - a minimal library for (differentiable) graphics
+
+```Python
+import gffx
+import torch
+import matplotlib.pyplot as plt
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Setup Camera
+camera = gffx.ray.Camera(
+    width   = 512,
+    height  = 512,
+    pos = [0, 0, 3],
+    dir = [0, 0, -1],
+
+    device = device
+)
+
+# Setup list of meshes
+vertices          = torch.load('path/to/vertices.pt').to(device)
+faces             = torch.load('path/to/vertices.pt').to(device)
+
+object_list = [
+    gffx.obj.mesh_from_vertices_and_faces(
+        vertices             = vertices,
+        faces                = faces,
+        init_translation     = [0, 0.0, 2],
+        init_rotation        = [0, 0, 0],
+        init_scale           = [1, 1, 1],
+        
+        ambient_color        = [0.5, 0.5, 0.5],
+        diffuse_color        = [0.5, 0.5, 0.5],
+        specular_color       = [0.5, 0.5, 0.5],
+        specular_coefficient = 1,
+        
+        device           = device
+    )
+]
+
+# Ray Trace Render
+images = gffx.ray.mesh_render(
+    meshes = object_list,
+    camera = camera,
+    light_intensity = 1.0,
+    ambient_intensity = 0.2,
+    light_pos = [5, 5, 5],
+    background_color = [0, 0, 0],
+    ray_chunk_size = 4096,
+    
+    device = device
+)
+
+plt.imshow((images[0].cpu()).permute(1, 0, 2))
+plt.gca().invert_yaxis()
+plt.show()
+```
+
+### [WIP] CLI
+```
+python -m gffx --vertices '/path/to/vertices.pt' --faces 'path/to/faces.pt'
+```
+
+### Motivation
+1. Installing PyTorch3D everytime in colab is a hassle. This library's first aim is to render a mesh using native operations in PyTorch.
+

gffx-0.1.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+gffx/__init__.py,sha256=nsP_q2t5aSWHcTXGRmqFkjlomO72Z7VXkRBKRxgdAG4,79
+gffx/__main__.py,sha256=-tUJcq8bapS4yoQ8RKkaj4zT_FAgQcvbTXpThLG1ORo,3964
+gffx/context/__init__.py,sha256=yIRcexatet4RZ_I0wac-O1qED2yY8tAhYbBKehEW64A,288
+gffx/linalg/__init__.py,sha256=eDf2sfLy-CuLDBuvQiskDXIIRsJaDuU9TP1bv4NUgD8,4196
+gffx/obj/__init__.py,sha256=_r8ysXwFzQYhGK9NilYd8oKGRB1R89owaATbbjMhNB8,13559
+gffx/ray/__init__.py,sha256=HeV_z5zkEIT9EV9F9HZCZplUn10Wmv2lQ2WIIjFnmIc,10865
+gffx/ray/camera.py,sha256=PBcJ0a62g7K7on4J9YqTrynAvV1OFXW0c6cpqY8nUf0,2837
+gffx-0.1.0.dist-info/LICENSE,sha256=ye96oZM4dBrUCrY2b9pyp7w07J6-SV72Jzq9yv2z5ro,1070
+gffx-0.1.0.dist-info/METADATA,sha256=ZrZDxc51AM6d7_8PVcOXye0avVXSQwUKCHJxmpdaISU,2255
+gffx-0.1.0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+gffx-0.1.0.dist-info/RECORD,,

gffx-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.1.1
+Root-Is-Purelib: true
+Tag: py3-none-any