bug fix: resize alpha channel when using non original image size

2022-04-09 12:31:48 +08:00 · 2022-04-09 12:31:48 +08:00 · ca65c95d22
commit ca65c95d22
parent 1b1aade067
1 changed files with 5 additions and 0 deletions
--- a/main.py
+++ b/main.py
@ -92,7 +92,12 @@ def process():
    print(f"process time: {(time.time() - start) * 1000}ms")

    torch.cuda.empty_cache()
+
    if alpha_channel is not None:
+        if alpha_channel.shape[:2] != res_np_img.shape[:2]:
+            alpha_channel = cv2.resize(
+                alpha_channel, dsize=(res_np_img.shape[1], res_np_img.shape[0])
+            )
        res_np_img = np.concatenate(
            (res_np_img, alpha_channel[:, :, np.newaxis]), axis=-1
        )