1
0
Fork 0

convert images to cv2

This commit is contained in:
Sean Sube 2023-12-18 21:53:20 -06:00
parent 1163accb49
commit 2263093c4d
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 5 additions and 4 deletions

View File

@ -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)