feat(api): add strength param to denoise stage

ssube · Nov 13, 2023 · 768f478 · 768f478
1 parent 21d1240
commit 768f478
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/api/onnx_web/chain/blend_denoise.py b/api/onnx_web/chain/blend_denoise.py
@@ -24,6 +24,7 @@ def run(
         _params: ImageParams,
         sources: List[Image.Image],
         *,
+        strength: int = 10,
         stage_source: Optional[Image.Image] = None,
         callback: Optional[ProgressCallback] = None,
         **kwargs,
@@ -33,7 +34,7 @@ def run(
         results = []
         for source in sources:
             data = cv2.cvtColor(np.array(source), cv2.COLOR_RGB2BGR)
-            data = cv2.fastNlMeansDenoisingColored(data)
+            data = cv2.fastNlMeansDenoisingColored(data, None, strength, strength)
             results.append(Image.fromarray(cv2.cvtColor(data, cv2.COLOR_BGR2RGB)))
 
         return results