Update utils.py

utils.py

@@ -1,6 +1,7 @@
 import requests
 import numpy as np
-from models.deep_colorization.colorizers import postprocess_tens, preprocess_img
+from models.deep_colorization.colorizers import postprocess_tens, preprocess_img, load_img
+from PIL import Image
 
 
 # Define a function that we can use to load lottie files from a link.
@@ -36,3 +37,16 @@ def format_time(seconds: float) -> str:
 def colorize_frame(frame, colorizer) -> np.ndarray:
     tens_l_orig, tens_l_rs = preprocess_img(frame, HW=(256, 256))
     return postprocess_tens(tens_l_orig, colorizer(tens_l_rs).cpu())
+
+
+def colorize_image(file, loaded_model):
+    img = load_img(file)
+    # If user input a colored image with 4 channels, discard the fourth channel
+    if img.shape[2] == 4:
+        img = img[:, :, :3]
+
+    tens_l_orig, tens_l_rs = preprocess_img(img, HW=(256, 256))
+    out_img = postprocess_tens(tens_l_orig, loaded_model(tens_l_rs).cpu())
+    new_img = Image.fromarray((out_img * 255).astype(np.uint8))
+
+    return out_img, new_img