Update src/util.py

src/util.py

@@ -1,6 +1,7 @@
 import concurrent.futures
 import io
 import os
+import time
 
 import numpy as np
 import oss2
@@ -20,6 +21,15 @@ oss_path = "hejunjie.hjj/TransferAnythingHF"
 oss_path_img_gallery = "hejunjie.hjj/TransferAnythingHF_img_gallery"
 
 
+def resize(image, short_side_length=768):
+    width, height = image.size
+    ratio = short_side_length / min(width, height)
+    new_width = int(width * ratio)
+    new_height = int(height * ratio)
+    resized_image = image.resize((new_width, new_height))
+    return resized_image
+
+    
 def download_img_pil(index, img_url):
     # print(img_url)
     r = requests.get(img_url, stream=True)
@@ -50,11 +60,13 @@ def download_images(img_urls, batch_size):
 
 def upload_np_2_oss(input_image, name="cache.png", gallery=False):
     assert name.lower().endswith((".png", ".jpg")), name
+    if name.lower().endswith(".png"):
+        name = name[:-4] + ".jpg"
     imgByteArr = io.BytesIO()
     if name.lower().endswith(".png"):
-        Image.fromarray(input_image).save(imgByteArr, format="PNG")
+        resize(Image.fromarray(input_image)).save(imgByteArr, format="PNG")
     else:
-        Image.fromarray(input_image).save(imgByteArr, format="JPEG", quality=95)
+        resize(Image.fromarray(input_image)).save(imgByteArr, format="JPEG", quality=95)
     imgByteArr = imgByteArr.getvalue()
 
     if gallery:
@@ -62,8 +74,10 @@ def upload_np_2_oss(input_image, name="cache.png", gallery=False):
     else:
         path = oss_path
 
+    start_time = time.perf_counter()
     bucket.put_object(path + "/" + name, imgByteArr)  # data为数据，可以是图片
     ret = bucket.sign_url('GET', path + "/" + name, 60 * 60 * 24)  # 返回值为链接，参数依次为，方法/oss上文件路径/过期时间(s)
+    logger.info(f"upload cost: {time.perf_counter() - start_time} s.")
     del imgByteArr
     return ret