py3dcal 1.0.0__py3-none-any.whl

py3DCal/model_training/lib/validate_device.py

@@ -0,0 +1,22 @@
+import torch
+
+def validate_device(device: str):
+        """
+        Validates the device by converting it to a torch.device object.
+        Args:
+            device (str): Device to run the model on.
+        Returns:
+            None.
+        Raises:
+            ValueError: If the device is not specified or invalid.
+        """
+        try:
+            device = torch.device(device)
+        except Exception as e:
+            raise ValueError(
+                f"Invalid device '{device}'. Valid options include:\n"
+                "  - 'cpu': CPU processing\n"
+                "  - 'cuda' or 'cuda:0': NVIDIA GPU\n"
+                "  - 'mps': Apple Silicon GPU\n"
+                "See: https://pytorch.org/docs/stable/tensor_attributes.html#torch.device"
+            ) from e

py3DCal/model_training/lib/validate_parameters.py

@@ -0,0 +1,45 @@
+import torch
+from typing import Union
+from pathlib import Path
+
+def validate_device(device: str):
+        """
+        Validates the device by converting it to a torch.device object.
+        Args:
+            device (str): Device to run the model on.
+        Returns:
+            None.
+        Raises:
+            ValueError: If the device is not specified or invalid.
+        """
+        try:
+            device = torch.device(device)
+        except Exception as e:
+            raise ValueError(
+                f"Invalid device '{device}'. Valid options include:\n"
+                "  - 'cpu': CPU processing\n"
+                "  - 'cuda' or 'cuda:0': NVIDIA GPU\n"
+                "  - 'mps': Apple Silicon GPU\n"
+                "See: https://pytorch.org/docs/stable/tensor_attributes.html#torch.device"
+            ) from e
+        
+def validate_root(root):
+    """
+    Validates the root path specified by the user.
+
+    Args:
+        root: root path specified by the user.
+    Returns:
+        None.
+    Raises:
+        ValueError: If the root is not specified or invalid.
+    """
+    if root is None :
+       raise ValueError(
+           "root directory cannot be None.\n"
+       )
+    
+    if not isinstance(root, (str, Path)):
+       raise ValueError(
+           "root directory must be a valid file system path as a string or pathlib.Path object\n"
+       )

py3DCal/model_training/models/__init__.py

@@ -0,0 +1 @@
+from .touchnet import SensorType, TouchNet

py3DCal/model_training/models/touchnet.py

@@ -0,0 +1,211 @@
+import os
+import requests
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from enum import Enum
+from typing import Union
+from pathlib import Path
+
+class SensorType(Enum):
+    """
+    SensorType: Available sensor types with pretrained weights and compiled datasets
+    """
+    DIGIT = "DIGIT"
+    GELSIGHTMINI = "GelSightMini"
+
+class TouchNet(nn.Module):
+    """
+    TouchNet: A PyTorch neural network for producing surface normal maps from tactile sensor images.
+
+    Args:
+        load_pretrained (bool): If True, loads pretrained weights for the specified sensor type.
+        sensor_type (SensorType): The type of tactile sensor. Must be specified if load_pretrained is True.
+        root (str or pathlib.Path): The root directory for saving/loading the pretrained_weights (.pth) file.
+    """
+    def __init__(self, load_pretrained: bool = False, sensor_type: SensorType = None, root: Union[str, Path] = "."):
+        super().__init__()
+
+        self._validate_parameters(load_pretrained, sensor_type, root)
+
+        self.conv1 = nn.Conv2d(5, 32, kernel_size=7, padding=3)
+        self.bn1 = nn.BatchNorm2d(32)
+        self.dropout1 = nn.Dropout2d(0.2)
+
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=7, padding=3)
+        self.bn2 = nn.BatchNorm2d(64)
+        self.dropout2 = nn.Dropout2d(0.2)
+
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=7, padding=3)
+        self.bn3 = nn.BatchNorm2d(128)
+        self.dropout3 = nn.Dropout2d(0.2)
+
+        self.conv4 = nn.Conv2d(128, 256, kernel_size=5, padding=2)
+        self.bn4 = nn.BatchNorm2d(256)
+        self.dropout4 = nn.Dropout2d(0.3)
+
+        self.conv5 = nn.Conv2d(256, 256, kernel_size=5, padding=2)
+        self.bn5 = nn.BatchNorm2d(256)
+        self.dropout5 = nn.Dropout2d(0.3)
+
+        self.conv6 = nn.Conv2d(256, 128, kernel_size=5, padding=2)
+        self.bn6 = nn.BatchNorm2d(128)
+        self.dropout6 = nn.Dropout2d(0.2)
+
+        self.conv7 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+        self.bn7 = nn.BatchNorm2d(64)   
+        self.dropout7 = nn.Dropout2d(0.2)
+
+        self.conv8 = nn.Conv2d(64, 32, kernel_size=3, padding=1)
+        self.bn8 = nn.BatchNorm2d(32)
+        self.dropout8 = nn.Dropout2d(0.2)
+
+        self.conv9 = nn.Conv2d(32, 2, kernel_size=1)
+
+        if load_pretrained:
+            self._load_pretrained_model(root, sensor_type)
+    
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = self.dropout1(x)
+
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = self.dropout2(x)
+
+        x = F.relu(self.bn3(self.conv3(x)))
+        x = self.dropout3(x)
+
+        x = F.relu(self.bn4(self.conv4(x)))
+        x = self.dropout4(x)
+
+        x = F.relu(self.bn5(self.conv5(x)))
+        x = self.dropout5(x)
+
+        x = F.relu(self.bn6(self.conv6(x)))
+        x = self.dropout6(x)
+
+        x = F.relu(self.bn7(self.conv7(x)))
+        x = self.dropout7(x)
+
+        x = F.relu(self.bn8(self.conv8(x)))
+        x = self.dropout8(x)
+
+        x = self.conv9(x)
+
+        return x
+    
+    def _validate_parameters(self, load_pretrained, sensor_type, root):
+        if load_pretrained and sensor_type is None:
+            raise ValueError("sensor_type must be specified when load_pretrained is True. sensor_type must be either SensorType.DIGIT or SensorType.GELSIGHTMINI.")
+
+        if load_pretrained and not isinstance(sensor_type, SensorType):
+            raise ValueError("sensor_type must be either SensorType.DIGIT or SensorType.GELSIGHTMINI.")
+        
+        if load_pretrained and root is None:
+            raise ValueError("root directory for storing/loading model cannot be None when load_pretrained is True.")
+        
+        if load_pretrained and not isinstance(root, (str, Path)):
+            raise ValueError("root directory must be a valid file system path as a string or pathlib.Path object when load_pretrained is True.")
+
+        if not load_pretrained and sensor_type is not None:
+            print("Warning: sensor_type parameter is ignored when load_pretrained is False.")
+
+        if not load_pretrained and root is not ".":
+            print("Warning: root parameter is ignored when load_pretrained is False.")
+
+    def _load_pretrained_model(self, root, sensor_type):
+        """
+        Loads a pretrained model for either the DIGIT or GelSightMini sensor.
+        Args:
+            None.
+        Returns:
+            None.
+        """
+        
+        if sensor_type == SensorType.DIGIT:
+            file_path = os.path.join(root, "digit_pretrained_weights.pth")
+
+            # Check if DIGIT pretrained weights exist locally, if not download them
+            if not os.path.exists(file_path):
+
+                print(f"Downloading DIGIT pretrained weights ...")
+                response = requests.get('https://zenodo.org/records/17517028/files/digit_pretrained_weights.pth?download=1', stream=True)
+                response.raise_for_status()
+
+                total_size = int(response.headers.get('content-length', 0))
+                block_size = 1024
+
+                # Save file in chunks to handle large datasets
+                with open(file_path, 'wb') as f, tqdm(
+                    total=total_size,
+                    unit='B',
+                    unit_scale=True,
+                    desc="Downloading",
+                    ncols=80
+                ) as progress_bar:
+                    for chunk in response.iter_content(chunk_size=block_size):
+                        if chunk:
+                            f.write(chunk)
+                            progress_bar.update(len(chunk))
+
+                print(f"Download complete!")
+            else:
+                print(f"DIGIT pretrained weights already exists at: {file_path}/")
+
+        elif sensor_type == SensorType.GELSIGHTMINI:
+            file_path = os.path.join(root, "gsmini_pretrained_weights.pth")
+
+            # Check if GelSight Mini pretrained weights exist locally, if not download them
+            if not os.path.exists(file_path):
+
+                print(f"Downloading GelSight Mini pretrained weights ...")
+                response = requests.get('https://zenodo.org/records/17517028/files/gsmini_pretrained_weights.pth?download=1', stream=True)
+                response.raise_for_status()
+
+                total_size = int(response.headers.get('content-length', 0))
+                block_size = 1024
+
+                # Save file in chunks to handle large datasets
+                with open(file_path, 'wb') as f, tqdm(
+                    total=total_size,
+                    unit='B',
+                    unit_scale=True,
+                    desc="Downloading",
+                    ncols=80
+                ) as progress_bar:
+                    for chunk in response.iter_content(chunk_size=block_size):
+                        if chunk:
+                            f.write(chunk)
+                            progress_bar.update(len(chunk))
+
+                print(f"Download complete!")
+            else:
+                print(f"GelSight Mini pretrained weights already exists at: {file_path}/")
+
+        state_dict = torch.load(file_path, map_location="cpu")
+
+        self.load_state_dict(state_dict)
+
+    def load_weights(self, weights_path: Union[str, Path]):
+        """
+        Loads model weights from a specified .pth file.
+
+        Args:
+            weights_path (str or pathlib.Path): The file path to the .pth file containing the model weights.
+        Returns:
+            None.
+        Raises:
+            ValueError: If the weights_path is not specified or invalid.
+        """
+        if weights_path is None:
+            raise ValueError("weights_path cannot be None.")
+
+        if not isinstance(weights_path, (str, Path)):
+            raise ValueError("weights_path must be a valid file system path as a string or pathlib.Path object.")
+
+        if not os.path.exists(weights_path):
+            raise ValueError(f"The specified weights_path does not exist: {weights_path}")
+
+        state_dict = torch.load(weights_path, map_location="cpu")
+        self.load_state_dict(state_dict)

py3DCal/model_training/touchnet/dataset.py

@@ -0,0 +1,78 @@
+import os
+import sys
+import torch
+import numpy as np
+import pandas as pd
+from PIL import Image
+from torch.utils.data import Dataset
+from ..lib.precompute_gradients import precompute_gradients
+from ..lib.get_gradient_map import get_gradient_map
+
+
+class TactileSensorDataset(Dataset):
+    """
+    Tactile Sensor dataset.
+    """
+    def __init__(self, dataset_path, annotation_path, blank_image_path, transform=None, radius=36):
+        self.dataset_path = dataset_path
+        self.annotation_path = annotation_path
+        self.blank_image_path = blank_image_path
+        self.transform = transform
+        self.csv_file = dataset_path  # Store the CSV filename
+        
+        # Load the CSV data
+        self.data = pd.read_csv(annotation_path, comment='#')
+        
+        self.precomputed_gradients = precompute_gradients(dataset_path, annotation_path, r=radius)
+        self.blank_tensor = self.transform(Image.open(blank_image_path).convert("RGB"))
+    
+    def _add_coordinate_channels(self, image):
+        """
+        Add x and y coordinate channels to the input image.
+        X channel: column indices (1s in first column, 2s in second column, etc.)
+        Y channel: row indices (1s in first row, 2s in second row, etc.)
+        """
+        # Get image dimensions
+        _, height, width = image.shape
+        
+        # Create x coordinate channel (column indices)
+        x_coords = torch.arange(1, width + 1, dtype=torch.float32).unsqueeze(0).repeat(height, 1)
+        x_channel = x_coords.unsqueeze(0)  # Add channel dimension
+        
+        # Create y coordinate channel (row indices)
+        y_coords = torch.arange(1, height + 1, dtype=torch.float32).unsqueeze(1).repeat(1, width)
+        y_channel = y_coords.unsqueeze(0)  # Add channel dimension
+        
+        # Concatenate original image with coordinate channels
+        image_with_coords = torch.cat([image, x_channel, y_channel], dim=0)
+        
+        return image_with_coords
+        
+    def __len__(self):
+        return len(self.data)  # Use the DataFrame length
+
+    def __getitem__(self, idx):
+        # Check if index is valid
+        if idx < 0 or idx >= len(self.data):
+            raise IndexError("Index out of range")
+        
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+
+        img_name = os.path.join(self.dataset_path, self.data.iloc[idx, 0])
+        img = Image.open(img_name).convert("RGB")
+        image = np.array(img, copy=True)
+
+        target = get_gradient_map(idx, annotation_path=self.annotation_path, precomputed_gradients=self.precomputed_gradients)
+
+        # if self.transform:
+        image = self.transform(image)
+        target = self.transform(target)
+
+        # Subtract pre-transformed blank tensor
+        image = image - self.blank_tensor
+        # Add coordinate channels
+        image = self._add_coordinate_channels(image)
+
+        sample = (image, target)
+        return sample