Add variables for default config path (#244)

SainsburyWellcomeCentre · Nov 13, 2024 · c3e653d · c3e653d
1 parent e371bf9
commit c3e653d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/crabs/detector/train_model.py b/crabs/detector/train_model.py
@@ -29,6 +29,10 @@
     log_data_augm_as_artifacts,
 )
 
+DEFAULT_DETECTOR_CONFIG = str(
+    Path(__file__).parent / "config" / "faster_rcnn.yaml"
+)
+
 
 class DetectorTrain:
     """Training class for detector algorithm.
@@ -300,7 +304,7 @@ def train_parse_args(args):
     parser.add_argument(
         "--config_file",
         type=str,
-        default=str(Path(__file__).parent / "config" / "faster_rcnn.yaml"),
+        default=DEFAULT_DETECTOR_CONFIG,
         help=(
             "Location of YAML config to control training. "
             "Default: "

diff --git a/crabs/tracker/track_video.py b/crabs/tracker/track_video.py
@@ -30,6 +30,10 @@
     format_and_filter_bbox_predictions_for_sort,
 )
 
+DEFAULT_TRACKING_CONFIG = str(
+    Path(__file__).parent / "config" / "tracking_config.yaml"
+)
+
 
 class Tracking:
     """Interface for detecting and tracking crabs on a video.
@@ -349,7 +353,7 @@ def tracking_parse_args(args):
     parser.add_argument(
         "--config_file",
         type=str,
-        default=str(Path(__file__).parent / "config" / "tracking_config.yaml"),
+        default=DEFAULT_TRACKING_CONFIG,
         help=(
             "Location of YAML config to control tracking. "
             "Default: "