Repository: Jordan-Pierce/CoralNet-Toolbox (language: python)

Issue:
Clean up



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 9980d455..0f672b99 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,18 +1,22 @@
 *.hd
 *.hdf5
-*.idea
 *.pt
 *.torchscript
 *.bin
 *.engine
 *.onnx
 *__pycache__/
-*gooey_config.json
-.amlignore
 *.amltmp
-.ipynb_aml_checkpoints
+*.amlignore
+*.ipynb_aml_checkpoints
+
+*egg-info
+*.idea
+*.vscode
+launch.json
+
 build
 dist
-*egg-info
+
 Data/
 toolbox/QtClustering.py
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index effdc644..ee3034c1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,12 +1,13 @@
 PyQt5>=5.15.11
 ultralytics==8.3.0
+supervision==0.24.0
 pycocotools>=2.0.6
 timm==0.9.2
-git+https://github.com/ChaoningZhang/MobileSAM.git
-git+https://github.com/facebookresearch/segment-anything.git
-git+https://github.com/facebookresearch/sam2.git
 autodistill==0.1.28
 autodistill-grounding-dino==0.1.4
 roboflow==1.1.48
+git+https://github.com/ChaoningZhang/MobileSAM.git
+git+https://github.com/facebookresearch/segment-anything.git
+git+https://github.com/facebookresearch/sam2.git
 rasterio>=1.4.1
 dill
\ No newline at end of file
diff --git a/toolbox/MachineLearning/ConfusionMatrix.py b/toolbox/MachineLearning/ConfusionMatrix.py
index a877505d..26a3f4ea 100644
--- a/toolbox/MachineLearning/ConfusionMatrix.py
+++ b/toolbox/MachineLearning/ConfusionMatrix.py
@@ -333,4 +333,4 @@ def save_results(self, directory):
 
         # Save confusion matrices (both normalized and real versions)
         self.save_normalized_confusion_matrix_png(directory)
-        self.save_real_confusion_matrix_png(directory)
\ No newline at end of file
+        self.save_real_confusion_matrix_png(directory)
diff --git a/toolbox/MachineLearning/QtBatchInference.py b/toolbox/MachineLearning/QtBatchInference.py
index 38925de4..6692994a 100644
--- a/toolbox/MachineLearning/QtBatchInference.py
+++ b/toolbox/MachineLearning/QtBatchInference.py
@@ -19,6 +19,13 @@
 
 
 class BatchInferenceDialog(QDialog):
+    """
+    Dialog for performing batch inference on images for image classification, object detection, 
+    and instance segmentation.
+    
+    :param main_window: MainWindow object
+    :param parent: Parent widget
+    """
     def __init__(self, main_window, parent=None):
         super().__init__(parent)
         self.main_window = main_window
@@ -78,16 +85,26 @@ def __init__(self, main_window, parent=None):
         self.main_window.uncertaintyChanged.connect(self.on_uncertainty_changed)
 
     def update_uncertainty_label(self):
+        """
+        Update the uncertainty threshold label based on the slider value.
+        """
         # Convert the slider value to a ratio (0-1)
         value = self.uncertainty_threshold_slider.value() / 100.0
         self.main_window.update_uncertainty_thresh(value)
 
     def on_uncertainty_changed(self, value):
-        # Update the slider and label when the shared data changes
+        """
+        Update the slider and label when the shared data changes.
+        
+        :param value: New uncertainty threshold value
+        """
         self.uncertainty_threshold_slider.setValue(int(value * 100))
         self.uncertainty_threshold_label.setText(f"{value:.2f}")
 
     def setup_classification_tab(self):
+        """
+        Set up the layout and widgets for the image classification tab.
+        """
         layout = QVBoxLayout()
 
         # Create a group box for annotation options
@@ -151,6 +168,9 @@ def setup_classification_tab(self):
         self.classification_tab.setLayout(layout)
 
     def setup_detection_tab(self):
+        """
+        Set up the layout and widgets for the object detection tab.
+        """
         layout = QVBoxLayout()
 
         # Create a group box for image options
@@ -188,6 +208,9 @@ def setup_detection_tab(self):
         self.detection_tab.setLayout(layout)
 
     def setup_segmentation_tab(self):
+        """
+        Set up the layout and widgets for the instance segmentation tab.
+        """
         layout = QVBoxLayout()
 
         # Create a group box for image options
@@ -225,6 +248,9 @@ def setup_segmentation_tab(self):
         self.segmentation_tab.setLayout(layout)
 
     def on_ok_clicked(self):
+        """
+        Handle the "OK" button click event.
+        """
         if self.classification_all_checkbox.isChecked():
             reply = QMessageBox.warning(self,
                                         "Warning",
@@ -237,6 +263,9 @@ def on_ok_clicked(self):
         self.accept()  # Close the dialog after applying the changes
 
     def apply(self):
+        """
+        Apply the selected batch inference options.
+        """
         # Pause the cursor
         QApplication.setOverrideCursor(Qt.WaitCursor)
 
@@ -260,6 +289,9 @@ def apply(self):
         QApplication.restoreOverrideCursor()
 
     def apply_classification(self):
+        """
+        Apply batch inference for image classification.
+        """
         # Get the Review Annotations
         if self.classification_review_checkbox.isChecked():
             for image_path in self.get_selected_image_paths():
@@ -274,14 +306,25 @@ def apply_classification(self):
         self.batch_inference('classify')
 
     def apply_detection(self):
+        """
+        Apply batch inference for object detection.
+        """
         self.image_paths = self.get_selected_image_paths()
         self.batch_inference('detect')
 
     def apply_segmentation(self):
+        """
+        Apply batch inference for instance segmentation.
+        """
         self.image_paths = self.get_selected_image_paths()
         self.batch_inference('segment')
 
     def get_selected_image_paths(self):
+        """
+        Get the selected image paths based on the current tab and options.
+        
+        :return: List of selected image paths
+        """
         if self.tab_widget.currentIndex() == 0:  # Classification
             if self.apply_filtered_checkbox.isChecked():
                 return self.image_window.filtered_image_paths
@@ -317,6 +360,9 @@ def get_selected_image_paths(self):
                 return self.image_window.image_paths
 
     def preprocess_patch_annotations(self):
+        """
+        Preprocess patch annotations by cropping the images based on the annotations.
+        """
         # Get unique image paths
         self.image_paths = list(set(a.image_path for a in self.annotations))
         if not self.image_paths:
@@ -352,6 +398,11 @@ def crop(image_path, image_annotations):
         progress_bar.close()
 
     def batch_inference(self, task):
+        """
+        Perform batch inference on the selected images and annotations.
+        
+        :param task: Task type ('classify', 'detect', or 'segment')
+        """
         # Make predictions on each image's annotations
         progress_bar = ProgressBar(self, title=f"Batch Inference")
         progress_bar.show()
@@ -385,4 +436,4 @@ def batch_inference(self, task):
             self.deploy_model_dialog.predict_segmentation(image_paths=self.image_paths)
 
         progress_bar.stop_progress()
-        progress_bar.close()
\ No newline at end of file
+        progress_bar.close()
diff --git a/toolbox/MachineLearning/QtDeployModel.py b/toolbox/MachineLearning/QtDeployModel.py
index 578f1492..3aa6f278 100644
--- a/toolbox/MachineLearning/QtDeployModel.py
+++ b/toolbox/MachineLearning/QtDeployModel.py
@@ -30,124 +30,124 @@
 
 
 class DeployModelDialog(QDialog):
+    """
+    Dialog for deploying machine learning models for image classification, object detection, 
+    and instance segmentation.
+    
+    :param main_window: MainWindow object
+    :param parent: Parent widget
+    """
     def __init__(self, main_window, parent=None):
         super().__init__(parent)
         self.main_window = main_window
         self.label_window = main_window.label_window
         self.annotation_window = main_window.annotation_window
-
         self.sam_dialog = None
 
         self.setWindowTitle("Deploy Model")
         self.resize(400, 300)
 
         self.layout = QVBoxLayout(self)
-
         self.model_paths = {'classify': None, 'detect': None, 'segment': None}
         self.loaded_models = {'classify': None, 'detect': None, 'segment': None}
         self.class_mappings = {'classify': None, 'detect': None, 'segment': None}
         self.use_sam = {'classify': None, 'detect': None, 'segment': None}
 
-        self.tab_widget = QTabWidget()
-        self.layout.addWidget(self.tab_widget)
-
-        self.classification_tab = QWidget()
-        self.detection_tab = QWidget()
-        self.segmentation_tab = QWidget()
-
-        self.tab_widget.addTab(self.classification_tab, "Image Classification")
-        self.tab_widget.addTab(self.detection_tab, "Object Detection")
-        self.tab_widget.addTab(self.segmentation_tab, "Instance Segmentation")
-
-        self.setup_classification_tab()
-        self.setup_detection_tab()
-        self.setup_segmentation_tab()
-
-        # Task-specific status bars
-        self.status_bars = {
-            'classify': QLabel("No model loaded"),
-            'detect': QLabel("No model loaded"),
-            'segment': QLabel("No model loaded")
-        }
-        self.layout.addWidget(self.status_bars['classify'])
-        self.layout.addWidget(self.status_bars['detect'])
-        self.layout.addWidget(self.status_bars['segment'])
+        self.setup_tabs()
+        self.setup_status_bars()
 
         self.tab_widget.currentChanged.connect(self.update_status_bar_visibility)
-
         self.setLayout(self.layout)
 
     def showEvent(self, event: QShowEvent):
+        """
+        Handle the show event to check and display class names and update status bar visibility.
+        
+        :param event: QShowEvent object
+        """
         super().showEvent(event)
         self.check_and_display_class_names()
-        self.update_status_bar_visibility(self.tab_widget.currentIndex())
+        self.update_status_bar_visibility()
+
+    def setup_tabs(self):
+        """
+        Set up the tabs for different tasks (classification, detection, segmentation).
+        """
+        self.tab_widget = QTabWidget()
+        self.layout.addWidget(self.tab_widget)
+
+        tasks = [("Image Classification", "classify"),
+                 ("Object Detection", "detect"),
+                 ("Instance Segmentation", "segment")]
+
+        for label, task in tasks:
+            tab = QWidget()
+            text_area = self.setup_tab(tab, task)
+            self.tab_widget.addTab(tab, label)
+            setattr(self, f"{task}_text_area", text_area)
 
     def setup_tab(self, tab, task):
+        """
+        Set up a single tab with text area and buttons for the given task.
+        
+        :param tab: QWidget object
+        :param task: Task identifier as a string
+        :return: QTextEdit widget
+        """
         layout = QVBoxLayout()
-
         text_area = QTextEdit()
         text_area.setReadOnly(True)
         layout.addWidget(text_area)
 
-        browse_model_button = QPushButton("Browse Model")
-        browse_model_button.clicked.connect(lambda: self.browse_file(task))
-        layout.addWidget(browse_model_button)
-
-        browse_class_mapping_button = QPushButton("Browse Class Mapping")
-        browse_class_mapping_button.clicked.connect(lambda: self.browse_class_mapping_file(task))
-        layout.addWidget(browse_class_mapping_button)
+        buttons = [
+            ("Browse Model", lambda: self.browse_file(task)),
+            ("Browse Class Mapping", lambda: self.browse_class_mapping_file(task)),
+            ("Load Model", lambda: self.load_model(task)),
+            ("Deactivate Model", lambda: self.deactivate_model(task))
+        ]
 
-        load_button = QPushButton("Load Model")
-        load_button.clicked.connect(lambda: self.load_model(task))
-        layout.addWidget(load_button)
-
-        deactivate_button = QPushButton("Deactivate Model")
-        deactivate_button.clicked.connect(lambda: self.deactivate_model(task))
-        layout.addWidget(deactivate_button)
+        for btn_text, callback in buttons:
+            button = QPushButton(btn_text)
+            button.clicked.connect(callback)
+            layout.addWidget(button)
 
         use_sam_checkbox = QCheckBox("Use SAM for creating Polygons")
-        use_sam_checkbox.stateChanged.connect(lambda: self.is_sam_model_deployed())
-
-        if task == 'classify':
-            use_sam_checkbox.setChecked(False)
-            use_sam_checkbox.setEnabled(False)
-
+        use_sam_checkbox.stateChanged.connect(self.is_sam_model_deployed)
+        use_sam_checkbox.setEnabled(task != 'classify')
         layout.addWidget(use_sam_checkbox)
-
-        # Store the checkbox in the dictionary
         self.use_sam[task] = use_sam_checkbox
 
         tab.setLayout(layout)
-
         return text_area
 
-    def setup_classification_tab(self):
-        self.classification_text_area = self.setup_tab(self.classification_tab, 'classify')
-
-    def setup_detection_tab(self):
-        self.detection_text_area = self.setup_tab(self.detection_tab, 'detect')
-
-    def setup_segmentation_tab(self):
-        self.segmentation_text_area = self.setup_tab(self.segmentation_tab, 'segment')
-
-    def get_current_task(self):
-        index = self.tab_widget.currentIndex()
-        return ['classify', 'detect', 'segment'][index]
-
-    def get_text_area(self, task):
-        if task == 'classify':
-            return self.classification_text_area
-        elif task == 'detect':
-            return self.detection_text_area
-        elif task == 'segment':
-            return self.segmentation_text_area
-
-    def update_status_bar_visibility(self, index):
+    def setup_status_bars(self):
+        """
+        Set up status bars for each task.
+        """
+        self.status_bars = {
+            'classify': QLabel("No model loaded"),
+            'detect': QLabel("No model loaded"),
+            'segment': QLabel("No model loaded")
+        }
+        for bar in self.status_bars.values():
+            self.layout.addWidget(bar)
+
+    def update_status_bar_visibility(self):
+        """
+        Update the visibility of status bars based on the selected tab.
+        
+        :param index: Index of the selected tab
+        """
         current_task = self.get_current_task()
         for task, status_bar in self.status_bars.items():
             status_bar.setVisible(task == current_task)
 
     def browse_file(self, task):
+        """
+        Browse and select a model file for the given task.
+        
+        :param task: Task identifier as a string
+        """
         options = QFileDialog.Options()
         file_path, _ = QFileDialog.getOpenFileName(self,
                                                    "Open Model File", "",
@@ -168,6 +168,11 @@ def browse_file(self, task):
                 self.load_class_mapping(task, class_mapping_path)
 
     def browse_class_mapping_file(self, task):
+        """
+        Browse and select a class mapping file for the given task.
+        
+        :param task: Task identifier as a string
+        """
         options = QFileDialog.Options()
         file_path, _ = QFileDialog.getOpenFileName(self,
                                                    "Open Class Mapping File", "",
@@ -176,7 +181,39 @@ def browse_class_mapping_file(self, task):
         if file_path:
             self.load_class_mapping(task, file_path)
 
+    def get_text_area(self, task):
+        """
+        Retrieves the text area widget associated with the given task.
+        
+        :param task: Task identifier as a string
+        :return: QTextEdit widget corresponding to the task
+        """
+        if task == "classify":
+            return self.classify_text_area
+        elif task == "detect":
+            return self.detect_text_area
+        elif task == "segment":
+            return self.segment_text_area
+        else:
+            raise ValueError(f"Unknown task: {task}")
+
+    def get_current_task(self):
+        """
+        Retrieves the current task based on the selected tab in the deployment dialog.
+        
+        :return: Selected task identifier as a string
+        """
+        current_index = self.tab_widget.currentIndex()
+        tasks = ["classify", "detect", "segment"]
+        return tasks[current_index]
+
     def load_class_mapping(self, task, file_path):
+        """
+        Load the class mapping file for the given task.
+        
+        :param task: Task identifier as a string
+        :param file_path: Path to the class mapping file
+        """
         try:
             with open(file_path, 'r') as f:
                 self.class_mappings[task] = json.load(f)
@@ -185,6 +222,11 @@ def load_class_mapping(self, task, file_path):
             QMessageBox.critical(self, "Error", f"Failed to load class mapping file: {str(e)}")
 
     def is_sam_model_deployed(self):
+        """
+        Check if the SAM model is deployed and update the checkbox state accordingly.
+        
+        :return: Boolean indicating whether the SAM model is deployed
+        """
         self.sam_dialog = self.main_window.sam_deploy_model_dialog
 
         if not self.sam_dialog.loaded_model:
@@ -196,6 +238,11 @@ def is_sam_model_deployed(self):
         return True
 
     def load_model(self, task):
+        """
+        Load the model for the given task.
+        
+        :param task: Task identifier as a string
+        """
         if not self.model_paths[task]:
             QMessageBox.warning(self, "Warning", f"No {task} model file selected")
             return
@@ -218,6 +265,11 @@ def load_model(self, task):
             QApplication.restoreOverrideCursor()
 
     def handle_missing_class_mapping(self, task):
+        """
+        Handle the case when the class mapping file is missing.
+        
+        :param task: Task identifier as a string
+        """
         reply = QMessageBox.question(self,
                                      'No Class Mapping Found',
                                      'Do you want to create generic labels automatically?',
@@ -229,6 +281,11 @@ def handle_missing_class_mapping(self, task):
             QMessageBox.information(self, "Model Loaded", f"{task.capitalize()} model loaded successfully.")
 
     def add_labels_to_label_window(self, task):
+        """
+        Add labels to the label window based on the class mapping.
+        
+        :param task: Task identifier as a string
+        """
         if self.class_mappings[task]:
             for label in self.class_mappings[task].values():
                 self.label_window.add_label_if_not_exists(label['short_label_code'],
@@ -236,6 +293,11 @@ def add_labels_to_label_window(self, task):
                                                           QColor(*label['color']))
 
     def create_generic_labels(self, task):
+        """
+        Create generic labels for the given task.
+        
+        :param task: Task identifier as a string
+        """
         class_names = list(self.loaded_models[task].names.values())
         for class_name in class_names:
             r = random.randint(0, 255)
@@ -256,6 +318,11 @@ def create_generic_labels(self, task):
         self.class_mappings[task] = class_mapping
 
     def check_and_display_class_names(self, task=None):
+        """
+        Check and display the class names for the given task.
+        
+        :param task: Task identifier as a string (optional)
+        """
         if task is None:
             task = self.get_current_task()
 
@@ -285,6 +352,11 @@ def check_and_display_class_names(self, task=None):
                                     f"\n{missing_labels_str}")
 
     def deactivate_model(self, task):
+        """
+        Deactivate the model for the given task.
+        
+        :param task: Task identifier as a string
+        """
         self.loaded_models[task] = None
         self.model_paths[task] = None
         self.class_mappings[task] = None
@@ -292,8 +364,22 @@ def deactivate_model(self, task):
         empty_cache()
         self.status_bars[task].setText(f"No {task} model loaded")
         self.get_text_area(task).setText(f"No {task} model file selected")
+        
+    def get_confidence_threshold(self):
+        """
+        Get the confidence threshold for predictions.
+        
+        :return: Confidence threshold as a float
+        """
+        threshold = self.main_window.get_uncertainty_thresh()
+        return threshold if threshold < 0.10 else 0.10
 
     def predict_classification(self, annotations=None):
+        """
+        Predict the classification results for the given annotations.
+        
+        :param annotations: List of annotations (optional)
+        """
         if self.loaded_models['classify'] is None:
             return
 
@@ -318,6 +404,11 @@ def predict_classification(self, annotations=None):
         empty_cache()
 
     def preprocess_classification_annotations(self, annotations):
+        """
+        Preprocess the given annotations for classification.
+        
+        :param annotations: List of annotations
+        """
         if not annotations:
             return
 
@@ -325,18 +416,22 @@ def preprocess_classification_annotations(self, annotations):
         for annotation in annotations:
             images_np.append(pixmap_to_numpy(annotation.cropped_image))
 
-        progress_bar = ProgressBar(self, title=f"Making Classification Predictions")
+        progress_bar = ProgressBar(self, title="Making Classification Predictions")
         progress_bar.show()
         progress_bar.start_progress(len(annotations))
 
+        # Predict the classification results
         results = self.loaded_models['classify'](images_np, stream=True, device=self.main_window.device)
 
+        # Process the classification results
         for annotation, result in zip(annotations, results):
             self.process_classification_result(annotation, result)
             progress_bar.update_progress()
 
+        # Update the image window and confidence window
         self.main_window.confidence_window.display_cropped_image(annotation)
 
+        # Update the annotations for each image in the image window
         image_paths = list(set([annotation.image_path for annotation in annotations]))
         for image_path in image_paths:
             self.main_window.image_window.update_image_annotations(image_path)
@@ -345,6 +440,12 @@ def preprocess_classification_annotations(self, annotations):
         progress_bar.close()
 
     def process_classification_result(self, annotation, results):
+        """
+        Process the classification result for a single annotation.
+        
+        :param annotation: Annotation object
+        :param results: Classification results
+        """
         predictions = {}
 
         try:
@@ -369,6 +470,11 @@ def process_classification_result(self, annotation, results):
                 annotation.update_label(label)
 
     def predict_detection(self, image_paths=None):
+        """
+        Predict the detection results for the given image paths.
+        
+        :param image_paths: List of image paths (optional)
+        """
         if self.loaded_models['detect'] is None:
             return
 
@@ -377,11 +483,7 @@ def predict_detection(self, image_paths=None):
         if not image_paths:
             image_paths = [self.annotation_window.current_image_path]
 
-        # Perform detection
-        if self.main_window.get_uncertainty_thresh() < 0.10:
-            conf = self.main_window.get_uncertainty_thresh()
-        else:
-            conf = 0.10  # Arbitrary value to prevent too many detections
+        conf = self.get_confidence_threshold()
 
         results = self.loaded_models['detect'](image_paths,
                                                agnostic_nms=True,
@@ -390,96 +492,133 @@ def predict_detection(self, image_paths=None):
                                                device=self.main_window.device,
                                                stream=True)
 
-        # Check if the user selected to use SAM
         if self.use_sam['detect'].isChecked() and self.sam_dialog.loaded_model:
-            # Convert the boxes to SAM masks, process as segmentations
             results = self.sam_dialog.boxes_to_masks(results)
             self.process_segmentation_results(results)
         else:
-            # Process as detections
             self.process_detection_results(results)
 
         QApplication.restoreOverrideCursor()
         gc.collect()
         empty_cache()
 
+
     def process_detection_results(self, results_generator):
-        # Get the class mapping for detection
+        """
+        Process the detection results from the results generator.
+        
+        :param results_generator: Generator yielding detection results
+        """
         class_mapping = self.class_mappings['detect']
-
-        progress_bar = ProgressBar(self, title=f"Making Detection Predictions")
+        progress_bar = ProgressBar(self, title="Making Detection Predictions")
         progress_bar.show()
 
         for results in results_generator:
             progress_bar.start_progress(len(results))
             for result in results:
-                try:
-                    # Get the image path
-                    image_path = result.path.replace("\\", "/")
-
-                    # Extract the results
-                    cls = int(result.boxes.cls.cpu().numpy()[0])
-                    cls_name = result.names[cls]
-                    conf = float(result.boxes.conf.cpu().numpy()[0])
-                    x_min, y_min, x_max, y_max = map(float, result.boxes.xyxy.cpu().numpy()[0])
-
-                    # Determine the short label
-                    short_label = 'Review'
-                    if conf > self.main_window.get_uncertainty_thresh():
-                        short_label = class_mapping.get(cls_name, {}).get('short_label_code', 'Review')
-
-                    # Prepare the annotation data
-                    label = self.label_window.get_label_by_short_code(short_label)
-                    top_left = QPointF(x_min, y_min)
-                    bottom_right = QPointF(x_max, y_max)
-
-                    # Create the rectangle annotation
-                    annotation = RectangleAnnotation(top_left,
-                                                     bottom_right,
-                                                     label.short_label_code,
-                                                     label.long_label_code,
-                                                     label.color,
-                                                     image_path,
-                                                     label.id,
-                                                     self.main_window.get_transparency_value(),
-                                                     show_msg=True)
-
-                    # Store the annotation and display the cropped image
-                    self.annotation_window.annotations_dict[annotation.id] = annotation
-
-                    # Connect update signals
-                    annotation.selected.connect(self.annotation_window.select_annotation)
-                    annotation.annotationDeleted.connect(self.annotation_window.delete_annotation)
-                    annotation.annotationUpdated.connect(self.main_window.confidence_window.display_cropped_image)
-
-                    # Add the prediction for the confidence window
-                    predictions = {self.label_window.get_label_by_short_code(cls_name): conf}
-                    annotation.update_machine_confidence(predictions)
-
-                    # Update label if confidence is below threshold
-                    if conf < self.main_window.get_uncertainty_thresh():
-                        review_label = self.label_window.get_label_by_id('-1')
-                        annotation.update_label(review_label)
-
-                    # Create the graphics and cropped image
-                    if image_path == self.annotation_window.current_image_path:
-                        annotation.create_graphics_item(self.annotation_window.scene)
-                        annotation.create_cropped_image(self.annotation_window.rasterio_image)
-                        self.main_window.confidence_window.display_cropped_image(annotation)
-
-                    # Update the image annotations
-                    self.main_window.image_window.update_image_annotations(image_path)
-
-                    # Update the progress bar
-                    progress_bar.update_progress()
-
-                except Exception as e:
-                    print(f"Warning: Failed to process detection result\n{e}")
+                self.process_single_detection_result(result, class_mapping)
+                progress_bar.update_progress()
 
         progress_bar.stop_progress()
         progress_bar.close()
 
+    def process_single_detection_result(self, result, class_mapping):
+        """
+        Process a single detection result.
+        
+        :param result: Detection result
+        :param class_mapping: Class mapping dictionary
+        """
+        try:
+            image_path = result.path.replace("\\", "/")
+            cls, cls_name, conf, x_min, y_min, x_max, y_max = self.extract_detection_result(result)
+            short_label = self.get_short_label_for_detection(cls_name, conf, class_mapping)
+            label = self.label_window.get_label_by_short_code(short_label)
+            annotation = self.create_rectangle_annotation(x_min, y_min, x_max, y_max, label, image_path)
+            self.store_and_display_annotation(annotation, image_path, cls_name, conf)
+        except Exception as e:
+            print(f"Warning: Failed to process detection result\n{e}")
+
+    def extract_detection_result(self, result):
+        """
+        Extract relevant information from a detection result.
+        
+        :param result: Detection result
+        :return: Tuple containing class, class name, confidence, and bounding box coordinates
+        """
+        cls = int(result.boxes.cls.cpu().numpy()[0])
+        cls_name = result.names[cls]
+        conf = float(result.boxes.conf.cpu().numpy()[0])
+        x_min, y_min, x_max, y_max = map(float, result.boxes.xyxy.cpu().numpy()[0])
+        return cls, cls_name, conf, x_min, y_min, x_max, y_max
+
+    def get_short_label_for_detection(self, cls_name, conf, class_mapping):
+        """
+        Get the short label for a detection result based on confidence and class mapping.
+        
+        :param cls_name: Class name
+        :param conf: Confidence score
+        :param class_mapping: Class mapping dictionary
+        :return: Short label as a string
+        """
+        if conf <= self.main_window.get_uncertainty_thresh():
+            return 'Review'
+        return class_mapping.get(cls_name, {}).get('short_label_code', 'Review')
+
+    def create_rectangle_annotation(self, x_min, y_min, x_max, y_max, label, image_path):
+        """
+        Create a rectangle annotation for the given bounding box coordinates and label.
+        
+        :param x_min: Minimum x-coordinate
+        :param y_min: Minimum y-coordinate
+        :param x_max: Maximum x-coordinate
+        :param y_max: Maximum y-coordinate
+        :param label: Label object
+        :param image_path: Path to the image
+        :return: RectangleAnnotation object
+        """
+        top_left = QPointF(x_min, y_min)
+        bottom_right = QPointF(x_max, y_max)
+        return RectangleAnnotation(top_left, 
+                                   bottom_right, 
+                                   label.short_label_code, 
+                                   label.long_label_code, 
+                                   label.color, 
+                                   image_path, 
+                                   label.id, 
+                                   self.main_window.get_transparency_value(), 
+                                   show_msg=True)
+
+    def store_and_display_annotation(self, annotation, image_path, cls_name, conf):
+        """
+        Store and display the annotation in the annotation window and image window.
+        
+        :param annotation: Annotation object
+        :param image_path: Path to the image
+        :param cls_name: Class name
+        :param conf: Confidence score
+        """
+        self.annotation_window.annotations_dict[annotation.id] = annotation
+        annotation.selected.connect(self.annotation_window.select_annotation)
+        annotation.annotationDeleted.connect(self.annotation_window.delete_annotation)
+        annotation.annotationUpdated.connect(self.main_window.confidence_window.display_cropped_image)
+        predictions = {self.label_window.get_label_by_short_code(cls_name): conf}
+        annotation.update_machine_confidence(predictions)
+        if conf < self.main_window.get_uncertainty_thresh():
+            review_label = self.label_window.get_label_by_id('-1')
+            annotation.update_label(review_label)
+        if image_path == self.annotation_window.current_image_path:
+            annotation.create_graphics_item(self.annotation_window.scene)
+            annotation.create_cropped_image(self.annotation_window.rasterio_image)
+            self.main_window.confidence_window.display_cropped_image(annotation)
+        self.main_window.image_window.update_image_annotations(image_path)
+
     def predict_segmentation(self, image_paths=None):
+        """
+        Predict the segmentation results for the given image paths.
+        
+        :param image_paths: List of image paths (optional)
+        """
         if self.loaded_models['segment'] is None:
             return
 
@@ -488,11 +627,7 @@ def predict_segmentation(self, image_paths=None):
         if not image_paths:
             image_paths = [self.annotation_window.current_image_path]
 
-        # Perform detection
-        if self.main_window.get_uncertainty_thresh() < 0.10:
-            conf = self.main_window.get_uncertainty_thresh()
-        else:
-            conf = 0.10  # Arbitrary value to prevent too many detections
+        conf = self.get_confidence_threshold()
 
         results = self.loaded_models['segment'](image_paths,
                                                 agnostic_nms=True,
@@ -501,12 +636,9 @@ def predict_segmentation(self, image_paths=None):
                                                 device=self.main_window.device,
                                                 stream=True)
 
-        # Check if the user selected to use SAM
         if self.use_sam['segment'].isChecked() and self.sam_dialog.loaded_model:
-            # Convert the boxes to SAM masks
             results = self.sam_dialog.boxes_to_masks(results)
 
-        # Process the segmentation results
         self.process_segmentation_results(results)
 
         QApplication.restoreOverrideCursor()
@@ -514,9 +646,14 @@ def predict_segmentation(self, image_paths=None):
         empty_cache()
 
     def process_segmentation_results(self, results_generator):
-        # If SAM is being used, and there is no class mapping for segmentation, use the detection class mapping
+        """
+        Process the segmentation results from the results generator.
+        
+        :param results_generator: Generator yielding segmentation results
+        """
         class_mapping = self.class_mappings['segment']
-        class_mapping = class_mapping if class_mapping else self.class_mappings['detect']
+        if not class_mapping:
+            class_mapping = self.class_mappings['detect']
 
         progress_bar = ProgressBar(self, title=f"Making Segmentation Predictions")
         progress_bar.show()
@@ -524,66 +661,57 @@ def process_segmentation_results(self, results_generator):
         for results in results_generator:
             progress_bar.start_progress(len(results))
             for result in results:
-                try:
-                    # Get the image path
-                    image_path = result.path.replace("\\", "/")
-
-                    # Extract the results
-                    cls = int(result.boxes.cls.cpu().numpy()[0])
-                    cls_name = result.names[cls]
-                    conf = float(result.boxes.conf.cpu().numpy()[0])
-                    points = result.masks.cpu().xy[0].astype(float)
-
-                    # Determine the short label
-                    short_label = 'Review'
-                    if conf > self.main_window.get_uncertainty_thresh():
-                        short_label = class_mapping.get(cls_name, {}).get('short_label_code', 'Review')
-
-                    # Prepare the annotation data
-                    label = self.label_window.get_label_by_short_code(short_label)
-                    points = [QPointF(x, y) for x, y in points]
-
-                    # Create the rectangle annotation
-                    annotation = PolygonAnnotation(points,
-                                                   label.short_label_code,
-                                                   label.long_label_code,
-                                                   label.color,
-                                                   image_path,
-                                                   label.id,
-                                                   self.main_window.get_transparency_value(),
-                                                   show_msg=True)
-
-                    # Store the annotation and display the cropped image
-                    self.annotation_window.annotations_dict[annotation.id] = annotation
-
-                    # Connect update signals
-                    annotation.selected.connect(self.annotation_window.select_annotation)
-                    annotation.annotationDeleted.connect(self.annotation_window.delete_annotation)
-                    annotation.annotationUpdated.connect(self.main_window.confidence_window.display_cropped_image)
-
-                    # Add the prediction for the confidence window
-                    predictions = {self.label_window.get_label_by_short_code(cls_name): conf}
-                    annotation.update_machine_confidence(predictions)
-
-                    # Update label if confidence is below threshold
-                    if conf < self.main_window.get_uncertainty_thresh():
-                        review_label = self.label_window.get_label_by_id('-1')
-                        annotation.update_label(review_label)
-
-                    # Create the graphics and cropped image
-                    if image_path == self.annotation_window.current_image_path:
-                        annotation.create_graphics_item(self.annotation_window.scene)
-                        annotation.create_cropped_image(self.annotation_window.rasterio_image)
-                        self.main_window.confidence_window.display_cropped_image(annotation)
-
-                    # Update the image annotations
-                    self.main_window.image_window.update_image_annotations(image_path)
-
-                    # Update the progress bar
-                    progress_bar.update_progress()
-
-                except Exception as e:
-                    print(f"Warning: Failed to process detection result\n{e}")
+                self.process_single_segmentation_result(result, class_mapping)
+                progress_bar.update_progress()
 
         progress_bar.stop_progress()
-        progress_bar.close()
\ No newline at end of file
+        progress_bar.close()
+
+    def process_single_segmentation_result(self, result, class_mapping):
+        """
+        Process a single segmentation result.
+        
+        :param result: Segmentation result
+        :param class_mapping: Class mapping dictionary
+        """
+        try:
+            image_path = result.path.replace("\\", "/")
+            cls, cls_name, conf, points = self.extract_segmentation_result(result)
+            short_label = self.get_short_label_for_detection(cls_name, conf, class_mapping)
+            label = self.label_window.get_label_by_short_code(short_label)
+            annotation = self.create_polygon_annotation(points, label, image_path)
+            self.store_and_display_annotation(annotation, image_path, cls_name, conf)
+        except Exception as e:
+            print(f"Warning: Failed to process segmentation result\n{e}")
+
+    def extract_segmentation_result(self, result):
+        """
+        Extract relevant information from a segmentation result.
+        
+        :param result: Segmentation result
+        :return: Tuple containing class, class name, confidence, and polygon points
+        """
+        cls = int(result.boxes.cls.cpu().numpy()[0])
+        cls_name = result.names[cls]
+        conf = float(result.boxes.conf.cpu().numpy()[0])
+        points = result.masks.cpu().xy[0].astype(float)
+        return cls, cls_name, conf, points
+
+    def create_polygon_annotation(self, points, label, image_path):
+        """
+        Create a polygon annotation for the given points and label.
+        
+        :param points: List of polygon points
+        :param label: Label object
+        :param image_path: Path to the image
+        :return: PolygonAnnotation object
+        """
+        points = [QPointF(x, y) for x, y in points]
+        return PolygonAnnotation(points, 
+                                 label.short_label_code, 
+                                 label.long_label_code, 
+                                 label.color, 
+                                 image_path, 
+                                 label.id, 
+                                 self.main_window.get_transparency_value(), 
+                                 show_msg=True)
diff --git a/toolbox/MachineLearning/QtEvaluateModel.py b/toolbox/MachineLearning/QtEvaluateModel.py
index e45c032f..ff27cd28 100644
--- a/toolbox/MachineLearning/QtEvaluateModel.py
+++ b/toolbox/MachineLearning/QtEvaluateModel.py
@@ -26,16 +26,34 @@
 
 
 class EvaluateModelWorker(QThread):
+    """
+    Worker thread for evaluating a model.
+
+    Signals:
+        evaluation_started: Emitted when the evaluation starts.
+        evaluation_completed: Emitted when the evaluation completes.
+        evaluation_error: Emitted when an error occurs during evaluation.
+    """
     evaluation_started = pyqtSignal()
     evaluation_completed = pyqtSignal()
     evaluation_error = pyqtSignal(str)
 
     def __init__(self, model, params):
+        """
+        Initialize the EvaluateModelWorker.
+
+        Args:
+            model: The model to be evaluated.
+            params: A dictionary of parameters for evaluation.
+        """
         super().__init__()
         self.model = model
         self.params = params
 
     def run(self):
+        """
+        Run the evaluation process in a separate thread.
+        """
         try:
             # Emit signal to indicate evaluation has started
             self.evaluation_started.emit()
@@ -65,6 +83,13 @@ def run(self):
 
 
 class EvaluateModelDialog(QDialog):
+    """
+    Dialog for evaluating machine learning models for image classification, object detection, 
+    and instance segmentation.
+
+    :param main_window: MainWindow object
+    :param parent: Parent widget
+    """
     def __init__(self, main_window, parent=None):
         super().__init__(parent)
         self.main_window = main_window
@@ -94,6 +119,9 @@ def __init__(self, main_window, parent=None):
         self.setLayout(self.main_layout)
 
     def setup_ui(self):
+        """
+        Set up the user interface for the dialog.
+        """
         # Create a QTabWidget for different model types
         self.tab_widget = QTabWidget()
 
@@ -114,6 +142,12 @@ def setup_ui(self):
         self.main_layout.addWidget(self.cancel_button)
 
     def setup_image_classification_tab(self):
+        """
+        Set up the layout and widgets for the image classification tab.
+
+        Returns:
+            QWidget: The image classification tab widget.
+        """
         tab = QWidget()
         layout = QVBoxLayout()
 
@@ -169,6 +203,12 @@ def setup_image_classification_tab(self):
         return tab
 
     def setup_object_detection_tab(self):
+        """
+        Set up the layout and widgets for the object detection tab.
+
+        Returns:
+            QWidget: The object detection tab widget.
+        """
         tab = QWidget()
         layout = QVBoxLayout()
 
@@ -224,6 +264,12 @@ def setup_object_detection_tab(self):
         return tab
 
     def setup_instance_segmentation_tab(self):
+        """
+        Set up the layout and widgets for the instance segmentation tab.
+
+        Returns:
+            QWidget: The instance segmentation tab widget.
+        """
         tab = QWidget()
         layout = QVBoxLayout()
 
@@ -279,6 +325,9 @@ def setup_instance_segmentation_tab(self):
         return tab
 
     def browse_model_file(self):
+        """
+        Browse and select a model file.
+        """
         file_path, _ = QFileDialog.getOpenFileName(self, "Select Model File")
         if file_path:
             if self.tab_widget.currentIndex() == 0:
@@ -292,11 +341,17 @@ def browse_model_file(self):
             dir_path = os.path.dirname(os.path.dirname(file_path))
 
     def browse_dataset_dir(self):
+        """
+        Browse and select a dataset directory.
+        """
         dir_path = QFileDialog.getExistingDirectory(self, "Select Dataset Directory")
         if dir_path:
             self.dataset_dir_edit.setText(dir_path)
 
     def browse_object_detection_yaml(self):
+        """
+        Browse and select a YAML file for object detection.
+        """
         file_path, _ = QFileDialog.getOpenFileName(self,
                                                    "Select Object Detection YAML File",
                                                    "",
@@ -305,6 +360,9 @@ def browse_object_detection_yaml(self):
             self.object_detection_yaml_edit.setText(file_path)
 
     def browse_instance_segmentation_yaml(self):
+        """
+        Browse and select a YAML file for instance segmentation.
+        """
         file_path, _ = QFileDialog.getOpenFileName(self,
                                                    "Select Instance Segmentation YAML File",
                                                    "",
@@ -313,6 +371,9 @@ def browse_instance_segmentation_yaml(self):
             self.instance_segmentation_yaml_edit.setText(file_path)
 
     def browse_save_dir(self):
+        """
+        Browse and select a save directory.
+        """
         dir_path = QFileDialog.getExistingDirectory(self, "Select Save Directory")
         if dir_path:
             if self.tab_widget.currentIndex() == 0:
@@ -323,6 +384,9 @@ def browse_save_dir(self):
                 self.save_dir_edit_instance_segmentation.setText(dir_path)
 
     def accept(self):
+        """
+        Handle the OK button click event.
+        """
         if self.tab_widget.currentIndex() == 0:
             # Image Classification
             if not self.dataset_dir_edit.text():
@@ -349,6 +413,12 @@ def accept(self):
         super().accept()
 
     def get_evaluation_parameters(self):
+        """
+        Get the evaluation parameters from the dialog widgets.
+
+        Returns:
+            dict: A dictionary of evaluation parameters.
+        """
         # Extract values from dialog widgets based on the current tab
         params = {
             'verbose': True,
@@ -399,6 +469,9 @@ def get_evaluation_parameters(self):
         return params
 
     def evaluate_model(self):
+        """
+        Evaluate the model based on the provided parameters.
+        """
         # Get evaluation parameters
         self.params = self.get_evaluation_parameters()
 
@@ -423,13 +496,25 @@ def evaluate_model(self):
             print(error_message)
 
     def on_evaluation_started(self):
+        """
+        Handle the event when the evaluation starts.
+        """
         message = "Model evaluation has commenced.\nMonitor the console for real-time progress."
         QMessageBox.information(self, "Model Evaluation Status", message)
 
     def on_evaluation_completed(self):
+        """
+        Handle the event when the evaluation completes.
+        """
         message = "Model evaluation has successfully been completed."
         QMessageBox.information(self, "Model Evaluation Status", message)
 
     def on_evaluation_error(self, error_message):
+        """
+        Handle the event when an error occurs during evaluation.
+
+        Args:
+            error_message (str): The error message.
+        """
         QMessageBox.critical(self, "Error", error_message)
-        print(error_message)
\ No newline at end of file
+        print(error_message)
diff --git a/toolbox/MachineLearning/QtExportDataset.py b/toolbox/MachineLearning/QtExportDataset.py
index 55d56ab7..081cf655 100644
--- a/toolbox/MachineLearning/QtExportDataset.py
+++ b/toolbox/MachineLearning/QtExportDataset.py
@@ -33,6 +33,13 @@
 
 class ExportDatasetDialog(QDialog):
     def __init__(self, main_window, parent=None):
+        """
+        Initialize the ExportDatasetDialog class.
+
+        Args:
+            main_window: The main window object.
+            parent: The parent widget.
+        """
         super().__init__(parent)
         self.main_window = main_window
         self.annotation_window = main_window.annotation_window
@@ -84,17 +91,33 @@ def __init__(self, main_window, parent=None):
         self.radio_segmentation.toggled.connect(self.update_annotation_type_checkboxes)
 
     def showEvent(self, event):
+        """
+        Handle the show event to update annotation type checkboxes, populate class filter list,
+        and update summary statistics.
+
+        Args:
+            event: The show event.
+        """
         super().showEvent(event)
         self.update_annotation_type_checkboxes()
         self.populate_class_filter_list()
         self.update_summary_statistics()
 
     def browse_output_dir(self):
+        """
+        Browse and select an output directory.
+        """
         dir_path = QFileDialog.getExistingDirectory(self, "Select Output Directory")
         if dir_path:
             self.output_dir_edit.setText(dir_path)
 
     def get_class_mapping(self):
+        """
+        Get the class mapping for the selected labels.
+
+        Returns:
+            dict: Dictionary containing class mappings.
+        """
         # Get the label objects for the selected labels
         labels = [l for l in self.main_window.label_window.labels if l.short_label_code in self.selected_labels]
 
@@ -107,6 +130,13 @@ def get_class_mapping(self):
 
     @staticmethod
     def save_class_mapping_json(class_mapping, output_dir_path):
+        """
+        Save the class mapping dictionary as a JSON file.
+
+        Args:
+            class_mapping (dict): Dictionary containing class mappings.
+            output_dir_path (str): Path to the output directory.
+        """
         # Save the class_mapping dictionary as a JSON file
         class_mapping_path = os.path.join(output_dir_path, "class_mapping.json")
         with open(class_mapping_path, 'w') as json_file:
@@ -114,6 +144,16 @@ def save_class_mapping_json(class_mapping, output_dir_path):
 
     @staticmethod
     def merge_class_mappings(existing_mapping, new_mapping):
+        """
+        Merge the new class mappings with the existing ones without duplicates.
+
+        Args:
+            existing_mapping (dict): Existing class mappings.
+            new_mapping (dict): New class mappings.
+
+        Returns:
+            dict: Merged class mappings.
+        """
         # Merge the new class mappings with the existing ones without duplicates
         merged_mapping = existing_mapping.copy()
         for key, value in new_mapping.items():
@@ -123,6 +163,9 @@ def merge_class_mappings(existing_mapping, new_mapping):
         return merged_mapping
 
     def setup_layout(self):
+        """
+        Set up the layout for the ExportDatasetDialog.
+        """
         # Dataset Name and Output Directory
         self.dataset_name_edit = QLineEdit()
         self.output_dir_edit = QLineEdit()
@@ -206,6 +249,9 @@ def setup_layout(self):
         self.layout.addWidget(self.shuffle_button)
 
     def update_annotation_type_checkboxes(self):
+        """
+        Update the state of annotation type checkboxes based on the selected dataset type.
+        """
         if self.radio_classification.isChecked():  # Classification
             self.include_patches_checkbox.setChecked(True)
             self.include_patches_checkbox.setEnabled(False)
@@ -229,6 +275,12 @@ def update_annotation_type_checkboxes(self):
             self.include_polygons_checkbox.setEnabled(False)
 
     def filter_annotations(self):
+        """
+        Filter annotations based on the selected annotation types and current tab.
+
+        Returns:
+            list: List of filtered annotations.
+        """
         annotations = list(self.annotation_window.annotations_dict.values())
         filtered_annotations = []
 
@@ -237,22 +289,39 @@ def filter_annotations(self):
         if self.include_rectangles_checkbox.isChecked():
             filtered_annotations += [a for a in annotations if isinstance(a, RectangleAnnotation)]
         if self.include_polygons_checkbox.isChecked():
-            filtered_annotations += [a for a in annotations if isinstance(a, PolygonAnnotation)]
+            filtered_annotations += [a for a in filtered_annotations if isinstance(a, PolygonAnnotation)]
 
         return [a for a in filtered_annotations if a.label.short_label_code in self.selected_labels]
 
     def on_include_checkbox_state_changed(self, state):
+        """
+        Handle the state change event of the include checkboxes.
+
+        Args:
+            state: The new state of the checkbox.
+        """
         if state == Qt.Checked:
             self.update_summary_statistics()
         elif state == Qt.Unchecked:
             self.update_summary_statistics()
 
     def set_cell_color(self, row, column, color):
+        """
+        Set the background color of a cell in the label counts table.
+
+        Args:
+            row: The row index of the cell.
+            column: The column index of the cell.
+            color: The color to set as the background.
+        """
         item = self.label_counts_table.item(row, column)
         if item is not None:
             item.setBackground(QBrush(color))
 
     def populate_class_filter_list(self):
+        """
+        Populate the class filter list with labels and their counts.
+        """
         try:
             # Temporarily disconnect the cellChanged signal
             self.label_counts_table.cellChanged.disconnect()
@@ -319,6 +388,9 @@ def populate_class_filter_list(self):
         self.label_counts_table.cellChanged.connect(self.update_summary_statistics)
 
     def split_data(self):
+        """
+        Split the data by images based on the specified ratios.
+        """
         self.train_ratio = self.train_ratio_spinbox.value()
         self.val_ratio = self.val_ratio_spinbox.value()
         self.test_ratio = self.test_ratio_spinbox.value()
@@ -343,11 +415,20 @@ def split_data(self):
             self.test_images = images[val_split:]
 
     def determine_splits(self):
+        """
+        Determine the splits for train, validation, and test annotations.
+        """
         self.train_annotations = [a for a in self.selected_annotations if a.image_path in self.train_images]
         self.val_annotations = [a for a in self.selected_annotations if a.image_path in self.val_images]
         self.test_annotations = [a for a in self.selected_annotations if a.image_path in self.test_images]
 
     def check_label_distribution(self):
+        """
+        Check the label distribution in the splits to ensure all labels are present.
+
+        Returns:
+            bool: True if all labels are present in all splits, False otherwise.
+        """
         # Get the ratios from the spinboxes
         train_ratio = self.train_ratio_spinbox.value()
         val_ratio = self.val_ratio_spinbox.value()
@@ -399,6 +480,9 @@ def check_label_distribution(self):
         return True
 
     def update_summary_statistics(self):
+        """
+        Update the summary statistics for the dataset creation.
+        """
         if self.updating_summary_statistics:
             return
 
@@ -460,6 +544,9 @@ def update_summary_statistics(self):
         self.updating_summary_statistics = False
 
     def accept(self):
+        """
+        Handle the OK button click event to create the dataset.
+        """
         dataset_name = self.dataset_name_edit.text()
         output_dir = self.output_dir_edit.text()
         train_ratio = self.train_ratio_spinbox.value()
@@ -532,7 +619,12 @@ def accept(self):
         super().accept()
 
     def create_classification_dataset(self, output_dir_path):
+        """
+        Create an image classification dataset.
 
+        Args:
+            output_dir_path (str): Path to the output directory.
+        """
         # Create the train, val, and test directories
         train_dir = os.path.join(output_dir_path, 'train')
         val_dir = os.path.join(output_dir_path, 'val')
@@ -559,6 +651,14 @@ def create_classification_dataset(self, output_dir_path):
         pd.DataFrame(df).to_csv(f"{output_dir_path}/dataset.csv", index=False)
 
     def process_classification_annotations(self, annotations, split_dir, split):
+        """
+        Process and save classification annotations.
+
+        Args:
+            annotations (list): List of annotations.
+            split_dir (str): Path to the split directory.
+            split (str): Split name (e.g., "Training", "Validation", "Testing").
+        """
         # Get unique image paths
         image_paths = list(set(a.image_path for a in annotations))
         if not image_paths:
@@ -616,7 +716,12 @@ def save_annotations(cropped_images):
         progress_bar.close()
 
     def create_detection_dataset(self, output_dir_path):
+        """
+        Create an object detection dataset.
 
+        Args:
+            output_dir_path (str): Path to the output directory.
+        """
         # Create the yaml file
         yaml_path = os.path.join(output_dir_path, 'data.yaml')
 
@@ -653,6 +758,14 @@ def create_detection_dataset(self, output_dir_path):
         self.process_detection_annotations(self.test_annotations, test_dir, "Testing")
 
     def process_detection_annotations(self, annotations, split_dir, split):
+        """
+        Process and save detection annotations.
+
+        Args:
+            annotations (list): List of annotations.
+            split_dir (str): Path to the split directory.
+            split (str): Split name (e.g., "Training", "Validation", "Testing").
+        """
         # Get unique image paths
         image_paths = list(set(a.image_path for a in annotations))
         if not image_paths:
@@ -691,6 +804,12 @@ def process_detection_annotations(self, annotations, split_dir, split):
         progress_bar.close()
 
     def create_segmentation_dataset(self, output_dir_path):
+        """
+        Create an instance segmentation dataset.
+
+        Args:
+            output_dir_path (str): Path to the output directory.
+        """
         # Create the yaml file
         yaml_path = os.path.join(output_dir_path, 'data.yaml')
 
@@ -722,6 +841,14 @@ def create_segmentation_dataset(self, output_dir_path):
         self.process_segmentation_annotations(self.test_annotations, test_dir, "Testing")
 
     def process_segmentation_annotations(self, annotations, split_dir, split):
+        """
+        Process and save segmentation annotations.
+
+        Args:
+            annotations (list): List of annotations.
+            split_dir (str): Path to the split directory.
+            split (str): Split name (e.g., "Training", "Validation", "Testing").
+        """
         # Get unique image paths
         image_paths = list(set(a.image_path for a in annotations))
         if not image_paths:
@@ -758,4 +885,3 @@ def process_segmentation_annotations(self, annotations, split_dir, split):
 
         progress_bar.stop_progress()
         progress_bar.close()
-
diff --git a/toolbox/MachineLearning/QtImportDataset.py b/toolbox/MachineLearning/QtImportDataset.py
index 6ae0088b..e2e76d39 100644
--- a/toolbox/MachineLearning/QtImportDataset.py
+++ b/toolbox/MachineLearning/QtImportDataset.py
@@ -29,6 +29,12 @@
 
 
 class ImportDatasetDialog(QDialog):
+    """
+    Dialog for importing datasets for object detection and instance segmentation.
+
+    :param main_window: MainWindow object
+    :param parent: Parent widget
+    """
     def __init__(self, main_window, parent=None):
         super(ImportDatasetDialog, self).__init__(parent)
         self.main_window = main_window
@@ -40,6 +46,9 @@ def __init__(self, main_window, parent=None):
         self.init_ui()
 
     def init_ui(self):
+        """
+        Initialize the user interface for the ImportDatasetDialog.
+        """
         main_layout = QVBoxLayout()
         main_layout.setSpacing(10)
         main_layout.setContentsMargins(20, 20, 20, 20)
@@ -106,6 +115,9 @@ def init_ui(self):
         self.setLayout(main_layout)
 
     def browse_data_yaml(self):
+        """
+        Browse and select a data.yaml file.
+        """
         options = QFileDialog.Options()
         file_path, _ = QFileDialog.getOpenFileName(
             self, "Select data.yaml", "", "YAML Files (*.yaml);;All Files (*)", options=options
@@ -114,12 +126,18 @@ def browse_data_yaml(self):
             self.yaml_path_label.setText(file_path)
 
     def browse_output_dir(self):
+        """
+        Browse and select an output directory.
+        """
         options = QFileDialog.Options()
         dir_path = QFileDialog.getExistingDirectory(self, "Select Output Directory")
         if dir_path:
             self.output_dir_label.setText(dir_path)
 
     def accept(self):
+        """
+        Handle the OK button click event to validate and process the dataset.
+        """
         # Perform validation and processing here
         if not self.yaml_path_label.text():
             QMessageBox.warning(self, "Error", "Please select a data.yaml file.")
@@ -138,17 +156,21 @@ def accept(self):
         super().accept()
 
     def reject(self):
-        # Handle cancel action if needed
+        """
+        Handle the cancel action.
+        """
         super().reject()
 
     def process_dataset(self):
+        """
+        Process the dataset based on the selected data.yaml file and output directory.
+        """
         if not self.yaml_path_label.text():
             QMessageBox.warning(self,
                                 "No File Selected",
                                 "Please select a data.yaml file.")
             return
 
-
         output_folder = os.path.join(self.output_dir_label.text(), self.output_folder_name.text())
         os.makedirs(f"{output_folder}/images", exist_ok=True)
 
@@ -322,6 +344,12 @@ def process_dataset(self):
                                 "Dataset has been successfully imported.")
 
     def export_annotations(self, annotations, output_dir):
+        """
+        Export the annotations as a JSON file in the specified output directory.
+
+        :param annotations: List of annotations to export
+        :param output_dir: Path to the output directory
+        """
         QApplication.setOverrideCursor(Qt.WaitCursor)
 
         progress_bar = ProgressBar(self.annotation_window, title="Exporting Annotations")
@@ -371,4 +399,4 @@ def export_annotations(self, annotations, output_dir):
         progress_bar.close()
 
         # Make the cursor normal again
-        QApplication.restoreOverrideCursor()
\ No newline at end of file
+        QApplication.restoreOverrideCursor()
diff --git a/toolbox/MachineLearning/QtMergeDatasets.py b/toolbox/MachineLearning/QtMergeDatasets.py
index 7ae14c59..0ef841ac 100644
--- a/toolbox/MachineLearning/QtMergeDatasets.py
+++ b/toolbox/MachineLearning/QtMergeDatasets.py
@@ -21,27 +21,39 @@
 
 
 class MergeDatasetsDialog(QDialog):
+    """
+    Dialog for merging multiple datasets into a single dataset for machine learning tasks such as
+    image classification, object detection, and instance segmentation.
+    """
+
     def __init__(self, parent=None):
+        """
+        Initializes the MergeDatasetsDialog.
+
+        :param parent: Parent widget, default is None.
+        """
         super().__init__(parent)
         self.setWindowTitle("Merge Datasets")
         self.resize(400, 200)
 
         self.layout = QVBoxLayout(self)
 
-        # Create tabs
+        # Create tabs for different dataset types
         self.tabs = QTabWidget()
         self.tab_classification = QWidget()
         self.tab_detection = QWidget()
         self.tab_segmentation = QWidget()
 
         self.tabs.addTab(self.tab_classification, "Image Classification")
+        # Uncomment the following lines to add tabs for detection and segmentation
         # self.tabs.addTab(self.tab_detection, "Object Detection")
         # self.tabs.addTab(self.tab_segmentation, "Instance Segmentation")
 
         self.layout.addWidget(self.tabs)
 
-        # Setup tabs
+        # Setup each tab with its respective layout and widgets
         self.setup_tab(self.tab_classification, "classification")
+        # Uncomment the following lines to setup detection and segmentation tabs
         # self.setup_tab(self.tab_detection, "detection")
         # self.setup_tab(self.tab_segmentation, "segmentation")
 
@@ -51,21 +63,27 @@ def __init__(self, parent=None):
         self.buttons.rejected.connect(self.reject)
         self.layout.addWidget(self.buttons)
 
-        # Track valid directories and their class mappings
+        # List to track valid directories and their class mappings
         self.valid_directories = []
 
     def setup_tab(self, tab, tab_type):
+        """
+        Sets up the layout and widgets for a given tab.
+
+        :param tab: The QWidget representing the tab.
+        :param tab_type: The type of task (e.g., "classification").
+        """
         layout = QVBoxLayout()
 
-        # Dataset Name
+        # Dataset Name Input
         dataset_name_edit = QLineEdit()
-        dataset_name_edit.setObjectName("dataset_name_edit")  # Set object name
+        dataset_name_edit.setObjectName("dataset_name_edit")  # Set object name for access
         layout.addWidget(QLabel("Dataset Name:"))
         layout.addWidget(dataset_name_edit)
 
         # Output Directory Chooser
         output_dir_edit = QLineEdit()
-        output_dir_edit.setObjectName("output_dir_edit")  # Set object name
+        output_dir_edit.setObjectName("output_dir_edit")  # Set object name for access
         output_dir_button = QPushButton("Browse...")
         output_dir_button.clicked.connect(lambda: self.browse_output_directory(output_dir_edit))
         output_dir_layout = QHBoxLayout()
@@ -79,11 +97,11 @@ def setup_tab(self, tab, tab_type):
         layout.addWidget(QLabel("Existing Dataset Directories:"))
         layout.addLayout(existing_dirs_layout)
 
-        # Add two default choosers
+        # Add two default directory choosers
         self.add_chooser(existing_dirs_layout, tab, tab_type)
         self.add_chooser(existing_dirs_layout, tab, tab_type)
 
-        # Add Directory Button
+        # Button to add additional dataset directories
         add_dir_button = QPushButton("Add Dataset")
         add_dir_button.clicked.connect(lambda: self.add_chooser(existing_dirs_layout, tab, tab_type))
         layout.addWidget(add_dir_button)
@@ -91,35 +109,47 @@ def setup_tab(self, tab, tab_type):
         tab.setLayout(layout)
 
     def add_chooser(self, existing_dirs_layout, tab, tab_type):
+        """
+        Adds a directory chooser widget to the existing directories layout.
+
+        :param existing_dirs_layout: The QVBoxLayout to add the chooser to.
+        :param tab: The current tab QWidget.
+        :param tab_type: The type of task (e.g., "classification").
+        """
         chooser = QWidget()
         chooser_layout = QHBoxLayout(chooser)
 
+        # Status label to indicate validation status
         status_label = QLabel()
         chooser_layout.addWidget(status_label)
 
         if tab_type == "classification":
+            # Directory chooser for classification tasks
             dir_edit = QLineEdit()
             dir_button = QPushButton("Browse Directory...")
             dir_button.clicked.connect(lambda: self.browse_existing_directory(dir_edit, status_label, tab))
             chooser_layout.addWidget(dir_edit)
             chooser_layout.addWidget(dir_button)
 
+            # Class mapping chooser
             class_mapping_edit = QLineEdit()
-            class_mapping_edit.setObjectName("class_mapping_edit")  # Set object name
+            class_mapping_edit.setObjectName("class_mapping_edit")  # Set object name for access
             class_mapping_button = QPushButton("Select Class Mapping")
             class_mapping_button.clicked.connect(
                 lambda: self.browse_class_mapping(class_mapping_edit, status_label, tab))
             chooser_layout.addWidget(class_mapping_edit)
             chooser_layout.addWidget(class_mapping_button)
         else:
+            # Directory chooser for other task types (e.g., detection, segmentation)
             yaml_edit = QLineEdit()
             yaml_button = QPushButton("Select YAML")
             yaml_button.clicked.connect(lambda: self.browse_data_yaml(yaml_edit, status_label, tab))
             chooser_layout.addWidget(yaml_edit)
             chooser_layout.addWidget(yaml_button)
 
+            # Class mapping chooser
             class_mapping_edit = QLineEdit()
-            class_mapping_edit.setObjectName("class_mapping_edit")  # Set object name
+            class_mapping_edit.setObjectName("class_mapping_edit")  # Set object name for access
             class_mapping_button = QPushButton("Select Class Mapping")
             class_mapping_button.clicked.connect(
                 lambda: self.browse_class_mapping(class_mapping_edit, status_label, tab))
@@ -129,11 +159,23 @@ def add_chooser(self, existing_dirs_layout, tab, tab_type):
         existing_dirs_layout.addWidget(chooser)
 
     def browse_output_directory(self, output_dir_edit):
+        """
+        Opens a dialog to select the output directory and sets the selected path.
+
+        :param output_dir_edit: The QLineEdit widget to display the selected directory.
+        """
         dir_path = QFileDialog.getExistingDirectory(self, "Select Output Directory")
         if dir_path:
             output_dir_edit.setText(dir_path)
 
     def browse_existing_directory(self, dir_edit, status_label, tab):
+        """
+        Opens a dialog to select an existing dataset directory, sets the path, and validates it.
+
+        :param dir_edit: The QLineEdit widget to display the selected directory.
+        :param status_label: The QLabel to display validation status.
+        :param tab: The current tab QWidget.
+        """
         dir_path = QFileDialog.getExistingDirectory(self, "Select Existing Dataset Directory")
         if dir_path:
             dir_edit.setText(dir_path)
@@ -148,6 +190,13 @@ def browse_existing_directory(self, dir_edit, status_label, tab):
                     self.validate_class_mapping(class_mapping_path, status_label, tab)
 
     def browse_data_yaml(self, yaml_edit, status_label, tab):
+        """
+        Opens a dialog to select a YAML file, sets the path, and validates it.
+
+        :param yaml_edit: The QLineEdit widget to display the selected YAML file.
+        :param status_label: The QLabel to display validation status.
+        :param tab: The current tab QWidget.
+        """
         options = QFileDialog.Options()
         file_path, _ = QFileDialog.getOpenFileName(
             self, "Select data.yaml", "", "YAML Files (*.yaml);;All Files (*)", options=options
@@ -166,6 +215,13 @@ def browse_data_yaml(self, yaml_edit, status_label, tab):
                     self.validate_class_mapping(class_mapping_path, status_label, tab)
 
     def browse_class_mapping(self, class_mapping_edit, status_label, tab):
+        """
+        Opens a dialog to select a class mapping JSON file, sets the path, and validates it.
+
+        :param class_mapping_edit: The QLineEdit widget to display the selected JSON file.
+        :param status_label: The QLabel to display validation status.
+        :param tab: The current tab QWidget.
+        """
         options = QFileDialog.Options()
         file_path, _ = QFileDialog.getOpenFileName(
             self, "Select class_mapping.json", "", "JSON Files (*.json);;All Files (*)", options=options
@@ -175,14 +231,29 @@ def browse_class_mapping(self, class_mapping_edit, status_label, tab):
             self.validate_class_mapping(file_path, status_label, tab)
 
     def validate_directory(self, dir_path, status_label, tab):
+        """
+        Validates if the selected directory exists and updates the status.
+
+        :param dir_path: Path to the directory to validate.
+        :param status_label: The QLabel to display validation status.
+        :param tab: The current tab QWidget.
+        """
         if os.path.exists(dir_path):
             status_label.setText("✅")
             self.valid_directories.append((dir_path, None, tab))
         else:
             status_label.setText("❌")
+            # Remove invalid directory from the list
             self.valid_directories = [(d, c, t) for d, c, t in self.valid_directories if d != dir_path]
 
     def validate_yaml(self, yaml_path, status_label, tab):
+        """
+        Validates the selected YAML file for required fields.
+
+        :param yaml_path: Path to the YAML file to validate.
+        :param status_label: The QLabel to display validation status.
+        :param tab: The current tab QWidget.
+        """
         try:
             with open(yaml_path, 'r') as file:
                 data = yaml.safe_load(file)
@@ -191,17 +262,27 @@ def validate_yaml(self, yaml_path, status_label, tab):
                     self.valid_directories.append((yaml_path, None, tab))
                 else:
                     status_label.setText("❌")
+                    # Remove invalid YAML from the list
                     self.valid_directories = [(d, c, t) for d, c, t in self.valid_directories if d != yaml_path]
         except Exception as e:
             status_label.setText("❌")
+            # Remove invalid YAML from the list
             self.valid_directories = [(d, c, t) for d, c, t in self.valid_directories if d != yaml_path]
 
     def validate_class_mapping(self, class_mapping_path, status_label, tab):
+        """
+        Validates the selected class mapping JSON file.
+
+        :param class_mapping_path: Path to the class mapping JSON file.
+        :param status_label: The QLabel to display validation status.
+        :param tab: The current tab QWidget.
+        """
         try:
             with open(class_mapping_path, 'r') as file:
                 data = json.load(file)
                 if isinstance(data, dict):
                     status_label.setText("✅")
+                    # Update class mapping in the valid directories list
                     for d, c, t in self.valid_directories:
                         if t == tab and c is None:
                             self.valid_directories.remove((d, c, t))
@@ -213,36 +294,47 @@ def validate_class_mapping(self, class_mapping_path, status_label, tab):
             status_label.setText("❌")
 
     def merge_datasets(self):
+        """
+        Merges the selected datasets into a single output directory.
+        """
         current_tab = self.tabs.currentWidget()
         if current_tab not in [self.tab_classification]:
             QMessageBox.warning(self, "Warning", "Only Image Classification merging has been implemented.")
             return
 
+        # Retrieve output directory from the current tab
         output_dir = current_tab.findChild(QLineEdit, "output_dir_edit").text()
         if not output_dir:
             QMessageBox.warning(self, "Input Error", "Output directory must be specified.")
             return
 
+        # Retrieve dataset name from the current tab
         dataset_name = current_tab.findChild(QLineEdit, "dataset_name_edit").text()
         if not dataset_name:
             QMessageBox.warning(self, "Input Error", "Dataset name must be specified.")
             return
 
-        # Make cursor busy
+        # Make cursor busy to indicate processing
         QApplication.setOverrideCursor(Qt.WaitCursor)
 
+        # Create the output directory
         output_dir_path = os.path.join(output_dir, dataset_name)
         os.makedirs(output_dir_path, exist_ok=True)
 
         merged_class_mapping = {}
 
         def copy_directory(src, dest):
+            """
+            Copies contents from the source directory to the destination directory.
+
+            :param src: Source directory path.
+            :param dest: Destination directory path.
+            """
             shutil.copytree(src, dest, dirs_exist_ok=True)
 
         with ThreadPoolExecutor() as executor:
             futures = []
             for dir_path, class_mapping_path, tab in self.valid_directories:
-
                 if tab == self.tab_classification:
                     if class_mapping_path:
                         try:
@@ -252,6 +344,7 @@ def copy_directory(src, dest):
                         except Exception as e:
                             print(f"Error reading class mapping: {e}")
 
+                    # Copy train, val, and test splits
                     for split in ['train', 'val', 'test']:
                         src_split_dir = os.path.join(dir_path, split)
                         dest_split_dir = os.path.join(output_dir_path, split)
@@ -259,20 +352,21 @@ def copy_directory(src, dest):
                             future = executor.submit(copy_directory, src_split_dir, dest_split_dir)
                             futures.append(future)
 
-            # Create a progress dialog
+            # Create and display a progress bar
             progress_bar = ProgressBar(self, title=f"Merging Datasets")
             progress_bar.show()
             progress_bar.start_progress(len(futures))
 
-            # Wait for all copying tasks to complete
+            # Wait for all copy operations to complete
             for i, future in enumerate(as_completed(futures)):
                 future.result()
                 progress_bar.update_progress()
 
+            # Close the progress bar
             progress_bar.stop_progress()
             progress_bar.close()
 
-        # Check if the merged class mapping is empty
+        # Save the merged class mapping if available
         if merged_class_mapping:
             merged_class_mapping_path = os.path.join(output_dir_path, "class_mapping.json")
             with open(merged_class_mapping_path, 'w') as json_file:
@@ -280,9 +374,12 @@ def copy_directory(src, dest):
 
         QMessageBox.information(self, "Success", "Datasets merged successfully!")
 
-        # Restore cursor
+        # Restore the cursor to default
         QApplication.restoreOverrideCursor()
 
     def accept(self):
+        """
+        Overrides the accept method to perform dataset merging before closing the dialog.
+        """
         self.merge_datasets()
         super().accept()
\ No newline at end of file
diff --git a/toolbox/MachineLearning/QtOptimizeModel.py b/toolbox/MachineLearning/QtOptimizeModel.py
index 9938e6df..fb8e45e0 100644
--- a/toolbox/MachineLearning/QtOptimizeModel.py
+++ b/toolbox/MachineLearning/QtOptimizeModel.py
@@ -15,7 +15,17 @@
 
 
 class OptimizeModelDialog(QDialog):
+    """
+    A dialog for optimizing and exporting a YOLO model with custom parameters.
+    """
+
     def __init__(self, main_window, parent=None):
+        """
+        Initialize the dialog with the main window and optional parent widget.
+
+        param main_window: The main window of the application.
+        param parent: The parent widget of the dialog.
+        """
         super().__init__(parent)
         self.main_window = main_window
 
@@ -29,26 +39,23 @@ def __init__(self, main_window, parent=None):
         # Create a QLabel with explanatory text and hyperlink
         info_label = QLabel("Details on different production formats can be found "
                             "<a href='https://docs.ultralytics.com/modes/export/#export-formats'>here</a>.")
-
         info_label.setOpenExternalLinks(True)
         info_label.setWordWrap(True)
         self.layout.addWidget(info_label)
 
+        # Button to browse for model file
         browse_button = QPushButton("Browse")
         browse_button.clicked.connect(self.browse_file)
         self.layout.addWidget(browse_button)
 
+        # Text area to display selected model file path
         self.model_text_area = QTextEdit("No model file selected")
         self.model_text_area.setReadOnly(True)
         self.layout.addWidget(self.model_text_area)
 
         # Export Format Dropdown
         self.export_format_combo = QComboBox()
-        self.export_format_combo.addItems(["torchscript",
-                                           "onnx",
-                                           "openvino",
-                                           "engine"])
-
+        self.export_format_combo.addItems(["torchscript", "onnx", "openvino", "engine"])
         self.export_format_combo.setEditable(True)
         self.layout.addWidget(QLabel("Select or Enter Export Format:"))
         self.layout.addWidget(self.export_format_combo)
@@ -67,6 +74,7 @@ def __init__(self, main_window, parent=None):
 
         self.layout.addLayout(self.form_layout)
 
+        # Accept button to start optimization
         accept_button = QPushButton("Accept")
         accept_button.clicked.connect(self.optimize_model)
         self.layout.addWidget(accept_button)
@@ -74,6 +82,9 @@ def __init__(self, main_window, parent=None):
         self.setLayout(self.layout)
 
     def add_parameter_pair(self):
+        """
+        Add a new pair of parameter name and value input fields.
+        """
         param_layout = QHBoxLayout()
         param_name = QLineEdit()
         param_value = QLineEdit()
@@ -84,20 +95,30 @@ def add_parameter_pair(self):
         self.custom_params_layout.addLayout(param_layout)
 
     def browse_file(self):
+        """
+        Open a file dialog to select a model file and display its path.
+        """
         options = QFileDialog.Options()
-        file_path, _ = QFileDialog.getOpenFileName(self,
-                                                   "Open Model File", "",
-                                                   "Model Files (*.pt)", options=options)
+        file_path, _ = QFileDialog.getOpenFileName(self, 
+                                                   "Open Model File", 
+                                                   "", 
+                                                   "Model Files (*.pt)", 
+                                                   options=options)
         if file_path:
             self.model_path = file_path
             self.model_text_area.setText("Model file selected")
 
     def accept(self):
+        """
+        Override the accept method to optimize the model before closing the dialog.
+        """
         self.optimize_model()
         super().accept()
 
     def get_optimization_parameters(self):
-        # Extract values from dialog widgets
+        """
+        Extract and return the optimization parameters from the dialog widgets.
+        """
         params = {'format': self.export_format_combo.currentText()}
 
         for param_name, param_value in self.custom_params:
@@ -117,19 +138,19 @@ def get_optimization_parameters(self):
                         except ValueError:
                             params[name] = value
 
-        # Return the dictionary of parameters
         return params
 
     def optimize_model(self):
-
-        # Get training parameters
+        """
+        Optimize and export the model using the specified parameters.
+        """
         params = self.get_optimization_parameters()
 
         # Set the cursor to waiting (busy) cursor
         QApplication.setOverrideCursor(Qt.WaitCursor)
 
         try:
-            # Initialize the model, export given params
+            # Initialize the model and export with given params
             YOLO(self.model_path).export(**params)
 
             message = "Model export successful."
diff --git a/toolbox/MachineLearning/QtTrainModel.py b/toolbox/MachineLearning/QtTrainModel.py
index 3af63879..446cd244 100644
--- a/toolbox/MachineLearning/QtTrainModel.py
+++ b/toolbox/MachineLearning/QtTrainModel.py
@@ -30,17 +30,35 @@
 
 
 class TrainModelWorker(QThread):
+    """
+    Worker thread for training a model.
+
+    Signals:
+        training_started: Emitted when the training starts.
+        training_completed: Emitted when the training completes.
+        training_error: Emitted when an error occurs during training.
+    """
     training_started = pyqtSignal()
     training_completed = pyqtSignal()
     training_error = pyqtSignal(str)
 
     def __init__(self, params, device):
+        """
+        Initialize the TrainModelWorker.
+
+        Args:
+            params: A dictionary of parameters for training.
+            device: The device to use for training (e.g., 'cpu' or 'cuda').
+        """
         super().__init__()
         self.params = params
         self.device = device
         self.model = None
 
     def run(self):
+        """
+        Run the training process in a separate thread.
+        """
         try:
             # Emit signal to indicate training has started
             self.training_started.emit()
@@ -69,6 +87,9 @@ def run(self):
             self._cleanup()
 
     def evaluate_model(self):
+        """
+        Evaluate the model after training.
+        """
         try:
             # Create an instance of EvaluateModelWorker and start it
             eval_params = {
@@ -88,22 +109,43 @@ def evaluate_model(self):
             self.training_error.emit(str(e))
 
     def on_evaluation_started(self):
+        """
+        Handle the event when the evaluation starts.
+        """
         pass
 
     def on_evaluation_completed(self):
+        """
+        Handle the event when the evaluation completes.
+        """
         pass
 
     def on_evaluation_error(self, error_message):
-        # Handle any errors that occur during evaluation
+        """
+        Handle the event when an error occurs during evaluation.
+
+        Args:
+            error_message (str): The error message.
+        """
         self.training_error.emit(error_message)
 
     def _cleanup(self):
+        """
+        Clean up resources after training.
+        """
         del self.model
         gc.collect()
         empty_cache()
 
 
 class TrainModelDialog(QDialog):
+    """
+    Dialog for training machine learning models for image classification, object detection, 
+    and instance segmentation.
+
+    :param main_window: MainWindow object
+    :param parent: Parent widget
+    """
     def __init__(self, main_window, parent=None):
         super().__init__(parent)
         self.main_window = main_window
@@ -145,6 +187,9 @@ def __init__(self, main_window, parent=None):
         self.layout().addWidget(scroll_area)
 
     def setup_ui(self):
+        """
+        Set up the user interface for the dialog.
+        """
         # Create a QLabel with explanatory text and hyperlink
         info_label = QLabel("Details on different hyperparameters can be found "
                             "<a href='https://docs.ultralytics.com/modes/train/#train-settings'>here</a>.")
@@ -314,6 +359,9 @@ def setup_ui(self):
         self.main_layout.addWidget(self.cancel_button)
 
     def add_parameter_pair(self):
+        """
+        Add a new pair of parameter name and value input fields.
+        """
         param_layout = QHBoxLayout()
         param_name = QLineEdit()
         param_value = QLineEdit()
@@ -324,6 +372,9 @@ def add_parameter_pair(self):
         self.custom_params_layout.addLayout(param_layout)
 
     def browse_dataset_dir(self):
+        """
+        Browse and select a dataset directory.
+        """
         dir_path = QFileDialog.getExistingDirectory(self, "Select Dataset Directory")
         if dir_path:
             # Load the class mapping if it exists
@@ -335,6 +386,9 @@ def browse_dataset_dir(self):
             self.classify_dataset_edit.setText(dir_path)
 
     def browse_dataset_yaml(self):
+        """
+        Browse and select a dataset YAML file.
+        """
         file_path, _ = QFileDialog.getOpenFileName(self,
                                                    "Select Dataset YAML File",
                                                    "",
@@ -354,6 +408,9 @@ def browse_dataset_yaml(self):
                 self.segmentation_mapping_edit.setText(class_mapping_path)
 
     def browse_class_mapping_file(self):
+        """
+        Browse and select a class mapping file.
+        """
         file_path, _ = QFileDialog.getOpenFileName(self,
                                                    "Select Class Mapping File",
                                                    "",
@@ -369,16 +426,25 @@ def browse_class_mapping_file(self):
                 self.segmentation_mapping_edit.setText(file_path)
 
     def browse_project_dir(self):
+        """
+        Browse and select a project directory.
+        """
         dir_path = QFileDialog.getExistingDirectory(self, "Select Project Directory")
         if dir_path:
             self.project_edit.setText(dir_path)
 
     def browse_model_file(self):
+        """
+        Browse and select a model file.
+        """
         file_path, _ = QFileDialog.getOpenFileName(self, "Select Model File")
         if file_path:
             self.model_edit.setText(file_path)
 
     def setup_classification_tab(self):
+        """
+        Set up the layout and widgets for the image classification tab.
+        """
         layout = QVBoxLayout()
 
         # Dataset Directory
@@ -416,6 +482,9 @@ def setup_classification_tab(self):
         self.tab_classification.setLayout(layout)
 
     def setup_detection_tab(self):
+        """
+        Set up the layout and widgets for the object detection tab.
+        """
         layout = QVBoxLayout()
 
         self.detection_dataset_edit = QLineEdit()
@@ -452,6 +521,9 @@ def setup_detection_tab(self):
         self.tab_detection.setLayout(layout)
 
     def setup_segmentation_tab(self):
+        """
+        Set up the layout and widgets for the instance segmentation tab.
+        """
         layout = QVBoxLayout()
 
         self.segmentation_dataset_edit = QLineEdit()
@@ -488,11 +560,19 @@ def setup_segmentation_tab(self):
         self.tab_segmentation.setLayout(layout)
 
     def accept(self):
+        """
+        Handle the OK button click event.
+        """
         self.train_model()
         super().accept()
 
     def get_parameters(self):
+        """
+        Get the training parameters from the dialog widgets.
 
+        Returns:
+            dict: A dictionary of training parameters.
+        """
         # Determine the selected tab
         selected_tab = self.tabs.currentWidget()
         if selected_tab == self.tab_classification:
@@ -566,7 +646,9 @@ def get_parameters(self):
         return params
 
     def train_model(self):
-
+        """
+        Train the model based on the provided parameters.
+        """
         # Get training parameters
         self.params = self.get_parameters()
         # Create and start the worker thread
@@ -577,6 +659,9 @@ def train_model(self):
         self.worker.start()
 
     def on_training_started(self):
+        """
+        Handle the event when the training starts.
+        """
         # Save the class mapping JSON file
         output_dir_path = os.path.join(self.params['project'], self.params['name'])
         os.makedirs(output_dir_path, exist_ok=True)
@@ -589,9 +674,18 @@ def on_training_started(self):
         QMessageBox.information(self, "Model Training Status", message)
 
     def on_training_completed(self):
+        """
+        Handle the event when the training completes.
+        """
         message = "Model training has successfully been completed."
         QMessageBox.information(self, "Model Training Status", message)
 
     def on_training_error(self, error_message):
+        """
+        Handle the event when an error occurs during training.
+
+        Args:
+            error_message (str): The error message.
+        """
         QMessageBox.critical(self, "Error", error_message)
         print(error_message)
diff --git a/toolbox/QtAutoDistill.py b/toolbox/QtAutoDistill.py
index 6ea7890a..787f45ae 100644
--- a/toolbox/QtAutoDistill.py
+++ b/toolbox/QtAutoDistill.py
@@ -24,7 +24,19 @@
 
 
 class AutoDistillDeployModelDialog(QDialog):
+    """
+    Dialog for deploying and managing AutoDistill models.
+    Allows users to load, configure, and deactivate models, as well as make predictions on images.
+    """
+
     def __init__(self, main_window, parent=None):
+        """
+        Initialize the AutoDistillDeployModelDialog.
+
+        Args:
+            main_window: The main application window.
+            parent: The parent widget, default is None.
+        """
         super().__init__(parent)
         self.main_window = main_window
         self.label_window = main_window.label_window
@@ -122,22 +134,47 @@ def __init__(self, main_window, parent=None):
         self.main_layout.addWidget(self.status_bar)
 
     def showEvent(self, event):
+        """
+        Handle the show event to update label options.
+
+        Args:
+            event: The event object.
+        """
         super().showEvent(event)
         self.update_label_options()
 
     def update_area_label(self, min_val, max_val):
+        """
+        Update the area threshold label based on slider values.
+
+        Args:
+            min_val: Minimum slider value.
+            max_val: Maximum slider value.
+        """
         # Update the area threshold values
         self.area_thresh_min = min_val / 100.0
         self.area_thresh_max = max_val / 100.0
         self.area_threshold_label.setText(f"Area Threshold: {self.area_thresh_min:.2f} - {self.area_thresh_max:.2f}")
 
     def get_area_thresh(self, image_path):
+        """
+        Calculate area thresholds based on image dimensions.
+
+        Args:
+            image_path: Path to the image.
+
+        Returns:
+            Tuple of (min_area_thresh, max_area_thresh).
+        """
         h, w = self.main_window.image_window.rasterio_open(image_path).shape
         area_thresh_min = (h * w) * self.area_thresh_min
         area_thresh_max = (h * w) * self.area_thresh_max
         return area_thresh_min, area_thresh_max
 
     def update_uncertainty_label(self):
+        """
+        Update the uncertainty threshold label based on slider value.
+        """
         # Convert the slider value to a ratio (0-1)
         value = self.uncertainty_threshold_slider.value() / 100.0
         self.main_window.update_uncertainty_thresh(value)
@@ -145,12 +182,21 @@ def update_uncertainty_label(self):
         self.uncertainty_thresh = self.uncertainty_threshold_slider.value() / 100.0
 
     def on_uncertainty_changed(self, value):
+        """
+        Update the slider and label when the uncertainty threshold changes.
+
+        Args:
+            value: New uncertainty threshold value.
+        """
         # Update the slider and label when the shared data changes
         self.uncertainty_threshold_slider.setValue(int(value * 100))
         self.uncertainty_threshold_label.setText(f"{value:.2f}")
         self.uncertainty_thresh = self.uncertainty_threshold_slider.value() / 100.0
 
     def update_label_options(self):
+        """
+        Update the label options in ontology pairs based on available labels.
+        """
         label_options = [label.short_label_code for label in self.label_window.labels]
         for _, label_dropdown in self.ontology_pairs:
             previous_label = label_dropdown.currentText()
@@ -160,6 +206,9 @@ def update_label_options(self):
                 label_dropdown.setCurrentText(previous_label)
 
     def add_ontology_pair(self):
+        """
+        Add a new ontology pair input (text input and label dropdown).
+        """
         pair_layout = QHBoxLayout()
 
         text_input = QLineEdit()
@@ -174,6 +223,9 @@ def add_ontology_pair(self):
         self.ontology_layout.insertLayout(self.ontology_layout.count() - 1, pair_layout)
 
     def remove_ontology_pair(self):
+        """
+        Remove the last ontology pair input if more than one exists.
+        """
         if len(self.ontology_pairs) > 1:
             pair = self.ontology_pairs.pop()
             pair[0].deleteLater()
@@ -184,6 +236,9 @@ def remove_ontology_pair(self):
             item.layout().deleteLater()
 
     def load_model(self):
+        """
+        Load the selected model with the current configuration.
+        """
         # Make cursor busy
         QApplication.setOverrideCursor(Qt.WaitCursor)
         # Show a progress bar
@@ -191,30 +246,19 @@ def load_model(self):
         progress_bar.show()
         try:
             # Get the ontology mapping
-            ontology_mapping = {}
-            for text_input, label_dropdown in self.ontology_pairs:
-                if text_input.text() != "":
-                    ontology_mapping[text_input.text()] = label_dropdown.currentText()
+            ontology_mapping = self.get_ontology_mapping()
 
             # Set the ontology
             self.ontology = CaptionOntology(ontology_mapping)
 
             # Threshold for confidence
-            if self.main_window.get_uncertainty_thresh() < 0.10:
-                uncertainty_thresh = self.main_window.get_uncertainty_thresh()
-            else:
-                uncertainty_thresh = 0.10  # Arbitrary value to prevent too many detections
+            uncertainty_thresh = self.get_uncertainty_threshold()
 
             # Get the name of the model to load
             model_name = self.model_dropdown.currentText()
 
             if model_name != self.model_name:
-                if model_name == "GroundingDINO":
-                    from autodistill_grounding_dino import GroundingDINO
-                    self.model_name = model_name
-                    self.loaded_model = GroundingDINO(ontology=self.ontology,
-                                                      box_threshold=uncertainty_thresh,
-                                                      text_threshold=uncertainty_thresh)
+                self.load_new_model(model_name, uncertainty_thresh)
             else:
                 # Update the model with the new ontology
                 self.loaded_model.ontology = self.ontology
@@ -233,7 +277,53 @@ def load_model(self):
         # Exit the dialog box
         self.accept()
 
+    def get_ontology_mapping(self):
+        """
+        Retrieve the ontology mapping from user inputs.
+
+        Returns:
+            Dictionary mapping texts to label codes.
+        """
+        ontology_mapping = {}
+        for text_input, label_dropdown in self.ontology_pairs:
+            if text_input.text() != "":
+                ontology_mapping[text_input.text()] = label_dropdown.currentText()
+        return ontology_mapping
+
+    def get_uncertainty_threshold(self):
+        """
+        Get the uncertainty threshold, limiting it to a maximum of 0.10.
+
+        Returns:
+            Adjusted uncertainty threshold value.
+        """
+        if self.main_window.get_uncertainty_thresh() < 0.10:
+            return self.main_window.get_uncertainty_thresh()
+        else:
+            return 0.10  # Arbitrary value to prevent too many detections
+
+    def load_new_model(self, model_name, uncertainty_thresh):
+        """
+        Load a new model based on the selected model name.
+
+        Args:
+            model_name: Name of the model to load.
+            uncertainty_thresh: Threshold for uncertainty.
+        """
+        if model_name == "GroundingDINO":
+            from autodistill_grounding_dino import GroundingDINO
+            self.model_name = model_name
+            self.loaded_model = GroundingDINO(ontology=self.ontology,
+                                              box_threshold=uncertainty_thresh,
+                                              text_threshold=uncertainty_thresh)
+
     def predict(self, image_paths=None):
+        """
+        Make predictions on the given image paths using the loaded model.
+
+        Args:
+            image_paths: List of image paths to process. If None, uses the current image.
+        """
         if not self.loaded_model:
             QMessageBox.critical(self, "Error", "No model loaded")
             return
@@ -260,6 +350,13 @@ def predict(self, image_paths=None):
         empty_cache()
 
     def process_results(self, image_path, results):
+        """
+        Process the prediction results and create annotations.
+
+        Args:
+            image_path: Path to the image being processed.
+            results: Prediction results to process.
+        """
         progress_bar = ProgressBar(self, title=f"Making Predictions")
         progress_bar.show()
         progress_bar.start_progress(len(results))
@@ -329,6 +426,9 @@ def process_results(self, image_path, results):
         progress_bar.close()
 
     def deactivate_model(self):
+        """
+        Deactivate the currently loaded model and clean up resources.
+        """
         self.loaded_model = None
         self.model_name = None
         gc.collect()
diff --git a/toolbox/QtSAM.py b/toolbox/QtSAM.py
index 113ccb30..9a8e54a4 100644
--- a/toolbox/QtSAM.py
+++ b/toolbox/QtSAM.py
@@ -10,7 +10,7 @@
 
 from PyQt5.QtCore import Qt
 from PyQt5.QtWidgets import (QDialog, QVBoxLayout, QFormLayout, QSpinBox, QSlider, QLabel, QHBoxLayout, QPushButton,
-                             QTabWidget, QComboBox, QMessageBox, QApplication, QWidget, QCheckBox)
+                             QTabWidget, QComboBox, QMessageBox, QApplication, QWidget)
 
 from torch.cuda import empty_cache
 from mobile_sam import SamPredictor as MobileSamPredictor
@@ -25,6 +25,8 @@
 from ultralytics.utils.downloads import attempt_download_asset
 
 from toolbox.QtProgressBar import ProgressBar
+from toolbox.utilities import preprocess_image
+from toolbox.utilities import rasterio_to_numpy
 
 
 # ----------------------------------------------------------------------------------------------------------------------
@@ -32,19 +34,17 @@
 # ----------------------------------------------------------------------------------------------------------------------
 
 
-def custom_postprocess(masks, scores, logits, resized_image, original_image):
+def to_ultralytics(masks, scores, original_image):
     """
     Post-processes SAM's inference outputs to generate object detection masks and bounding boxes.
 
     Args:
         masks (torch.Tensor): Predicted masks with shape (N, 1, H, W).
         scores (torch.Tensor): Confidence scores for each mask with shape (N, 1).
-        logits (torch.Tensor): Logits for each mask with shape (N, 1, H, W).
-        resized_image (np.ndarray): The resized image used for inference.
         original_image (np.ndarray): The original, unprocessed image.
 
     Returns:
-        (Results): Results object containing detection masks, bounding boxes, and other metadata.
+        (Results): Ultralytics Results object containing detection masks, bounding boxes, and other metadata.
     """
     # Ensure the original image is in the correct format
     if not isinstance(original_image, np.ndarray):
@@ -87,9 +87,7 @@ def custom_postprocess(masks, scores, logits, resized_image, original_image):
     names = dict(enumerate(str(i) for i in range(len(masks))))
 
     # Create Results object
-    result = Results(original_image, path="", names=names, masks=scaled_masks, boxes=pred_bboxes)
-
-    return result
+    return Results(original_image, path="", names=names, masks=scaled_masks, boxes=pred_bboxes)
 
 
 # ----------------------------------------------------------------------------------------------------------------------
@@ -99,6 +97,9 @@ def custom_postprocess(masks, scores, logits, resized_image, original_image):
 
 class SAMDeployModelDialog(QDialog):
     def __init__(self, main_window, parent=None):
+        """
+        Initialize the SAM Deploy Model dialog.
+        """
         super().__init__(parent)
         self.main_window = main_window
         self.annotation_window = main_window.annotation_window
@@ -167,6 +168,9 @@ def __init__(self, main_window, parent=None):
         self.main_layout.addWidget(self.status_bar)
 
     def update_uncertainty_label(self):
+        """
+        Update the uncertainty threshold label when the slider value changes.
+        """
         # Convert the slider value to a ratio (0-1)
         value = self.uncertainty_threshold_slider.value() / 100.0
         self.main_window.update_uncertainty_thresh(value)
@@ -174,12 +178,21 @@ def update_uncertainty_label(self):
         self.conf = self.uncertainty_threshold_slider.value() / 100.0
 
     def on_uncertainty_changed(self, value):
+        """
+        Update the uncertainty threshold slider and label when the shared data changes.
+        
+        Args:
+            value (float): The new value of the uncertainty threshold.
+        """
         # Update the slider and label when the shared data changes
         self.uncertainty_threshold_slider.setValue(int(value * 100))
         self.uncertainty_threshold_label.setText(f"{value:.2f}")
         self.conf = self.uncertainty_threshold_slider.value() / 100.0
 
     def setup_tabs(self):
+        """
+        Set up the tabs for the different models.
+        """
         self.tabs = QTabWidget()
 
         # Create tabs
@@ -195,6 +208,12 @@ def setup_tabs(self):
         self.main_layout.addWidget(self.tabs)
 
     def create_model_tab(self, model_name):
+        """
+        Create a tab for the specified model.
+        
+        Args:
+            model_name (str): The name of the model to create a tab for.
+        """
         tab = QWidget()
         layout = QVBoxLayout(tab)
         combo_box = QComboBox()
@@ -216,6 +235,9 @@ def create_model_tab(self, model_name):
         return tab
 
     def load_model(self):
+        """
+        Load the model selected in the combo box.
+        """
         # Make the cursor busy
         QApplication.setOverrideCursor(Qt.WaitCursor)
         # Show a progress bar
@@ -286,6 +308,12 @@ def load_model(self):
         self.accept()
 
     def resize_image(self, image):
+        """
+        Resize the image to the specified size.
+        
+        Args:
+            image (np.ndarray): The image to resize.
+        """
         # Get the current image size
         imgsz = self.imgsz_spinbox.value()
         # Determine the target shape based on the long side
@@ -300,28 +328,20 @@ def resize_image(self, image):
         return resized_image
 
     def set_image(self, image):
+        """
+        Set the image in the predictor.
+        
+        Args:
+            image (np.ndarray): The image to set.
+        """
         QApplication.setOverrideCursor(Qt.WaitCursor)
         progress_bar = ProgressBar(self.annotation_window, title="Setting Image")
         progress_bar.show()
 
         try:
             if self.predictor is not None:
-                # Ensure image has correct dimensions (h, w, 3)
-                if len(image.shape) == 2:  # Grayscale image
-                    image = np.stack((image,) * 3, axis=-1)
-                elif len(image.shape) == 3:
-                    if image.shape[2] == 4:  # RGBA image
-                        image = image[..., :3]  # Drop alpha channel
-                    elif image.shape[2] != 3:  # If channels are not last
-                        # Check if channels are first (c, h, w)
-                        if image.shape[0] == 3:
-                            image = np.transpose(image, (1, 2, 0))
-                        elif image.shape[0] == 4:  # RGBA in channels-first format
-                            image = np.transpose(image, (1, 2, 0))[..., :3]
-                        else:
-                            raise ValueError("Image must have 3 or 4 color channels")
-                else:
-                    raise ValueError("Image must be 2D or 3D array")
+                # Reshape image if needed
+                image = preprocess_image(image)
 
                 # Save the original image
                 self.original_image = image
@@ -350,9 +370,18 @@ def set_image(self, image):
             QApplication.restoreOverrideCursor()
 
     def predict(self, bbox, points, labels):
+        """
+        Make predictions using the currently loaded model.
+        
+        Args:
+            bbox (list): List of bounding box coordinates in the form [xmin, ymin, xmax, ymax].
+            points (np.ndarray): Array of points in the form [[x1, y1], [x2, y2], ...].
+            labels (list): List of class labels for each point.
+        """
         if not self.predictor:
             QMessageBox.critical(self, "Model Not Loaded", "Model not loaded, cannot make predictions")
             return None
+        
         try:
             point_labels = None
             point_coords = None
@@ -384,7 +413,7 @@ def predict(self, bbox, points, labels):
                 scaled_points[:, :, 1] *= scale_y
                 point_coords = scaled_points.long()  # Cast back to int64
 
-                if self.model_path.startswith("sam_"):
+                if not self.model_path.startswith("sam2"):
                     # Apply the scaled points to the predictor
                     point_coords = self.predictor.transform.apply_coords_torch(point_coords,
                                                                                self.resized_image.shape[:2])
@@ -401,29 +430,28 @@ def predict(self, bbox, points, labels):
                 scaled_bbox[:, 3] *= scale_y
                 bbox_coords = scaled_bbox.long()  # Cast back to int64
 
-                if self.model_path.startswith("sam_"):
+                if not self.model_path.startswith("sam2"):
                     # Apply the scaled boxes to the predictor
                     bbox_coords = self.predictor.transform.apply_boxes_torch(bbox_coords,
                                                                              self.resized_image.shape[:2])
 
             if self.model_path.startswith("sam2"):
-                mask, score, logit = self.predictor.predict(box=bbox_coords,
+                mask, score, _ = self.predictor.predict(box=bbox_coords,
                                                             point_coords=point_coords,
                                                             point_labels=point_labels,
                                                             multimask_output=False)
 
                 mask = torch.tensor(mask).unsqueeze(0)
                 score = torch.tensor(score).unsqueeze(0)
-                logit = torch.tensor(logit).unsqueeze(0)
 
             else:
-                mask, score, logit = self.predictor.predict_torch(boxes=bbox_coords,
+                mask, score, _ = self.predictor.predict_torch(boxes=bbox_coords,
                                                                   point_coords=point_coords,
                                                                   point_labels=point_labels,
                                                                   multimask_output=False)
 
             # Post-process the results
-            results = custom_postprocess(mask, score, logit, self.resized_image, self.original_image)[0]
+            results = to_ultralytics(mask, score, self.original_image)[0]
 
         except Exception as e:
             QMessageBox.critical(self, "Prediction Error", f"Error predicting: {e}")
@@ -432,7 +460,12 @@ def predict(self, bbox, points, labels):
         return results
 
     def boxes_to_masks(self, results_generator):
-
+        """
+        Convert bounding boxes to masks using the currently loaded model.
+        
+        Args:
+            results_generator (generator): Generator of Results objects containing bounding boxes.
+        """
         results_dict = {}
         for results in results_generator:
             path = results.path.replace("\\", "/")
@@ -451,7 +484,7 @@ def boxes_to_masks(self, results_generator):
         for image_path in results_dict:
             # Convert rasterio image to numpy array
             image = self.main_window.image_window.rasterio_open(image_path)
-            image = image.read().transpose(1, 2, 0)
+            image = rasterio_to_numpy(image)
 
             # Set the image
             self.set_image(image)
@@ -476,7 +509,7 @@ def boxes_to_masks(self, results_generator):
             scaled_bbox[:, 3] *= scale_y
             bbox_coords = scaled_bbox.long()  # Cast back to int64
 
-            if self.model_path.startswith("sam_"):
+            if not self.model_path.startswith("sam2"):
                 # Apply the scaled boxes to the predictor
                 bbox_coords = self.predictor.transform.apply_boxes_torch(bbox_coords,
                                                                          self.resized_image.shape[:2])
@@ -489,16 +522,15 @@ def boxes_to_masks(self, results_generator):
 
                 mask = torch.tensor(mask).unsqueeze(0)
                 score = torch.tensor(score).unsqueeze(0)
-                logit = torch.tensor(logit).unsqueeze(0)
 
             else:
-                mask, score, logit = self.predictor.predict_torch(boxes=bbox_coords,
+                mask, score, _ = self.predictor.predict_torch(boxes=bbox_coords,
                                                                   point_coords=None,
                                                                   point_labels=None,
                                                                   multimask_output=False)
 
             # Post-process the results
-            sam_results = custom_postprocess(mask, score, logit, self.resized_image, self.original_image)
+            sam_results = to_ultralytics(mask, score, self.original_image)
             sam_results.boxes = results.boxes
             sam_results.names = results.names
             sam_results.path = image_path
@@ -506,6 +538,9 @@ def boxes_to_masks(self, results_generator):
             yield sam_results
 
     def deactivate_model(self):
+        """
+        Deactivate the currently loaded model.
+        """
         self.loaded_model = None
         self.predictor = None
         self.model_path = None
diff --git a/toolbox/main.py b/toolbox/main.py
index 98f2cc14..42adeca1 100644
--- a/toolbox/main.py
+++ b/toolbox/main.py
@@ -1,11 +1,6 @@
 # TODO Later:
 #   - CoralNet tools
-#   - Merge Datasets
 #   - Resample polygon to be smoother, uniformly spaced
-#   - Don't access all points when updating polygon
-#   - Fix the cuda device string
-#   - Show a loading model window when using AutoDistill (like SAM)
-#   - Look into SAM resizing images bigger than 1024
 
 import traceback
 
diff --git a/toolbox/utilities.py b/toolbox/utilities.py
index 05d4ae4b..222dcb80 100644
--- a/toolbox/utilities.py
+++ b/toolbox/utilities.py
@@ -39,6 +39,42 @@ def get_available_device():
     return devices
 
 
+def preprocess_image(image):
+    """
+    Ensure the image has correct dimensions (h, w, 3).
+
+    :param image:
+    :return:
+    """
+    if len(image.shape) == 2:  # Grayscale image
+        image = np.stack((image,) * 3, axis=-1)
+    elif len(image.shape) == 3:
+        if image.shape[2] == 4:  # RGBA image
+            image = image[..., :3]  # Drop alpha channel
+        elif image.shape[2] != 3:  # If channels are not last
+            # Check if channels are first (c, h, w)
+            if image.shape[0] == 3:
+                image = np.transpose(image, (1, 2, 0))
+            elif image.shape[0] == 4:  # RGBA in channels-first format
+                image = np.transpose(image, (1, 2, 0))[..., :3]
+            else:
+                raise ValueError("Image must have 3 or 4 color channels")
+    else:
+        raise ValueError("Image must be 2D or 3D array")
+
+    return image
+
+
+def rasterio_to_numpy(rasterio_src):
+    """
+    Convert a Rasterio dataset to a NumPy array.
+
+    :param rasterio_src:
+    :return:
+    """
+    return rasterio_src.read().transpose(1, 2, 0)
+
+
 def pixmap_to_numpy(pixmap):
     """
     Convert a QPixmap to a NumPy array.


Test patch:
diff --git a/tests/Subset_Dataset.py b/tests/Subset_Dataset.py
new file mode 100644
index 00000000..e180528a
--- /dev/null
+++ b/tests/Subset_Dataset.py
@@ -0,0 +1,93 @@
+import os
+import shutil
+import argparse
+import random
+import matplotlib.pyplot as plt
+import pandas as pd
+from collections import Counter
+from tqdm import tqdm
+
+
+def count_images_per_class(src_dir):
+    class_counts = {}
+    for root, dirs, files in os.walk(src_dir):
+        if not dirs:  # We are in a class directory
+            class_name = os.path.basename(root)
+            class_counts[class_name] = len(files)
+    return class_counts
+
+
+def subset_dataset(src_dir, dst_dir, min_count):
+    for root, dirs, files in os.walk(src_dir):
+        for dir_name in dirs:
+            class_dir = os.path.join(root, dir_name)
+            subset_class_dir = class_dir.replace(src_dir, dst_dir)
+            os.makedirs(subset_class_dir, exist_ok=True)
+
+        for file_name in tqdm(files, desc=f"Processing {os.path.basename(root)}"):
+            class_name = os.path.basename(root)
+            class_dir = os.path.join(dst_dir, class_name)
+            os.makedirs(class_dir, exist_ok=True)
+
+            if class_name in min_count and len(os.listdir(class_dir)) < min_count[class_name]:
+                src_file = os.path.join(root, file_name)
+                dst_file = os.path.join(class_dir, file_name)
+                shutil.copy(src_file, dst_file)
+
+
+def plot_distribution(dst_dir, dataset_type):
+    class_counts = count_images_per_class(os.path.join(dst_dir, dataset_type))
+    df = pd.DataFrame(list(class_counts.items()), columns=['Class', 'Count'])
+
+    plt.figure(figsize=(10, 6))
+    plt.bar(df['Class'], df['Count'], color='skyblue')
+    plt.xlabel('Class')
+    plt.ylabel('Number of Images')
+    plt.title(f'Distribution of {dataset_type.capitalize()} Set')
+    plt.xticks(rotation=45)
+    plt.tight_layout()
+
+    plot_path = os.path.join(dst_dir, f'{dataset_type}_distribution.png')
+    plt.savefig(plot_path)
+    plt.close()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Subset a YOLO formatted dataset.")
+    parser.add_argument("src", type=str, help="Source directory of the YOLO formatted dataset.")
+    parser.add_argument("dst", type=str, help="Destination directory for the subsetted dataset.")
+    args = parser.parse_args()
+
+    src_dir = args.src
+    dst_dir = os.path.join(os.path.dirname(src_dir), "Datasets_Subsetted", args.dst)
+
+    # Count images per class
+    class_counts = count_images_per_class(src_dir)
+
+    # Find the minority class count greater than 1
+    min_count = min([count for count in class_counts.values() if count > 1])
+
+    # Create a dictionary with the minimum count for each class
+    min_count_dict = {class_name: min_count for class_name in class_counts}
+
+    # Print the amount per class
+    print("Original Class Counts:")
+    for class_name, count in class_counts.items():
+        print(f"Class: {class_name}, Images: {count}")
+
+    print("\nSubsetted Class Counts:")
+    for class_name, count in min_count_dict.items():
+        print(f"Class: {class_name}, Images: {count}")
+
+    # Subset the dataset
+    subset_dataset(src_dir, dst_dir, min_count_dict)
+
+    # Plot and save distribution for train, valid, and test sets
+    for dataset_type in ['train', 'valid', 'test']:
+        plot_distribution(dst_dir, dataset_type)
+
+    print(f"\nSubsetted dataset created at {dst_dir}")
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file