fix(api): convert images to ndarray before passing to GFPGAN

ssube · Jan 31, 2023 · 39f84c4 · 39f84c4
1 parent d5e5588
commit 39f84c4
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/api/onnx_web/chain/correct_gfpgan.py b/api/onnx_web/chain/correct_gfpgan.py
@@ -17,6 +17,8 @@
     load_resrgan,
 )
 
+import numpy as np
+
 logger = getLogger(__name__)
 
 
@@ -56,7 +58,7 @@ def correct_gfpgan(
     ctx: ServerContext,
     _stage: StageParams,
     _params: ImageParams,
-    image: Image.Image,
+    source_image: Image.Image,
     *,
     upscale: UpscaleParams,
     upsampler: Optional[RealESRGANer] = None,
@@ -69,7 +71,9 @@ def correct_gfpgan(
     logger.info('correcting faces with GFPGAN model: %s', upscale.correction_model)
     gfpgan = load_gfpgan(ctx, upscale, upsampler=upsampler)
 
+    output = np.array(source_image)
     _, _, output = gfpgan.enhance(
         image, has_aligned=False, only_center_face=False, paste_back=True, weight=upscale.face_strength)
+    output = Image.fromarray(output, 'RGB')
 
     return output