Updated README

configs/classification_heavy_model.yaml

@@ -14,9 +14,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
 
   batch_size: 8
   epochs: 200

configs/classification_light_model.yaml

@@ -14,9 +14,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
 
   batch_size: 8
   epochs: 200

configs/complex_model.yaml

@@ -15,30 +15,30 @@ model:
         - RepPANNeck
 
       losses:
-        name: EfficientKeypointBboxLoss
+        - name: EfficientKeypointBBoxLoss
 
       metrics:
         - name: ObjectKeypointSimilarity
           is_main_metric: true
         - name: MeanAveragePrecisionKeypoints
 
       visualizers:
-        name: MultiVisualizer
-        params:
-          visualizers:
-            - name: KeypointVisualizer
-              params:
-                nonvisible_color: blue
-            - name: BBoxVisualizer
-              params:
-                colors:
-                  person: "#FF5055"
+        - name: MultiVisualizer
+          params:
+            visualizers:
+              - name: KeypointVisualizer
+                params:
+                  nonvisible_color: blue
+              - name: BBoxVisualizer
+                params:
+                  colors:
+                    person: "#FF5055"
 
     - name: SegmentationHead
       inputs:
         - RepPANNeck
       losses:
-        name: BCEWithLogitsLoss
+        - name: BCEWithLogitsLoss
       metrics:
         - name: F1Score
           params:
@@ -47,9 +47,9 @@ model:
           params:
             task: binary
       visualizers:
-        name: SegmentationVisualizer
-        params:
-          colors: "#FF5055"
+        - name: SegmentationVisualizer
+          params:
+            colors: "#FF5055"
 
     - name: EfficientBBoxHead
       inputs:
@@ -58,18 +58,18 @@ model:
         conf_thres: 0.75
         iou_thres: 0.45
       losses:
-        name: AdaptiveDetectionLoss
+        - name: AdaptiveDetectionLoss
       metrics:
-        name: MeanAveragePrecision
+        - name: MeanAveragePrecision
       visualizers:
-        name: BBoxVisualizer
+        - name: BBoxVisualizer
 
 tracker:
   project_name: coco_test
   save_directory: output
-  is_tensorboard: True
-  is_wandb: False
-  is_mlflow: False
+  is_tensorboard: true
+  is_wandb: false
+  is_mlflow: false
 
 loader:
   train_view: train
@@ -86,23 +86,23 @@ trainer:
 
   n_sanity_val_steps: 1
   profiler: null
-  verbose: True
+  verbose: true
   batch_size: 8
   accumulate_grad_batches: 1
   epochs: &epochs 200
   n_workers: 8
   validation_interval: 10
   n_log_images: 8
-  skip_last_batch: True
-  log_sub_losses: True
+  skip_last_batch: true
+  log_sub_losses: true
   save_top_k: 3
 
   preprocessing:
     train_image_size: [&height 384, &width 384]
-    keep_aspect_ratio: True
-    train_rgb: True
+    keep_aspect_ratio: true
+    train_rgb: true
     normalize:
-      active: True
+      active: true
     augmentations:
       - name: Defocus
         params:
@@ -131,14 +131,15 @@ trainer:
         mode: min
         verbose: true
     - name: ExportOnTrainEnd
+    - name: ArchiveOnTrainEnd
     - name: TestOnTrainEnd
 
   optimizer:
     name: SGD
     params:
       lr: 0.02
       momentum: 0.937
-      nesterov: True
+      nesterov: true
       weight_decay: 0.0005
 
   scheduler:

configs/detection_heavy_model.yaml

@@ -14,9 +14,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
 
   batch_size: 8
   epochs: &epochs 200

configs/detection_light_model.yaml

@@ -14,9 +14,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
 
   batch_size: 8
   epochs: &epochs 200

configs/example_export.yaml

@@ -15,9 +15,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
 
   batch_size: 8
   epochs: &epochs 200
@@ -46,5 +46,5 @@ exporter:
   onnx:
     opset_version: 11
   blobconverter:
-    active: True
+    active: true
     shaves: 8

configs/example_tuning.yaml

@@ -2,9 +2,9 @@
 
 
 model:
-  name: segmentation_light
+  name: detection_light
   predefined_model:
-    name: SegmentationModel
+    name: DetectionModel
     params:
       variant: light
 
@@ -15,9 +15,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
     augmentations:
       - name: Defocus
         params:
@@ -40,7 +40,7 @@ trainer:
 
 
 tuner:
-  study_name: seg_study
+  study_name: det_study
   n_trials: 10
   storage:
     storage_type: local

configs/keypoint_bbox_heavy_model.yaml

@@ -14,9 +14,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
 
   batch_size: 8
   epochs: &epochs 200

configs/keypoint_bbox_light_model.yaml

@@ -14,9 +14,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
 
   batch_size: 8
   epochs: &epochs 200

configs/segmentation_heavy_model.yaml

@@ -14,9 +14,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
 
   batch_size: 8
   epochs: &epochs 200

configs/segmentation_light_model.yaml

@@ -14,9 +14,9 @@ loader:
 trainer:
   preprocessing:
     train_image_size: [384, 512]
-    keep_aspect_ratio: True
+    keep_aspect_ratio: true
     normalize:
-      active: True
+      active: true
 
   batch_size: 8
   epochs: &epochs 200

luxonis_train/__main__.py

@@ -41,6 +41,15 @@ class _ViewType(str, Enum):
     ),
 ]
 
+WeightsType = Annotated[
+    Path | None,
+    typer.Option(
+        help="Path to the model weights.",
+        show_default=False,
+        metavar="FILE",
+    ),
+]
+
 ViewType = Annotated[
     _ViewType, typer.Option(help="Which dataset view to use.")
 ]
@@ -77,12 +86,13 @@ def train(
 def test(
     config: ConfigType = None,
     view: ViewType = _ViewType.VAL,
+    weights: WeightsType = None,
     opts: OptsType = None,
 ):
     """Evaluate model."""
     from luxonis_train.core import LuxonisModel
 
-    LuxonisModel(config, opts).test(view=view.value)
+    LuxonisModel(config, opts).test(view=view.value, weights=weights)
 
 
 @app.command()
@@ -94,11 +104,21 @@ def tune(config: ConfigType = None, opts: OptsType = None):
 
 
 @app.command()
-def export(config: ConfigType = None, opts: OptsType = None):
+def export(
+    config: ConfigType = None,
+    save_path: Annotated[
+        Path | None,
+        typer.Option(help="Path where to save the exported model."),
+    ] = None,
+    weights: WeightsType = None,
+    opts: OptsType = None,
+):
     """Export model."""
     from luxonis_train.core import LuxonisModel
 
-    LuxonisModel(config, opts).export()
+    LuxonisModel(config, opts).export(
+        onnx_save_path=save_path, weights=weights
+    )
 
 
 @app.command()
@@ -107,13 +127,17 @@ def infer(
     view: ViewType = _ViewType.VAL,
     save_dir: SaveDirType = None,
     source_path: SourcePathType = None,
+    weights: WeightsType = None,
     opts: OptsType = None,
 ):
     """Run inference."""
     from luxonis_train.core import LuxonisModel
 
     LuxonisModel(config, opts).infer(
-        view=view.value, save_dir=save_dir, source_path=source_path
+        view=view.value,
+        save_dir=save_dir,
+        source_path=source_path,
+        weights=weights,
     )
 
 
@@ -138,7 +162,8 @@ def inspect(
             "-s",
             help=(
                 "Multiplier for the image size. "
-                "By default the images are shown in their original size."
+                "By default the images are shown in their original size. "
+                "Use this option to scale them."
             ),
             show_default=False,
         ),
@@ -223,19 +248,20 @@ def inspect(
 
 @app.command()
 def archive(
+    config: ConfigType = None,
     executable: Annotated[
         str | None,
         typer.Option(
             help="Path to the model file.", show_default=False, metavar="FILE"
         ),
     ] = None,
-    config: ConfigType = None,
+    weights: WeightsType = None,
     opts: OptsType = None,
 ):
     """Generate NN archive."""
     from luxonis_train.core import LuxonisModel
 
-    LuxonisModel(str(config), opts).archive(executable)
+    LuxonisModel(str(config), opts).archive(path=executable, weights=weights)
 
 
 def version_callback(value: bool):