@@ -47,33 +47,37 @@ def superanimal_analyze_images(
4747 images : str | Path | list [str ] | list [Path ],
4848 max_individuals : int ,
4949 out_folder : str ,
50+ bbox_threshold : float = 0.6 ,
5051 progress_bar : bool = True ,
5152 device : str | None = None ,
5253 customized_pose_checkpoint : str | None = None ,
5354 customized_detector_checkpoint : str | None = None ,
5455 customized_model_config : str | None = None ,
5556):
5657 """
57- This funciton inferences a superanimal model on a set of images and saves the results as labeled images.
58+ This funciton inferences a superanimal model on a set of images and saves the
59+ results as labeled images.
5860
5961 Parameters
6062 ----------
6163 superanimal_name: str
62- The name of the superanimal to analyze.
63- supported list:
64- superanimal_topviewmouse
65- superanimal_quadruped
64+ The name of the superanimal to analyze. Supported list:
65+ - "superanimal_topviewmouse"
66+ - "superanimal_quadruped"
6667 model_name: str
67- The name of the model to use for inference.
68- supported list:
69- hrnetw32
68+ The name of the model to use for inference. Supported list:
69+ - "hrnetw32"
7070 images: str | Path | list[str] | list[Path]
7171 The images to analyze. Can either be a directory containing images, or
7272 a list of paths of images.
7373 max_individuals: int
7474 The maximum number of individuals to detect in each image.
7575 out_folder: str
7676 The directory where the labeled images will be saved.
77+ bbox_threshold: float, default=0.1
78+ The minimum confidence score to keep bounding box detections. Must be in (0, 1).
79+ Only used when `customized_model_config=None` (otherwise, edit your
80+ `customized_model_config` with the desired bbox_threshold).
7781 progress_bar: bool
7882 Whether to display a progress bar when running inference.
7983 device: str | None
@@ -95,17 +99,19 @@ def superanimal_analyze_images(
9599 --------
96100 >>> import deeplabcut
97101 >>> from deeplabcut.pose_estimation_pytorch.apis.analyze_images import superanimal_analyze_images
98- >>> superanimal_name = ' superanimal_quadruped'
99- >>> model_name = ' hrnetw32'
100- >>> device = ' cuda'
102+ >>> superanimal_name = " superanimal_quadruped"
103+ >>> model_name = " hrnetw32"
104+ >>> device = " cuda"
101105 >>> max_individuals = 3
102- >>> test_images_folder = 'test_rodent_images'
103- >>> out_images_folder = 'vis_test_rodent_images'
104- >>> ret = superanimal_analyze_images(superanimal_name,
105- model_name,
106- test_images_folder,
107- max_individuals,
108- out_images_folder)
106+ >>> test_images_folder = "test_rodent_images"
107+ >>> out_images_folder = "vis_test_rodent_images"
108+ >>> ret = superanimal_analyze_images(
109+ >>> superanimal_name,
110+ >>> model_name,
111+ >>> test_images_folder,
112+ >>> max_individuals,
113+ >>> out_images_folder
114+ >>> )
109115 """
110116
111117 os .makedirs (out_folder , exist_ok = True )
@@ -119,6 +125,10 @@ def superanimal_analyze_images(
119125 snapshot_path ,
120126 detector_path ,
121127 ) = get_config_model_paths (superanimal_name , model_name )
128+
129+ if "detector" in model_cfg :
130+ model_cfg ["detector" ]["model" ]["box_score_thresh" ] = bbox_threshold
131+
122132 config = {** project_config , ** model_cfg }
123133 config = update_config (config , max_individuals , device )
124134 else :
@@ -146,9 +156,7 @@ def superanimal_analyze_images(
146156
147157 superanimal_colormaps = get_superanimal_colormaps ()
148158 colormap = superanimal_colormaps [superanimal_name ]
149-
150159 create_labeled_images_from_predictions (predictions , out_folder , colormap )
151-
152160 return predictions
153161
154162
@@ -164,8 +172,6 @@ def analyze_images(
164172 device : str | None = None ,
165173 max_individuals : int | None = None ,
166174 progress_bar : bool = True ,
167- superanimal_name = None ,
168- model_name = None ,
169175) -> dict [str , dict ]:
170176 """Runs analysis on images using a pose model.
171177
0 commit comments