demo/util.py

import concurrent.futures
import io
import os

import numpy as np
import oss2
import requests
from PIL import Image, ImageDraw, ImageFont

from .log import logger

# oss
access_key_id = os.getenv("ACCESS_KEY_ID")
access_key_secret = os.getenv("ACCESS_KEY_SECRET")
bucket_name = os.getenv("BUCKET_NAME")
endpoint = os.getenv("ENDPOINT")

oss_path = "VectorizeAnything"
oss_path_img_gallery = "VectorizeAnything"

bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)

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


def download_images(img_urls, batch_size, img_name):
    imgs_pil = [None] * batch_size
    # worker_results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
        to_do = []
        for i, url in enumerate(img_urls):
            future = executor.submit(download_img_pil, i, url, img_name)
            to_do.append(future)

        for future in concurrent.futures.as_completed(to_do):
            ret = future.result()
            # worker_results.append(ret)
            index, img_pil = ret
            imgs_pil[index] = img_pil  # 按顺序排列url，后续下载关联的图片或者svg需要使用

    return imgs_pil

def download_svg(index, img_url, svg_name):
    # print(img_url)
    r = requests.get(img_url, stream=True)
    if r.status_code == 200:
        img = io.BytesIO(r.content)
        svg_url = upload_svg_2_oss(img, f"{svg_name}_{index}.svg")
        return (index, svg_url)
    else:
        logger.error(f"Fail to download: {img_url}")


def download_svgs(img_urls, batch_size, svg_name):
    imgs_pil = [None] * batch_size
    # worker_results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
        to_do = []
        for i, url in enumerate(img_urls):
            future = executor.submit(download_svg, i, url, svg_name)
            to_do.append(future)

        for future in concurrent.futures.as_completed(to_do):
            ret = future.result()
            # worker_results.append(ret)
            index, img_pil = ret
            imgs_pil[index] = img_pil  # 按顺序排列url，后续下载关联的图片或者svg需要使用

    return imgs_pil

def upload_np_2_oss(input_image, name="cache.png", gallery=False):
    imgByteArr = io.BytesIO()
    Image.fromarray(input_image).save(imgByteArr, format="PNG")
    imgByteArr = imgByteArr.getvalue()

    if gallery:
        path = oss_path_img_gallery
    else:
        path = oss_path

    bucket.put_object(path + "/" + name, imgByteArr)  # data为数据，可以是图片
    ret = bucket.sign_url('GET', path + "/" + name, 60 * 60 * 24)  # 返回值为链接，参数依次为，方法/oss上文件路径/过期时间(s)
    del imgByteArr
    return ret


def upload_svg_2_oss(input_svg, name="cache.png", gallery=False):

    if gallery:
        path = oss_path_img_gallery
    else:
        path = oss_path

    bucket.put_object(path + "/" + name, input_svg)  # data为数据，可以是图片
    ret = bucket.sign_url('GET', path + "/" + name, 60 * 60 * 24)  # 返回值为链接，参数依次为，方法/oss上文件路径/过期时间(s)
    del input_svg
    return ret