From d79606aa7234c7cca8d8842a34d6fe9351ee4eff Mon Sep 17 00:00:00 2001
From: ashawkey <ashawkey1999@gmail.com>
Date: Thu, 28 Mar 2024 12:54:02 +0800
Subject: [PATCH] write_image: fix a silent bug when passing in batched images

---
 kiui/utils.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/kiui/utils.py b/kiui/utils.py
index 575b6a1..27eeb30 100644
--- a/kiui/utils.py
+++ b/kiui/utils.py
@@ -250,9 +250,14 @@ def write_image(
 
     if img.dtype == np.float32 or img.dtype == np.float64:
         img = (img * 255).astype(np.uint8)
-
-    # cvtColor
+    
+    if len(img.shape) == 4:
+        if img.shape[0] > 1:
+            raise ValueError(f'only support saving a single image! current image: {img.shape}')
+        img = img[0]
+        
     if len(img.shape) == 3:
+        # cvtColor
         if order == "RGB":
             if img.shape[-1] == 4:
                 img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)