| from PIL import Image | |
| import numpy as np | |
| import cv2 | |
| def postprocess_img(original_img: Image.Image, colored_img: Image.Image) -> Image.Image: | |
| original_np, colored_np = np.array(original_img), np.array(colored_img) | |
| original_yuv = cv2.cvtColor(original_np, cv2.COLOR_BGR2YUV) | |
| predicted_yuv = cv2.cvtColor(colored_np, cv2.COLOR_BGR2YUV) | |
| processed_img = original_yuv.copy() | |
| processed_img[:, :, 1:] = predicted_yuv[:, :, 1:] | |
| processed_img = cv2.cvtColor(processed_img, cv2.COLOR_YUV2BGR) | |
| processed_img = Image.fromarray(processed_img) | |
| return processed_img |