Update src/util.py

src/util.py

@@ -3,7 +3,6 @@ import io
 import os
 import time
 
-import cv2
 import oss2
 import requests
 from PIL import Image
@@ -20,17 +19,13 @@ 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 resize(image, short_side_length=512):
+    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):
@@ -60,13 +55,13 @@ def download_images(img_urls, batch_size):
 
 def upload_np_2_oss(input_image, name="cache.jpg"):
     assert name.lower().endswith((".png", ".jpg")), name
-    if name.lower().endswith(".png"):
+    if name.endswith(".png"):
         name = name[:-4] + ".jpg"
     imgByteArr = io.BytesIO()
     if name.lower().endswith(".png"):
-        Image.fromarray(resize(input_image)).save(imgByteArr, format="PNG")
+        resize(Image.fromarray(input_image)).save(imgByteArr, format="PNG")
     else:
-        Image.fromarray(resize(input_image)).save(imgByteArr, format="JPEG", quality=95)
+        resize(Image.fromarray(input_image)).save(imgByteArr, format="JPEG", quality=95)
     imgByteArr = imgByteArr.getvalue()
 
     start_time = time.perf_counter()