Merge pull request #10 from hdr-bgnn/9-fix-model-not-found

Fix model not found error
hdr-bgnn · Jan 18, 2023 · 8abd71d · 8abd71d
2 parents 699e400 + 18a912c
commit 8abd71d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
diff --git a/Dockerfile b/Dockerfile
@@ -36,4 +36,7 @@ COPY --from=model_fetcher /model/model_final.pth \
 
 COPY gen_metadata.py /pipeline
 
+# Default to use enhanced model added above (unset DM_CONFIG_FILENAME to use config.json)
+ENV DM_CONFIG_FILENAME config_enhance_no_joel.json
+
 CMD echo "python gen_metadata.py"
diff --git a/config/config_enhance_no_joel.json b/config/config_enhance_no_joel.json
@@ -0,0 +1,4 @@
+{
+  "ENHANCE": 1,
+  "JOEL": 0
+}
diff --git a/gen_metadata.py b/gen_metadata.py
@@ -27,7 +27,8 @@
 
 # Look for the config directory in the same directory as this script
 root_dir_path = os.path.join(os.path.dirname(__file__))
-main_config_path = os.path.join(root_dir_path, 'config', 'config.json')
+config_filename = os.environ.get('DM_CONFIG_FILENAME', 'config.json')
+main_config_path = os.path.join(root_dir_path, 'config', config_filename)
 mask_config_path = os.path.join(root_dir_path, 'config', 'mask_rcnn_R_50_FPN_3x.yaml')
 
 VAL_SCALE_FAC = 0.5