fix docs

Author: kozlov721

luxonis_train/__main__.py

@@ -490,6 +490,15 @@ def quantize(
     config: str | None = None,
     weights: str | None = None,
 ):
+    """Quantize the model using AIMET.
+
+    @type config: str
+    @param config: Path to the configuration file.
+    @type weights: str
+    @param weights: Path to the model weights.
+    @type opts: list[str]
+    @param opts: A list of optional CLI overrides of the config file.
+    """
     model = create_model(
         config, opts, weights=weights, allow_empty_dataset=True
     )

luxonis_train/callbacks/README.md

@@ -188,7 +188,15 @@ A callback that maintains an exponential moving average (EMA) of the model's par
 
 ## `AIMETCallback`
 
-# Callback to perform AIMET quantization at the end of the training.
+Callback to perform AIMET quantization at the end of the training.
+
+This callback runs AIMET post-training static quantization using the best checkpoint (by default based on the main metric) at the end of training.
+
+**Parameters:**
+
+| Key                    | Type                        | Default value | Description                                                                                                                                                     |
+| ---------------------- | --------------------------- | ------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `preferred_checkpoint` | `Literal["metric", "loss"]` | `"metric"`    | Which checkpoint should the callback use. If the preferred checkpoint is not available, the other option is used. If none is available, the callback is skipped |
 
 ## `TrainingProgressCallback`
 

luxonis_train/callbacks/luxonis_progress_bar.py

@@ -70,8 +70,8 @@ def print_table(
 
         @type title: str
         @param title: Title of the table
-        @type table: Mapping[str, int | str | float]
-        @param table: Table to print
+        @type table: Iterable[tuple[str | int | float, ...]]
+        @param table: Table to print as an iterable of rows, where each row is a tuple of values.
         @type column_names: list[str]
         @param column_names: Names of the columns in the table
         """
@@ -181,13 +181,10 @@ def print_table(
 
         @type title: str
         @param title: Title of the table
-        @type table: Mapping[str, int | str | float]
-        @param table: Table to print
-        @type key_name: str
-        @param key_name: Name of the key column. Defaults to C{"Name"}.
-        @type value_name: str
-        @param value_name: Name of the value column. Defaults to
-            C{"Value"}.
+        @type table: Iterable[tuple[str | int | float, ...]]
+        @param table: Table to print as an iterable of rows, where each row is a tuple of values.
+        @type column_names: list[str]
+        @param column_names: Names of the columns in the table
         """
         self._rule(title)
         formatted = tabulate(
@@ -332,13 +329,10 @@ def print_table(
 
         @type title: str
         @param title: Title of the table
-        @type table: Mapping[str, int | str | float]
-        @param table: Table to print
-        @type key_name: str
-        @param key_name: Name of the key column. Defaults to C{"Name"}.
-        @type value_name: str
-        @param value_name: Name of the value column. Defaults to
-            C{"Value"}.
+        @type table: Iterable[tuple[str | int | float, ...]]
+        @param table: Table to print as an iterable of rows, where each row is a tuple of values.
+        @type column_names: list[str]
+        @param column_names: Names of the columns in the table
         @param console: Console instance to use, if None use default
             console. Defaults to None.
         @type console: Console | None

luxonis_train/utils/dataset_metadata.py

@@ -122,9 +122,8 @@ def n_keypoints(self, task_name: str | None = None) -> int:
         @type task_name: str | None
         @param task_name: Task to get the number of keypoints for.
         @rtype: int
-        @return: Number of keypoints for the specified task type.
-        @raises ValueError: If the C{task} is not present in the
-            dataset.
+        @return: Number of keypoints for the specified task type or 0 if
+            the task does not involve keypoints.
         @raises RuntimeError: If the C{task} was not provided and the
             dataset contains different number of keypoints for different
             task types.