py3dcal 1.0.0__py3-none-any.whl → 1.0.6__py3-none-any.whl

py3DCal/__init__.py

@@ -1,12 +1,14 @@
 from .data_collection.Calibrator import Calibrator
-from .data_collection.Printers.Printer import Printer
-from .data_collection.Sensors.Sensor import Sensor
-from .data_collection.Printers.Ender3.Ender3 import Ender3
-from .data_collection.Sensors.DIGIT.DIGIT import DIGIT
-from .data_collection.Sensors.GelsightMini.GelsightMini import GelsightMini
+from .data_collection.printers.Printer import Printer
+from .data_collection.printers.Ender3.Ender3 import Ender3
+from .data_collection.sensors.Sensor import Sensor
+from .data_collection.sensors.DIGIT.DIGIT import DIGIT
+from .data_collection.sensors.GelsightMini.GelsightMini import GelsightMini
 from .model_training import datasets, models
 from .model_training.datasets.split_dataset import split_dataset
 from .model_training.models.touchnet import SensorType
+from .model_training.lib.annotate_dataset import annotate, visualize_annotations
 from .model_training.lib.train_model import train_model
 from .model_training.lib.depthmaps import get_depthmap, save_2d_depthmap, show_2d_depthmap
+from .model_training.lib.fast_poisson import fast_poisson
 from .utils.utils import list_com_ports

py3DCal/data_collection/Calibrator.py

@@ -6,9 +6,8 @@ from typing import Union
 from pathlib import Path
 from PIL import Image
 from tqdm import tqdm
-
-from .Printers.Printer import Printer
-from .Sensors.Sensor import Sensor
+from .printers.Printer import Printer
+from .sensors.Sensor import Sensor
 
 class Calibrator:
     """ Calibrator class to automatically probe a tactile sensor.
@@ -153,7 +152,7 @@ class Calibrator:
         # Connect to 3D printer if not already connected
         if not self.printer_connected:
             self.connect_printer()
-            self.printer.send_gcode("M117 Sensor Calibration In Progress")
+            self.printer.send_gcode("M117 Calibrating sensor...")
 
         # Connect to sensor
         if save_images == True:
@@ -257,11 +256,14 @@ class Calibrator:
                     with open(os.path.join(data_save_path, "annotations", "probe_data.csv"), 'a', newline='') as csv_file:
                         csv_writer = csv.writer(csv_file)
 
-                        for j in range(int(self.calibration_points[i][3])):
+                        # Flush frames to clear camera buffer
+                        self.sensor.flush_frames(n=5)
+
+                        for _ in range(int(self.calibration_points[i][3])):
                             frame = self.sensor.capture_image()
 
                             img_name =  str(img_idx) + "_" + "X" + str(self.calibration_points[i][0]) + "Y" + str(self.calibration_points[i][1]) + "Z" + str(self.calibration_points[i][2]) + ".png"
-                            img_path = os.path.join(data_save_path, "probe_images",img_name)
+                            img_path = os.path.join(data_save_path, "probe_images", img_name)
 
                             img = Image.fromarray(frame)
                             img.save(img_path)
@@ -283,7 +285,7 @@ class Calibrator:
         print("")
 
         # Update printer display
-        self.printer.send_gcode("M117 Sensor Calibration Complete!")
+        self.printer.send_gcode("M117 Calibration Done!")
 
         # Disconnect from 3D printer
         self.disconnect_printer()
@@ -292,7 +294,7 @@ class Calibrator:
         if save_images == True:
             self.disconnect_sensor()
 
-        print("Sensor calibration procedure complete!")
+        print("Calibration procedure complete!")
         print("")
 
         return True

py3DCal/data_collection/{Sensors → sensors}/GelsightMini/GelsightMini.py

@@ -10,10 +10,10 @@ except:
 
 class GelsightMini(Sensor):
     """
-    GelsightMini: A Sensor Class for the Gelsight Mini sensor
+    GelsightMini: A Sensor Class for the GelSight Mini sensor
     """
     def __init__(self):
-        self.name = "Gelsight Mini"
+        self.name = "GelSight Mini"
         self.x_offset = 108
         self.y_offset = 110
         self.z_offset = 67
@@ -23,7 +23,7 @@ class GelsightMini(Sensor):
 
     def connect(self):
         """
-        Connects to the Gelsight Mini sensor.
+        Connects to the GelSight Mini sensor.
         """
         # Code to connect to the sensor
         self.sensor = gsdevice.Camera("GelSight Mini")
@@ -31,14 +31,14 @@ class GelsightMini(Sensor):
         
     def disconnect(self):
         """
-        Disconnects from the Gelsight Mini sensor.
+        Disconnects from the GelSight Mini sensor.
         """
         # Code to disconnect from the sensor
         self.sensor.stop_video()
 
     def capture_image(self):
         """
-        Captures an image from the Gelsight Mini sensor.
+        Captures an image from the GelSight Mini sensor.
         """
         # Code to return an image from the sensor
         image = cv2.cvtColor(self.sensor.get_image(), cv2.COLOR_BGR2RGB)

py3DCal/data_collection/{Sensors → sensors}/Sensor.py

@@ -32,4 +32,13 @@ class Sensor(ABC):
         Returns:
             numpy.ndarray: The image from the sensor.
         """
-        pass
+        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

@@ -21,13 +21,15 @@ class DIGIT(TactileSensorDataset):
     def __init__(self, root: Union[str, Path] = Path("."), download=False, add_coordinate_embeddings=True, subtract_blank=True, transform=None):
         validate_root(root)
 
-        self.dataset_path = os.path.join(root, "digit_calibration_data")
+        self.root = root
 
-        super().__init__(root=self.dataset_path, add_coordinate_embeddings=add_coordinate_embeddings, subtract_blank=subtract_blank, transform=transform)
+        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.

py3DCal/model_training/datasets/GelSightMini_dataset.py

@@ -21,13 +21,15 @@ class GelSightMini(TactileSensorDataset):
     def __init__(self, root: Union[str, Path] = Path("."), download=False, add_coordinate_embeddings=True, subtract_blank=True, transform=None):
         validate_root(root)
 
-        self.dataset_path = os.path.join(root, "gsmini_calibration_data")
+        self.root = root
 
-        super().__init__(root=self.dataset_path, add_coordinate_embeddings=add_coordinate_embeddings, subtract_blank=subtract_blank, transform=transform)
+        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.

py3DCal/model_training/datasets/tactile_sensor_dataset.py

@@ -9,7 +9,7 @@ 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
+from ..lib.validate_parameters import validate_root, validate_dataset
 
 class TactileSensorDataset(Dataset):
     """
@@ -22,12 +22,13 @@ class TactileSensorDataset(Dataset):
         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_root(root)
+        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")
-        self.blank_image_path = os.path.join(root, "blank_images", "blank.png")
+        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