py3dcal 1.0.5__py3-none-any.whl

py3DCal/data_collection/sensors/Sensor.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+
+class Sensor(ABC):
+    """
+    Sensor: An abstract base class for tactile sensors.
+    """
+    def __init__(self):
+        self.name = ""
+        self.x_offset = 5
+        self.y_offset = 5
+        self.z_offset = 5
+        self.z_clearance = 2
+        self.max_penetration = 0
+        self.default_calibration_file = "calibration_procs/digit/default.csv"
+
+    @abstractmethod
+    def connect(self):
+        """ Connects to the sensor
+        """
+        pass
+
+    @abstractmethod
+    def disconnect(self):
+        """ Disconnects from the sensor
+        """
+        pass
+
+    @abstractmethod
+    def capture_image(self):
+        """ Captures an image from the sensor
+
+        Returns:
+            numpy.ndarray: The image from the sensor.
+        """
+        pass
+
+    def flush_frames(self, n: int = 5):
+        """Discards the next n frames to clear camera buffer.
+        
+        Args:
+            n (int): Number of frames to discard. Default is 5.
+        """
+        for _ in range(n):
+            self.capture_image()

py3DCal/model_training/datasets/DIGIT_dataset.py

@@ -0,0 +1,77 @@
+import os
+import requests
+import tarfile
+from tqdm import tqdm
+from typing import Union
+from pathlib import Path
+from .tactile_sensor_dataset import TactileSensorDataset
+from ..lib.validate_parameters import validate_root
+
+
+class DIGIT(TactileSensorDataset):
+    """
+    DIGIT: A Dataset Class for the DIGIT sensor
+    Args:
+        root (str or pathlib.Path): The root directory containing digit_calibration_data.
+        download (bool, optional): If True, downloads the dataset for the specified sensor type. Defaults to False.
+        add_coordinate_embeddings (bool, optional): If True, adds xy coordinate embeddings to each image. Defaults to True.
+        subtract_blank (bool, optional): If True, subtracts a blank image from each input image. Defaults to False.
+        transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. Default: ``transforms.ToTensor()``
+    """
+    def __init__(self, root: Union[str, Path] = Path("."), download=False, add_coordinate_embeddings=True, subtract_blank=True, transform=None):
+        validate_root(root)
+
+        self.root = root
+
+        self.dataset_path = os.path.join(self.root, "digit_calibration_data")
+
+        if download:
+            self._download_dataset()
+
+        super().__init__(root=self.dataset_path, add_coordinate_embeddings=add_coordinate_embeddings, subtract_blank=subtract_blank, transform=transform)
+
+    def _download_dataset(self):
+        """
+        Downloads the dataset for either the DIGIT sensor.
+
+        """
+        # Check if self.dataset_path exists
+        if not os.path.exists(self.dataset_path):
+            os.makedirs(self.root, exist_ok=True)
+
+            tar_path = os.path.join(self.root, "digit_calibration_data.tar.gz")
+
+            print(f"Downloading DIGIT dataset ...")
+            response = requests.get('https://zenodo.org/records/17517028/files/digit_calibration_data.tar.gz?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(tar_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!")
+
+            # Extract .tar.gz file
+            print("Extracting files ...")
+            with tarfile.open(tar_path, "r:gz") as tar:
+                tar.extractall(path=self.root)
+
+            os.remove(tar_path)
+
+            print(f"Extraction complete! Files are in: {self.root}/")
+
+        else:
+            print(f"DIGIT dataset already exists at: {self.dataset_path}/")
+

py3DCal/model_training/datasets/GelSightMini_dataset.py

@@ -0,0 +1,75 @@
+import os
+import requests
+import tarfile
+from tqdm import tqdm
+from typing import Union
+from pathlib import Path
+from .tactile_sensor_dataset import TactileSensorDataset
+from ..lib.validate_parameters import validate_root
+
+
+class GelSightMini(TactileSensorDataset):
+    """
+    GelSight Mini: A Dataset Class for the GelSight Mini sensor
+    Args:
+        root (str or pathlib.Path): The root directory containing gsmini_calibration_data.
+        download (bool, optional): If True, downloads the dataset for the specified sensor type. Defaults to False.
+        add_coordinate_embeddings (bool, optional): If True, adds xy coordinate embeddings to each image. Defaults to True.
+        subtract_blank (bool, optional): If True, subtracts a blank image from each input image. Defaults to False.
+        transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. Default: ``transforms.ToTensor()``
+    """
+    def __init__(self, root: Union[str, Path] = Path("."), download=False, add_coordinate_embeddings=True, subtract_blank=True, transform=None):
+        validate_root(root)
+
+        self.root = root
+
+        self.dataset_path = os.path.join(self.root, "gsmini_calibration_data")
+
+        if download:
+            self._download_dataset()
+
+        super().__init__(root=self.dataset_path, add_coordinate_embeddings=add_coordinate_embeddings, subtract_blank=subtract_blank, transform=transform)
+
+    def _download_dataset(self):
+        """
+        Downloads the dataset for the GelSight Mini sensor.
+        """
+        # Check if self.dataset_path exists
+        if not os.path.exists(self.dataset_path):
+            os.makedirs(self.root, exist_ok=True)
+            
+            tar_path = os.path.join(self.root, "gsmini_calibration_data.tar.gz")
+
+            print(f"Downloading GelSight Mini dataset ...")
+            response = requests.get('https://zenodo.org/records/17517028/files/gsmini_calibration_data.tar.gz?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(tar_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!")
+
+            # Extract .tar.gz file
+            print("Extracting files ...")
+            with tarfile.open(tar_path, "r:gz") as tar:
+                tar.extractall(path=self.root)
+
+            os.remove(tar_path)
+
+            print(f"Extraction complete! Files are in: {self.root}/")
+
+        else:
+            print(f"GelSight Mini dataset already exists at: {self.dataset_path}/")

py3DCal/model_training/datasets/__init__.py

@@ -0,0 +1,3 @@
+from .tactile_sensor_dataset import TactileSensorDataset
+from .DIGIT_dataset import DIGIT
+from .GelSightMini_dataset import GelSightMini

py3DCal/model_training/datasets/split_dataset.py

@@ -0,0 +1,38 @@
+import copy
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from .tactile_sensor_dataset import TactileSensorDataset
+
+def split_dataset(dataset, train_ratio=0.8):
+    """
+    Splits a dataset into training and validation sets.
+
+    Args:
+        dataset (py3DCal.datasets.TactileSensorDataset): The dataset to split.
+        train_ratio (float): The proportion of the dataset to include in the training set. Default is 0.8.
+
+    Returns:
+        tuple: A tuple containing the training and validation datasets.
+    """
+    if not isinstance(dataset, TactileSensorDataset):
+        raise TypeError("Expected dataset to be an instance of py3DCal.datasets.TactileSensorDataset")
+
+    df = dataset.data.copy()
+
+    unique_coords = df[['x_mm', 'y_mm']].drop_duplicates().reset_index(drop=True)
+
+    train_df, val_df = train_test_split(unique_coords, train_size=train_ratio, random_state=42)
+
+    # Merge back to get full rows
+    train_data = pd.merge(df, train_df, on=['x_mm', 'y_mm'])
+    val_data = pd.merge(df, val_df, on=['x_mm', 'y_mm'])
+
+    # Create two copies of the original dataset
+    train_dataset = copy.deepcopy(dataset)
+    val_dataset = copy.deepcopy(dataset)
+
+    # Update the data attribute of each dataset
+    train_dataset.data = train_data
+    val_dataset.data = val_data
+
+    return train_dataset, val_dataset

py3DCal/model_training/datasets/tactile_sensor_dataset.py

@@ -0,0 +1,83 @@
+import os
+import torch
+import pandas as pd
+from PIL import Image
+from typing import Union
+from pathlib import Path
+from torch.utils.data import Dataset
+from torchvision import transforms
+from ..lib.precompute_gradients import precompute_gradients
+from ..lib.get_gradient_map import get_gradient_map
+from ..lib.add_coordinate_embeddings import add_coordinate_embeddings
+from ..lib.validate_parameters import validate_root, validate_dataset
+
+class TactileSensorDataset(Dataset):
+    """
+    Tactile Sensor Dataset.
+
+    Args:
+        root (str or pathlib.Path): The root directory that contains the dataset folder.
+        add_coordinate_embeddings (bool, optional): If True, adds xy coordinate embeddings to each image. Defaults to True.
+        subtract_blank (bool, optional): If True, subtracts a blank image from each input image. Defaults to False.
+        transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version. Default: ``transforms.ToTensor()``
+    """
+    def __init__(self, root: Union[str, Path], add_coordinate_embeddings=True, subtract_blank=True, transform=None):
+        validate_dataset(root, subtract_blank)
+
+        self.root = root
+        self.annotation_path = os.path.join(root, "annotations", "annotations.csv")
+        self.metadata_path = os.path.join(root, "annotations", "metadata.json")
+        if subtract_blank:
+            self.blank_image_path = os.path.join(root, "blank_images", "blank.png")
+        self.add_coordinate_embeddings = add_coordinate_embeddings
+        self.subtract_blank = subtract_blank
+
+        if transform is None:
+            self.transform = transforms.ToTensor()
+        else:
+            self.transform = transform
+
+        # Load the CSV data
+        self.data = pd.read_csv(self.annotation_path)
+
+        # Get probe radius (in px) from metadata
+        metadata = pd.read_json(self.metadata_path, typ="series")
+        radius = metadata["probe_radius_mm"] * metadata["px_per_mm"]
+
+        # Precompute gradients
+        self.precomputed_gradients = precompute_gradients(dataset_path=self.root, annotation_path=self.annotation_path, r=radius)
+
+        # Load and transform blank image
+        if subtract_blank:
+            self.blank_image = self.transform(Image.open(self.blank_image_path).convert("RGB"))
+
+    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()
+
+        image_name = os.path.join(self.root, "probe_images", self.data.iloc[idx, 0])
+        image = Image.open(image_name).convert("RGB")
+
+        target = get_gradient_map(idx, annotation_path=self.annotation_path, precomputed_gradients=self.precomputed_gradients)
+
+        image = self.transform(image)
+        target = self.transform(target)
+
+        if self.subtract_blank:
+            # Subtract pre-transformed blank tensor
+            image = image - self.blank_image
+
+        if self.add_coordinate_embeddings:
+            # Add coordinate embeddings
+            image = add_coordinate_embeddings(image)
+
+        sample = (image, target)
+
+        return sample

py3DCal/model_training/lib/add_coordinate_embeddings.py

@@ -0,0 +1,29 @@
+import torch
+
+def add_coordinate_embeddings(image):
+        """
+        Add coordinate embeddings 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.)
+
+        Args:
+            image (torch.Tensor): Input image tensor of shape (C, H, W).
+
+        Returns:
+            torch.Tensor: Image tensor with added coordinate embeddings of shape (C+2, H, W).
+        """
+        # Get image dimensions
+        _, height, width = image.shape
+        
+        # Create x coordinate channel (column indices)
+        x_embedding = torch.arange(1, width + 1, dtype=torch.float32).unsqueeze(0).repeat(height, 1)
+        x_channel = x_embedding.unsqueeze(0)  # Add channel dimension
+
+        # Create y coordinate channel (row indices)
+        y_embedding = torch.arange(1, height + 1, dtype=torch.float32).unsqueeze(1).repeat(1, width)
+        y_channel = y_embedding.unsqueeze(0)  # Add channel dimension
+
+        # Concatenate original image with position embeddings
+        image_with_embeddings = torch.cat([image, x_channel, y_channel], dim=0)
+
+        return image_with_embeddings