From 2263093c4d3bb1f24568a71a7a35cf4292273a32 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Mon, 18 Dec 2023 21:53:20 -0600 Subject: [PATCH] convert images to cv2 --- api/onnx_web/chain/correct_codeformer.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/api/onnx_web/chain/correct_codeformer.py b/api/onnx_web/chain/correct_codeformer.py index b23a2fc0..5a99fec6 100644 --- a/api/onnx_web/chain/correct_codeformer.py +++ b/api/onnx_web/chain/correct_codeformer.py @@ -2,6 +2,7 @@ from logging import getLogger from os import path from typing import Optional +import cv2 import torch from PIL import Image from torchvision.transforms.functional import normalize @@ -67,7 +68,8 @@ class CorrectCodeformerStage(BaseStage): ) results = [] - for img in sources.as_image(): + for img in sources.as_numpy(): + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # clean all the intermediate results to process the next image face_helper.clean_all() face_helper.read_image(img) @@ -115,8 +117,7 @@ class CorrectCodeformerStage(BaseStage): face_helper.get_inverse_affine(None) # paste each restored face to the input image - results.append( - face_helper.paste_faces_to_input_image(upsample_img=img, draw_box=False) - ) + output = face_helper.paste_faces_to_input_image(upsample_img=img, draw_box=False) + results.append(Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))) return StageResult.from_images(results)