1
0
Fork 0

feat(api): add strength param to denoise stage

This commit is contained in:
Sean Sube 2023-11-12 21:30:01 -06:00
parent 21d1240e28
commit 768f478884
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 2 additions and 1 deletions

View File

@ -24,6 +24,7 @@ class BlendDenoiseStage(BaseStage):
_params: ImageParams,
sources: List[Image.Image],
*,
strength: int = 10,
stage_source: Optional[Image.Image] = None,
callback: Optional[ProgressCallback] = None,
**kwargs,
@ -33,7 +34,7 @@ class BlendDenoiseStage(BaseStage):
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