diff --git a/README.md b/README.md index ce9baca..bc366cb 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,6 @@ conda activate BranchPoseNet Then install the required libraries ``` -cd whorl_pose_detector pip install -r requirements.txt ``` @@ -31,7 +30,7 @@ For predicting on your own data just point the --dir_root to the path to the fol ### 🖥️ To run using command line (CLI): ``` -python whorl_pose_detect_CLI.py --dir_root data --my_model whorl_pose_nano_1000px/weights/best.pt --alpha 0.5 --min_internodal_d 0.3 --tree_id_label treeID --semantic_label semantic +python whorl_pose_detect.py --dir_root data --my_model models/whorl_pose_nano_1000px/weights/best.pt --alpha 0.5 --min_internodal_d 0.3 --tree_id_label treeID --semantic_label semantic ``` ### 🎮 Demo version diff --git a/demo_predict_whorl_pose.ipynb b/demo_predict_whorl_pose.ipynb index f84b691..3f89f19 100644 --- a/demo_predict_whorl_pose.ipynb +++ b/demo_predict_whorl_pose.ipynb @@ -149,7 +149,7 @@ "min_internodal_d= 0.3 # in m (it can be set to something like 0.01 to allow the full output to post-process later)\n", "\n", "# pose model\n", - "my_model=\"whorl_pose_nano_1000px/weights/best.pt\"\n", + "my_model=\"models/whorl_pose_nano_1000px/weights/best.pt\"\n", "\n", "# label for the column with tree instance unique identifiers\n", "tree_id_label='treeID' \n", diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..ac16299 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +pandas +numpy +laspy[lazrs] +Pillow +scipy +matplotlib +ultralytics diff --git a/whorl_pose_detect.py b/whorl_pose_detect.py index 5f45847..d152327 100644 --- a/whorl_pose_detect.py +++ b/whorl_pose_detect.py @@ -12,7 +12,7 @@ import matplotlib matplotlib.use('Agg') # Use Agg backend for non-GUI rendering import matplotlib.pyplot as plt - +import argparse import ultralytics from ultralytics import YOLO @@ -24,23 +24,37 @@ #################################################################################################################################### +# Parse CLI arguments first +parser = argparse.ArgumentParser() +parser.add_argument("--dir_root", default="data", help="Path to root folder where .las/.laz files are stored") +parser.add_argument("--my_model", default="models/whorl_pose_nano_1000px/weights/best.pt", help="Path to YOLO model weights") +parser.add_argument("--alpha", type=float, default=0.5, help="Alpha value for plotting") +parser.add_argument("--min_internodal_d", type=float, default=0.3, help="Minimum internodal distance in meters") +parser.add_argument("--tree_id_label", default="treeID", help="Column name for tree IDs") +parser.add_argument("--semantic_label", default="semantic", help="Column name for semantic labels") +args = parser.parse_args() + +# --------------------------------------------------------- +# Default parameters (can be overridden by CLI arguments) +# --------------------------------------------------------- + # directory to where YOLO temp predictions are stored -dir_root="data" # path to root folder where .las/.laz files are stored +dir_root = args.dir_root # path to root folder where .las/.laz files are stored # alpha for plotting -alpha= 0.5 +alpha = args.alpha # minimum distance between whorls (used to remove duplicates) -min_internodal_d= 0.3 # in m (it can be set to something like 0.01 to allow the full output to post-process later) +min_internodal_d = args.min_internodal_d # in m (can be set to 0.01 to allow full output post-processing later) # pose model -my_model="whorl_pose_nano_1000px/weights/best.pt" +my_model = args.my_model # label for the column with tree instance unique identifiers -tree_id_label='treeID' +tree_id_label = args.tree_id_label -# label for the column with semantic labels (non required) -semantic_label='semantic' +# label for the column with semantic labels (non-required) +semantic_label = args.semantic_label #################################################################################################################################### # files_to_process= list_las_laz_files(dir_root) @@ -145,3 +159,6 @@ # Cleanup temp files import shutil shutil.rmtree(dir_pred) + + +