src/util.py

import concurrent.futures
import io
import os
import time

import oss2
import requests
from PIL import Image

from .log import logger

# oss
oss_ak = os.getenv("OSS_AK")
oss_sk = os.getenv("OSS_SK")
oss_bucket = os.getenv("OSS_BUCKET")
oss_endpoint = os.getenv("OSS_ENDPOINT")
oss_path = os.getenv("OSS_PATH")

bucket = oss2.Bucket(oss2.Auth(oss_ak, oss_sk), oss_endpoint, oss_bucket)


def download_pil_img(index, img_url):
    r = requests.get(img_url, stream=True)
    if r.status_code == 200:
        img = Image.open(io.BytesIO(r.content))
        return (index, img)
    else:
        logger.error(f"Fail to download: {img_url}")


def download_images(img_urls, n=1):
    imgs_pil = [None] * n
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
        to_do = []
        for i, url in enumerate(img_urls):
            future = executor.submit(download_pil_img, i, url)
            to_do.append(future)

        for future in concurrent.futures.as_completed(to_do):
            ret = future.result()
            index, img_pil = ret
            imgs_pil[index] = img_pil

    return imgs_pil


def resize_img_with_max_size(image, max_side_length=512):
    width, height = image.size
    ratio = max_side_length / max(width, height)
    new_width = int(width * ratio)
    new_height = int(height * ratio)
    resized_image = image.resize((new_width, new_height))
    return resized_image


def upload_pil_2_oss(pil_image, name="cache.jpg"):
    image = resize_img_with_max_size(pil_image, max_side_length=512)

    imgByteArr = io.BytesIO()
    if name.lower().endswith(".png"):
        image.save(imgByteArr, format="PNG")
    else:
        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