Update src/util.py

src/util.py

@@ -1,7 +1,9 @@
 import concurrent.futures
 import io
 import os
+import time
 
+import cv2
 import oss2
 import requests
 from PIL import Image
@@ -18,6 +20,19 @@ bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, buck
 oss_path = os.getenv("OSS_PATH")
 
 
+def resize(img, short_side_length=512):
+    height, width, _ = img.shape
+    aspect_ratio = width / height
+    if width > height:
+        new_width = short_side_length
+        new_height = int(new_width / aspect_ratio)
+    else:
+        new_height = short_side_length
+        new_width = int(new_height * aspect_ratio)
+    resized_img = cv2.resize(img, (new_width, new_height))
+    return resized_img
+
+
 def download_img_pil(index, img_url):
     r = requests.get(img_url, stream=True)
     if r.status_code == 200:
@@ -43,16 +58,20 @@ def download_images(img_urls, batch_size):
     return imgs_pil
 
 
-def upload_np_2_oss(input_image, name="cache.png"):
+def upload_np_2_oss(input_image, name="cache.jpg"):
     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")
+        Image.fromarray(resize(input_image)).save(imgByteArr, format="PNG")
     else:
-        Image.fromarray(input_image).save(imgByteArr, format="JPEG", quality=95)
+        Image.fromarray(resize(input_image)).save(imgByteArr, format="JPEG", quality=95)
     imgByteArr = imgByteArr.getvalue()
 
+    start_time = time.perf_counter()
     bucket.put_object(oss_path + "/" + name, imgByteArr)
     ret = bucket.sign_url('GET', oss_path + "/" + name, 60 * 60 * 24)
+    logger.info(f"upload cost: {time.perf_counter() - start_time} s.")
     del imgByteArr
     return ret