pyg-nightly 2.7.0.dev20250701__py3-none-any.whl → 2.7.0.dev20250703__py3-none-any.whl

torch_geometric/datasets/teeth3ds.py

@@ -0,0 +1,269 @@
+import json
+import os
+import os.path as osp
+from glob import glob
+from typing import Callable, Dict, List, Optional
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from torch_geometric.data import (
+    Data,
+    InMemoryDataset,
+    download_url,
+    extract_zip,
+)
+
+
+class Teeth3DS(InMemoryDataset):
+    r"""The Teeth3DS+ dataset from the `"An Extended Benchmark for Intra-oral
+    3D Scans Analysis" <https://crns-smartvision.github.io/teeth3ds/>`_ paper.
+
+    This dataset is the first comprehensive public benchmark designed to
+    advance the field of intra-oral 3D scan analysis developed as part of the
+    3DTeethSeg 2022 and 3DTeethLand 2024 MICCAI challenges, aiming to drive
+    research in teeth identification, segmentation, labeling, 3D modeling,
+    and dental landmark identification.
+    The dataset includes at least 1,800 intra-oral scans (containing 23,999
+    annotated teeth) collected from 900 patients, covering both upper and lower
+    jaws separately.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str): The split name (one of :obj:`"Teeth3DS"`,
+            :obj:`"3DTeethSeg22_challenge"` or :obj:`"3DTeethLand_challenge"`).
+        train (bool, optional): If :obj:`True`, loads the training dataset,
+            otherwise the test dataset. (default: :obj:`True`)
+        num_samples (int, optional): Number of points to sample from each mesh.
+            (default: :obj:`30000`)
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before every access.
+            (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+    """
+    urls = {
+        'data_part_1.zip':
+        'https://osf.io/download/qhprs/',
+        'data_part_2.zip':
+        'https://osf.io/download/4pwnr/',
+        'data_part_3.zip':
+        'https://osf.io/download/frwdp/',
+        'data_part_4.zip':
+        'https://osf.io/download/2arn4/',
+        'data_part_5.zip':
+        'https://osf.io/download/xrz5f/',
+        'data_part_6.zip':
+        'https://osf.io/download/23hgq/',
+        'data_part_7.zip':
+        'https://osf.io/download/u83ad/',
+        'train_test_split':
+        'https://files.de-1.osf.io/v1/'
+        'resources/xctdy/providers/osfstorage/?zip='
+    }
+
+    sample_url = {
+        'teeth3ds_sample': 'https://osf.io/download/vr38s/',
+    }
+
+    landmarks_urls = {
+        '3DTeethLand_landmarks_train.zip': 'https://osf.io/download/k5hbj/',
+        '3DTeethLand_landmarks_test.zip': 'https://osf.io/download/sqw5e/',
+    }
+
+    def __init__(
+        self,
+        root: str,
+        split:
+        str = 'Teeth3DS',  # [3DTeethSeg22_challenge, 3DTeethLand_challenge]
+        train: bool = True,
+        num_samples: int = 30000,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ) -> None:
+
+        self.mode = 'training' if train else 'testing'
+        self.split = split
+        self.num_samples = num_samples
+
+        super().__init__(root, transform, pre_transform,
+                         force_reload=force_reload)
+
+    @property
+    def processed_dir(self) -> str:
+        return os.path.join(self.root, f'processed_{self.split}_{self.mode}')
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ['license.txt']
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        # Directory containing train/test split files:
+        split_subdir = 'teeth3ds_sample' if self.split == 'sample' else ''
+        split_dir = osp.join(
+            self.raw_dir,
+            split_subdir,
+            f'{self.split}_train_test_split',
+        )
+
+        split_files = glob(osp.join(split_dir, f'{self.mode}*.txt'))
+
+        # Collect all file names from the split files:
+        combined_list = []
+        for file_path in split_files:
+            with open(file_path) as file:
+                combined_list.extend(file.read().splitlines())
+
+        # Generate the list of processed file paths:
+        return [f'{file_name}.pt' for file_name in combined_list]
+
+    def download(self) -> None:
+        if self.split == 'sample':
+            for key, url in self.sample_url.items():
+                path = download_url(url, self.root, filename=key)
+                extract_zip(path, self.raw_dir)
+                os.unlink(path)
+        else:
+            for key, url in self.urls.items():
+                path = download_url(url, self.root, filename=key)
+                extract_zip(path, self.raw_dir)
+                os.unlink(path)
+            for key, url in self.landmarks_urls.items():
+                path = download_url(url, self.root, filename=key)
+                extract_zip(path, self.raw_dir)  # Extract each downloaded part
+                os.unlink(path)
+
+    def process_file(self, file_path: str) -> Optional[Data]:
+        """Processes the input file path to load mesh data, annotations,
+        and prepare the input features for a graph-based deep learning model.
+        """
+        import trimesh
+        from fpsample import bucket_fps_kdline_sampling
+
+        mesh = trimesh.load_mesh(file_path)
+
+        if isinstance(mesh, list):
+            # Handle the case where a list of Geometry objects is returned
+            mesh = mesh[0]
+
+        vertices = mesh.vertices
+        vertex_normals = mesh.vertex_normals
+
+        # Perform sampling on mesh vertices:
+        if len(vertices) < self.num_samples:
+            sampled_indices = np.random.choice(
+                len(vertices),
+                self.num_samples,
+                replace=True,
+            )
+        else:
+            sampled_indices = bucket_fps_kdline_sampling(
+                vertices,
+                self.num_samples,
+                h=5,
+                start_idx=0,
+            )
+
+        if len(sampled_indices) != self.num_samples:
+            raise RuntimeError(f"Sampled points mismatch, expected "
+                               f"{self.num_samples} points, but got "
+                               f"{len(sampled_indices)} for '{file_path}'")
+
+        # Extract features and annotations for the sampled points:
+        pos = torch.tensor(vertices[sampled_indices], dtype=torch.float)
+        x = torch.tensor(vertex_normals[sampled_indices], dtype=torch.float)
+
+        # Load segmentation annotations:
+        seg_annotation_path = file_path.replace('.obj', '.json')
+        if osp.exists(seg_annotation_path):
+            with open(seg_annotation_path) as f:
+                seg_annotations = json.load(f)
+            y = torch.tensor(
+                np.asarray(seg_annotations['labels'])[sampled_indices],
+                dtype=torch.float)
+            instances = torch.tensor(
+                np.asarray(seg_annotations['instances'])[sampled_indices],
+                dtype=torch.float)
+        else:
+            y = torch.empty(0, 3)
+            instances = torch.empty(0, 3)
+
+        # Load landmarks annotations:
+        landmarks_annotation_path = file_path.replace('.obj', '__kpt.json')
+
+        # Parse keypoint annotations into structured tensors:
+        keypoints_dict: Dict[str, List] = {
+            key: []
+            for key in [
+                'Mesial', 'Distal', 'Cusp', 'InnerPoint', 'OuterPoint',
+                'FacialPoint'
+            ]
+        }
+        keypoint_tensors: Dict[str, torch.Tensor] = {
+            key: torch.empty(0, 3)
+            for key in [
+                'Mesial', 'Distal', 'Cusp', 'InnerPoint', 'OuterPoint',
+                'FacialPoint'
+            ]
+        }
+        if osp.exists(landmarks_annotation_path):
+            with open(landmarks_annotation_path) as f:
+                landmarks_annotations = json.load(f)
+
+            for keypoint in landmarks_annotations['objects']:
+                keypoints_dict[keypoint['class']].extend(keypoint['coord'])
+
+            keypoint_tensors = {
+                k: torch.tensor(np.asarray(v),
+                                dtype=torch.float).reshape(-1, 3)
+                for k, v in keypoints_dict.items()
+            }
+
+        data = Data(
+            pos=pos,
+            x=x,
+            y=y,
+            instances=instances,
+            jaw=file_path.split('.obj')[0].split('_')[1],
+            mesial=keypoint_tensors['Mesial'],
+            distal=keypoint_tensors['Distal'],
+            cusp=keypoint_tensors['Cusp'],
+            inner_point=keypoint_tensors['InnerPoint'],
+            outer_point=keypoint_tensors['OuterPoint'],
+            facial_point=keypoint_tensors['FacialPoint'],
+        )
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        return data
+
+    def process(self) -> None:
+        for file in tqdm(self.processed_file_names):
+            name = file.split('.')[0]
+            path = osp.join(self.raw_dir, '**', '*', name + '.obj')
+            paths = glob(path)
+            if len(paths) == 1:
+                data = self.process_file(paths[0])
+                torch.save(data, osp.join(self.processed_dir, file))
+
+    def len(self) -> int:
+        return len(self.processed_file_names)
+
+    def get(self, idx: int) -> Data:
+        return torch.load(
+            osp.join(self.processed_dir, self.processed_file_names[idx]),
+            weights_only=False,
+        )
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({len(self)}, '
+                f'mode={self.mode}, split={self.split})')

torch_geometric/nn/aggr/patch_transformer.py

@@ -32,6 +32,8 @@ class PatchTransformerAggregation(Aggregation):
         aggr (str or list[str], optional): The aggregation module, *e.g.*,
             :obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`,
             :obj:`"var"`, :obj:`"std"`. (default: :obj:`"mean"`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -43,6 +45,7 @@ class PatchTransformerAggregation(Aggregation):
         heads: int = 1,
         dropout: float = 0.0,
         aggr: Union[str, List[str]] = 'mean',
+        device: Optional[torch.device] = None,
     ) -> None:
         super().__init__()
 
@@ -55,12 +58,13 @@ class PatchTransformerAggregation(Aggregation):
         for aggr in self.aggrs:
             assert aggr in ['sum', 'mean', 'min', 'max', 'var', 'std']
 
-        self.lin = torch.nn.Linear(in_channels, hidden_channels)
+        self.lin = torch.nn.Linear(in_channels, hidden_channels, device=device)
         self.pad_projector = torch.nn.Linear(
             patch_size * hidden_channels,
             hidden_channels,
+            device=device,
         )
-        self.pe = PositionalEncoding(hidden_channels)
+        self.pe = PositionalEncoding(hidden_channels, device=device)
 
         self.blocks = torch.nn.ModuleList([
             MultiheadAttentionBlock(
@@ -68,12 +72,14 @@ class PatchTransformerAggregation(Aggregation):
                 heads=heads,
                 layer_norm=True,
                 dropout=dropout,
+                device=device,
             ) for _ in range(num_transformer_blocks)
         ])
 
         self.fc = torch.nn.Linear(
             hidden_channels * len(self.aggrs),
             out_channels,
+            device=device,
         )
 
     def reset_parameters(self) -> None:

torch_geometric/nn/aggr/utils.py

@@ -26,9 +26,11 @@ class MultiheadAttentionBlock(torch.nn.Module):
             normalization. (default: :obj:`True`)
         dropout (float, optional): Dropout probability of attention weights.
             (default: :obj:`0`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(self, channels: int, heads: int = 1, layer_norm: bool = True,
-                 dropout: float = 0.0):
+                 dropout: float = 0.0, device: Optional[torch.device] = None):
         super().__init__()
 
         self.channels = channels
@@ -40,10 +42,13 @@ class MultiheadAttentionBlock(torch.nn.Module):
             heads,
             batch_first=True,
             dropout=dropout,
+            device=device,
         )
-        self.lin = Linear(channels, channels)
-        self.layer_norm1 = LayerNorm(channels) if layer_norm else None
-        self.layer_norm2 = LayerNorm(channels) if layer_norm else None
+        self.lin = Linear(channels, channels, device=device)
+        self.layer_norm1 = LayerNorm(channels,
+                                     device=device) if layer_norm else None
+        self.layer_norm2 = LayerNorm(channels,
+                                     device=device) if layer_norm else None
 
     def reset_parameters(self):
         self.attn._reset_parameters()

torch_geometric/nn/conv/meshcnn_conv.py

@@ -1,7 +1,7 @@
 # The below is to suppress the warning on torch.nn.conv.MeshCNNConv::update
 # pyright: reportIncompatibleMethodOverride=false
+import warnings
 from typing import Optional
-from warnings import warn
 
 import torch
 from torch.nn import Linear, Module, ModuleList
@@ -456,13 +456,10 @@ class MeshCNNConv(MessagePassing):
                 {type(network)}"
             if not hasattr(network, "in_channels") and \
                     not hasattr(network, "in_features"):
-                warn(
-                    f"kernel[{i}] does not have attribute \
-                            'in_channels' nor 'out_features'. The \
-                            network must take as input a \
-                            {self.in_channels}-dimensional tensor. \
-                            Still, assuming user configured \
-                            correctly. Continuing..", stacklevel=2)
+                warnings.warn(
+                    f"kernel[{i}] does not have attribute 'in_channels' nor "
+                    f"'out_features'. The network must take as input a "
+                    f"{self.in_channels}-dimensional tensor.", stacklevel=2)
             else:
                 input_dimension = getattr(network, "in_channels",
                                           network.in_features)
@@ -475,13 +472,10 @@ class MeshCNNConv(MessagePassing):
 
             if not hasattr(network, "out_channels") and \
                     not hasattr(network, "out_features"):
-                warn(
-                    f"kernel[{i}] does not have attribute \
-                                'in_channels' nor 'out_features'. The \
-                                network must take as input a \
-                                {self.in_channels}-dimensional tensor. \
-                                Still, assuming user configured \
-                                correctly. Continuing..", stacklevel=2)
+                warnings.warn(
+                    f"kernel[{i}] does not have attribute 'in_channels' nor "
+                    f"'out_features'. The network must take as input a "
+                    f"{self.in_channels}-dimensional tensor.", stacklevel=2)
             else:
                 output_dimension = getattr(network, "out_channels",
                                            network.out_features)

torch_geometric/nn/encoding.py

@@ -1,4 +1,5 @@
 import math
+from typing import Optional
 
 import torch
 from torch import Tensor
@@ -23,12 +24,15 @@ class PositionalEncoding(torch.nn.Module):
         granularity (float, optional): The granularity of the positions. If
             set to smaller value, the encoder will capture more fine-grained
             changes in positions. (default: :obj:`1.0`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
         out_channels: int,
         base_freq: float = 1e-4,
         granularity: float = 1.0,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -40,7 +44,8 @@ class PositionalEncoding(torch.nn.Module):
         self.base_freq = base_freq
         self.granularity = granularity
 
-        frequency = torch.logspace(0, 1, out_channels // 2, base_freq)
+        frequency = torch.logspace(0, 1, out_channels // 2, base_freq,
+                                   device=device)
         self.register_buffer('frequency', frequency)
 
         self.reset_parameters()
@@ -75,13 +80,17 @@ class TemporalEncoding(torch.nn.Module):
 
     Args:
         out_channels (int): Size :math:`d` of each output sample.
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, out_channels: int):
+    def __init__(self, out_channels: int,
+                 device: Optional[torch.device] = None):
         super().__init__()
         self.out_channels = out_channels
 
         sqrt = math.sqrt(out_channels)
-        weight = 1.0 / sqrt**torch.linspace(0, sqrt, out_channels).view(1, -1)
+        weight = 1.0 / sqrt**torch.linspace(0, sqrt, out_channels,
+                                            device=device).view(1, -1)
         self.register_buffer('weight', weight)
 
         self.reset_parameters()

torch_geometric/nn/models/__init__.py

@@ -32,6 +32,7 @@ from .visnet import ViSNet
 from .g_retriever import GRetriever
 from .git_mol import GITMol
 from .molecule_gpt import MoleculeGPT
+from .protein_mpnn import ProteinMPNN
 from .glem import GLEM
 from .sgformer import SGFormer
 # Deprecated:
@@ -86,6 +87,7 @@ __all__ = classes = [
     'GRetriever',
     'GITMol',
     'MoleculeGPT',
+    'ProteinMPNN',
     'GLEM',
     'SGFormer',
     'ARLinkPredictor',

torch_geometric/nn/models/glem.py

@@ -8,6 +8,13 @@ from torch_geometric.loader import DataLoader, NeighborLoader
 from torch_geometric.nn.models import GraphSAGE, basic_gnn
 
 
+def deal_nan(x):
+    if isinstance(x, torch.Tensor):
+        x = x.clone()
+        x[torch.isnan(x)] = 0.0
+    return x
+
+
 class GLEM(torch.nn.Module):
     r"""This GNN+LM co-training model is based on GLEM from the `"Learning on
     Large-scale Text-attributed Graphs via Variational Inference"
@@ -379,9 +386,6 @@ class GLEM(torch.nn.Module):
             is_augmented: use EM or just train GNN and LM with gold data
 
         """
-        def deal_nan(x):
-            return 0 if torch.isnan(x) else x
-
         if is_augmented and (sum(~is_gold) > 0):
             mle_loss = deal_nan(loss_func(logits[is_gold], labels[is_gold]))
             # all other labels beside from ground truth(gold labels)