py3dcal 1.0.5__py3-none-any.whl

py3DCal/__init__.py

@@ -0,0 +1,14 @@
+from .data_collection.Calibrator import Calibrator
+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
+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

@@ -0,0 +1,300 @@
+import numpy as np
+import time
+import csv
+import os
+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
+
+class Calibrator:
+    """ Calibrator class to automatically probe a tactile sensor.
+    Args:
+        printer (Printer): An instance of a Printer class.
+        sensor (Sensor): An instance of a Sensor class.
+    """
+    def __init__(self, printer: Printer, sensor: Sensor):
+        self.printer = printer
+        self.sensor = sensor
+
+        self.printer_connected = False
+        self.sensor_connected = False
+
+        try:
+            self.printer_name = self.printer.name
+        except:
+            self.printer_name = "printer"
+        try:
+            self.sensor_name = self.sensor.name
+        except:
+            self.sensor_name = "tactile"
+    
+    def connect_printer(self):
+        """ Connects to the 3D Printer
+
+        Returns:
+            bool: Returns True if connection was successful.
+        """
+        print("Connecting to " + str(self.printer_name) + "...")
+
+        try:
+            self.printer.connect()
+            self.printer_connected = True
+
+            print("Connected to " + str(self.printer_name) + "!")
+            print("")
+            return True
+        except:
+            self.printer_connected = False
+            print("Error connecting to " + str(self.printer_name) + ".")
+            print("")
+            return False
+        
+    def disconnect_printer(self):
+        """ Disconnects from the 3D Printer
+
+        Returns:
+            bool: Returns True if disconnection was successful.
+        """
+        print("Disconnecting from " + str(self.printer_name) + "...")
+
+        try:
+            self.printer.disconnect()
+            self.printer_connected = False
+            print("Disconnected from " + str(self.printer_name) + "!")
+            print("")
+            return True
+        except:
+            print("Error disconnecting from " + str(self.printer_name) + ".")
+            print("")
+            return False
+
+    def initialize_printer(self):
+        """ Sends gcode to configure and home 3D Printer
+
+        Returns:
+            bool: Returns True if initialization was successful.
+        """
+        # Probe must be detached to home printer
+        print("Make sure probe is not attached to print head. ", end="")
+        input("Press Enter to continue...")
+        print("")
+        
+        if not self.printer_connected:
+            self.connect_printer()
+        
+        try:
+            print("Initializing printer...")
+
+            self.printer.initialize()
+
+            print("Printer initialization complete!")
+            print("")
+            return True
+        except:
+            print("Error sending initialization gcode to printer.")
+            print("")
+            return False
+        
+    def connect_sensor(self):
+        """ Connects to the sensor
+
+        Returns:
+            bool: Returns True if connection was successful.
+        """
+        print("Connecting to " + self.sensor_name + " sensor...")
+
+        try:
+            # Connect to sensor
+            self.sensor.connect()
+            print("Connected to sensor!")
+            print("")
+        except:
+            print("Error connecting to sensor.")
+            print("")
+
+    def disconnect_sensor(self):
+        """ Disconnects from the sensor
+
+        Returns:
+            bool: Returns True if disconnection was successful.
+        """
+        print("Disconnecting from " + self.sensor.name + " sensor...")
+
+        self.sensor.connect()
+        try:
+            self.printer.disconnect()
+            self.printer_connected = False
+            print("Disconnected from sensor!")
+            print("")
+            return True
+        except:
+            print("Error disconnecting from sensor.")
+            print("")
+            return False
+
+    def probe(self, home_printer: bool = True, save_images: bool = True, calibration_file_path: Union[str, Path] = None, data_save_path: str = "."):
+        """ Executes the probing procedure on 3D printer
+
+        Args:
+            home_printer (bool, optional): Determines whether to home the printer prior to probing. Defaults to True.
+            save_images (bool, optional): Determines whether sensor images are saved. Defaults to True.
+            calibration_file_path (str, optional): The path of the calibration file. For the DIGIT and GelSight Mini,
+                if no file is specified, a default calibration file will be used.
+            data_save_path (str, optional): The folder in which the data should be saved. If no folder is specified,
+                data will be stored in a directory named "sensor_calibration_data" within the current working directory.
+
+        Returns:
+            bool: Returns True when the probing procedure is complete.
+        """
+        # Connect to 3D printer if not already connected
+        if not self.printer_connected:
+            self.connect_printer()
+            self.printer.send_gcode("M117 Calibrating sensor...")
+
+        # Connect to sensor
+        if save_images == True:
+            self.connect_sensor()
+
+            for i in range(30):
+                self.sensor.capture_image()
+
+        # Send initialization gcode to printer
+        if home_printer == True:
+            self.initialize_printer()
+
+        # If no data path was provided, set default path to a folder called "sensor_calibration_data" in the Downloads folder
+        if data_save_path is not None:
+            data_save_path = os.path.join(data_save_path, "sensor_calibration_data")
+        
+        # Create necessary directories if they don't exist
+        Path(data_save_path).mkdir(parents=True, exist_ok=True)
+        Path(os.path.join(data_save_path, "annotations")).mkdir(parents=True, exist_ok=True)
+        Path(os.path.join(data_save_path, "blank_images")).mkdir(parents=True, exist_ok=True)
+        Path(os.path.join(data_save_path, "probe_images")).mkdir(parents=True, exist_ok=True)
+
+        # Open a csv file to write calibration data
+        with open(os.path.join(data_save_path, "annotations", "probe_data.csv"), 'w', newline='') as csv_file:
+            csv_writer = csv.writer(csv_file)
+            csv_writer.writerow(['img_name', 'x_mm', 'y_mm', 'penetration_depth_mm'])
+
+        # Save blank image
+        if save_images == True:
+            blank_img = self.sensor.capture_image()
+            img = Image.fromarray(blank_img)
+            img.save(os.path.join(data_save_path, "blank_images", "blank.png"))
+
+        # If no calibration file path was provided, use the default calibration file for the specified sensor
+        if calibration_file_path == None:
+            calibration_file_path = self.sensor.default_calibration_file
+
+        # Load CSV file into numpy array
+        self.calibration_points = np.genfromtxt(calibration_file_path, delimiter=',', skip_header=1)
+
+        # Get number of rows (i.e. calibration points)
+        N = self.calibration_points.shape[0]
+
+        # Variable to keep track of image number
+        img_idx = 0
+
+        # Move to offset Z position
+        self.printer.send_gcode("G0  Z" + str(self.sensor.z_offset + self.sensor.z_clearance))
+
+        time.sleep(1 + (self.sensor.z_offset + self.sensor.z_clearance) / 4)
+
+        print("Attach probe to printer head. ", end="")
+        input("Press Enter to continue...")
+        print("")
+
+        # Move to offset XY position
+        self.printer.send_gcode("G0 X" + str(self.sensor.x_offset) + " Y" + str(self.sensor.y_offset))
+
+        time.sleep(1 + max(self.sensor.x_offset, self.sensor.y_offset) / 10)
+
+        print("Beginning sensor calibration procedure...")
+        print("")
+
+        x_prev = self.sensor.x_offset
+        y_prev = self.sensor.y_offset
+        z_prev = self.sensor.z_offset + self.sensor.z_clearance
+
+        # Loop through every calibration point
+        for i in tqdm(range(N), desc="Sensor Calibration Progress"):
+            # If specified penetration depth exceeds maximum value, print message
+            if abs(self.calibration_points[i][2]) > self.sensor.max_penetration:
+                print("Line " + str(i+1) + ": Maximum penetration depth for sensor exceeded. Skipping calibration point.")
+            # If penetration depth does not exceed maximum value, move to calibration point
+            else:
+                # Get absolute XYZ coordinates
+                x = self.sensor.x_offset + self.calibration_points[i][0]
+                y = self.sensor.y_offset + self.calibration_points[i][1]
+                z = self.sensor.z_offset - abs(self.calibration_points[i][2])
+
+                # Move to Z clearance height
+                self.printer.send_gcode("G0 Z" + str(self.sensor.z_offset + self.sensor.z_clearance))
+
+                # Move to desired XY location
+                self.printer.send_gcode("G0 X" + str(x) + " Y" + str(y))
+
+                # Move to desired Z penetration
+                self.printer.send_gcode("G0 Z" + str(z))
+                
+                # Calculate time required to reach position
+                # Assumes x and y speed of 10 mm/s, z speed of 4 mm/s
+                travel_time = abs(self.sensor.z_offset + self.sensor.z_clearance - z_prev) / 4 + max(abs(x - x_prev), abs(y - y_prev)) / 10 + abs(z - (self.sensor.z_offset + self.sensor.z_clearance)) / 4
+                time.sleep(travel_time + 1)
+
+                # Update variables
+                x_prev = x
+                y_prev = y
+                z_prev = z
+
+                # Take desired number of images
+                if save_images == True:
+                    with open(os.path.join(data_save_path, "annotations", "probe_data.csv"), 'a', newline='') as csv_file:
+                        csv_writer = csv.writer(csv_file)
+
+                        # 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 = Image.fromarray(frame)
+                            img.save(img_path)
+
+                            csv_writer.writerow([img_name, self.calibration_points[i][0], self.calibration_points[i][1], self.calibration_points[i][2]])
+
+                            img_idx += 1
+
+                            time.sleep(0.5)
+
+                else:
+                    with open(os.path.join(data_save_path, "annotations", "probe_data.csv"), 'a', newline='') as csv_file:
+                        csv_writer = csv.writer(csv_file)
+                        csv_writer.writerow(["---", self.calibration_points[i][0], self.calibration_points[i][1], self.calibration_points[i][2]])
+
+        # Move to Z clearance height
+        self.printer.send_gcode("G0 Z" + str(self.sensor.z_offset + self.sensor.z_clearance))
+
+        print("")
+
+        # Update printer display
+        self.printer.send_gcode("M117 Calibration Done!")
+
+        # Disconnect from 3D printer
+        self.disconnect_printer()
+
+        # Disconnect from sensor
+        if save_images == True:
+            self.disconnect_sensor()
+
+        print("Calibration procedure complete!")
+        print("")
+
+        return True

py3DCal/data_collection/printers/Ender3/Ender3.py

@@ -0,0 +1,82 @@
+from ..Printer import Printer
+import serial
+from typing import Union
+from pathlib import Path
+
+class Ender3(Printer):
+    """
+    Ender3: A Printer Class for the Ender 3
+    Args:
+        port (str or pathlib.Path): The COM port the printer is connected to.
+    """
+    def __init__(self, port: Union[str, Path]):
+        self.port = port
+        self.name = "Ender 3"
+
+    def connect(self):
+        """
+        Connects to the Ender 3 printer.
+        """
+        # Code to connect to the printer
+        self.ser = serial.Serial(self.port, 115200)
+
+    def disconnect(self):
+        """
+        Disconnects from the Ender 3 printer.
+        """
+        # Code to disconnect from the printer
+        self.ser.close()
+
+    def send_gcode(self, command: str):
+        """
+        Sends a G-code command to the Ender 3 printer.
+        Args:
+            command (str): The G-code command to be sent to the 3D printer.
+        """
+        # Code to execute gcode command on the printer
+        self.ser.write(str.encode(command + "\r\n"))
+
+    def get_response(self):
+        """
+        Gets messages sent by the Ender 3 printer.
+
+        Returns:
+            response (str): The message sent by the printer.
+        """
+        # Code to return message from the printer
+        response = self.ser.readline().decode('utf-8')
+
+        return response
+
+    def initialize(self, xy_only: bool = False):
+        """
+        Initializes the Ender 3 printer (homes the printer, sets units, adjusts fans, etc).
+        Args:
+            xy_only (bool): If True, only homes the X and Y axes.
+        """
+        # Code to initialize printer (home, set units, set absolute/relative movements, adjust fan speeds, etc.)
+
+        # Use Metric Values
+        self.send_gcode("G21")
+
+        # Absolute Positioning
+        self.send_gcode("G90")
+
+        # Fan Off
+        self.send_gcode("M107")
+
+        if xy_only:
+            # Home Printer X Y
+            self.send_gcode("G28 X Y")
+        else:
+            # Home Printer X Y Z
+            self.send_gcode("G28")
+
+        # Check if homing is complete
+        ok_count = 0
+
+        while ok_count < 4:
+            if "ok" in self.get_response():
+                ok_count += 1
+
+        return True

py3DCal/data_collection/printers/Printer.py

@@ -0,0 +1,63 @@
+from abc import ABC, abstractmethod
+
+class Printer(ABC):
+    """
+    Printer: An abstract base class for 3D printers.
+    """
+    def __init__(self):
+        self.name = ""
+        
+    @abstractmethod
+    def connect(self):
+        """ Connects to the 3D printer.
+        """
+        pass
+
+    @abstractmethod
+    def disconnect(self):
+        """ Disconnects from the 3D printer.
+        """
+        pass
+
+    @abstractmethod
+    def send_gcode(self, command: str):
+        """ Sends a G-code command to the 3D printer.
+
+        Args:
+            command (str): The G-code command to be sent to the 3D printer.
+        """
+        pass
+
+    @abstractmethod
+    def get_response(self):
+        """ Gets messages sent by the 3D printer.
+
+        Returns:
+            str: The message sent by the printer.
+        """
+        pass
+
+    @abstractmethod
+    def initialize(self):
+        """ Initializes the 3D printer (homes the printer, sets units, adjusts fans, etc).
+        """
+        pass
+
+    def go_to(self, x: float = None, y: float = None, z: float = None):
+        """
+        Moves the printer to a specific position.
+        Args:
+            x (float, optional): The X coordinate to move to.
+            y (float, optional): The Y coordinate to move to.
+            z (float, optional): The Z coordinate to move to.
+        """
+        # Code to move the printer to a specific position
+        command = "G0"
+        if x is not None:
+            command += f" X{x}"
+        if y is not None:
+            command += f" Y{y}"
+        if z is not None:
+            command += f" Z{z}"
+
+        self.send_gcode(command)

py3DCal/data_collection/sensors/DIGIT/DIGIT.py

@@ -0,0 +1,47 @@
+from ..Sensor import Sensor
+import cv2
+import os
+
+try:
+    from digit_interface import Digit
+except:
+    pass
+
+class DIGIT(Sensor):
+    """
+    DIGIT: A Sensor Class for the DIGIT sensor
+    Args:
+        serial_number (str): The serial number of the DIGIT sensor.
+    """
+    def __init__(self, serial_number: str):
+        self.serial_number = serial_number
+        self.name = "DIGIT"
+        self.x_offset = 110
+        self.y_offset = 111.5
+        self.z_offset = 137
+        self.z_clearance = 2
+        self.max_penetration = 4
+        self.default_calibration_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "default.csv")
+
+    def connect(self):
+        """
+        Connects to the DIGIT sensor.
+        """
+        # Code to connect to the sensor
+        self.sensor = Digit(self.serial_number)
+        self.sensor.connect()
+        self.sensor.set_fps(30)
+        
+    def disconnect(self):
+        """
+        Disconnects from the DIGIT sensor.
+        """
+        # Code to disconnect from the sensor
+        self.sensor.disconnect()
+
+    def capture_image(self):
+        """
+        Captures an image from the DIGIT sensor.
+        """
+        # Code to return an image from the sensor
+        return cv2.flip(self.sensor.get_frame(), 1)