diff --git a/annotator/teed/__init__.py b/annotator/teed/__init__.py index 9aed08a..4751704 100644 --- a/annotator/teed/__init__.py +++ b/annotator/teed/__init__.py @@ -36,7 +36,7 @@ class TEEDDector: if self.model is not None: self.model.cpu() - def __call__(self, image: np.ndarray, is_safe: bool = True) -> np.ndarray: + def __call__(self, image: np.ndarray, safe_steps: int = 2) -> np.ndarray: self.model.to(self.device) @@ -49,7 +49,7 @@ class TEEDDector: edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges] edges = np.stack(edges, axis=2) edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64))) - if is_safe: - edge = safe_step(edge) + if safe_steps != 0: + edge = safe_step(edge, safe_steps) edge = (edge * 255.0).clip(0, 255).astype(np.uint8) return edge \ No newline at end of file diff --git a/scripts/global_state.py b/scripts/global_state.py index e4522ca..c2f47c5 100644 --- a/scripts/global_state.py +++ b/scripts/global_state.py @@ -111,8 +111,7 @@ cn_preprocessor_modules = { "anime_face_segment": anime_face_segment, "densepose": functools.partial(densepose, cmap="viridis"), "densepose_parula": functools.partial(densepose, cmap="parula"), - "te_hed":functools.partial(te_hed, is_safe=False), - "te_hed_safe":functools.partial(te_hed, is_safe=True), + "te_hed":te_hed, } cn_preprocessor_unloadable = { @@ -176,7 +175,6 @@ preprocessor_aliases = { "densepose": "densepose (pruple bg & purple torso)", "densepose_parula": "densepose_parula (black bg & blue torso)", "te_hed": "softedge_teed", - "te_hed_safe": "softedge_teed_safe", } ui_preprocessor_keys = ['none', preprocessor_aliases['invert']] diff --git a/scripts/processor.py b/scripts/processor.py index f624c1d..eac9777 100644 --- a/scripts/processor.py +++ b/scripts/processor.py @@ -692,13 +692,13 @@ def unload_densepose(): model_te_hed = None -def te_hed(img, res=512, is_safe=True,**kwargs): +def te_hed(img, res=512, thr_a=2, **kwargs): img, remove_pad = resize_image_with_pad(img, res) global model_te_hed if model_te_hed is None: from annotator.teed import TEEDDector model_te_hed = TEEDDector() - result = model_te_hed(img, is_safe=is_safe) + result = model_te_hed(img, safe_steps=int(thr_a)) return remove_pad(result), True def unload_te_hed(): @@ -1200,15 +1200,14 @@ preprocessor_sliders_config = { "value": 512, "min": 64, "max": 2048 - } - ], - "te_hed_safe": [ + }, { - "name": flag_preprocessor_resolution, - "value": 512, - "min": 64, - "max": 2048 - } + "name": "Safe Steps", + "min": 0, + "max": 10, + "value": 2, + "step": 1, + }, ], }