优化接口图片传输方式

util/image_util.py

@@ -1,12 +1,16 @@
 import logging
 import math
+import os
 import urllib.request
 
 import cv2
 import numpy
+import requests
 from paddleclas import PaddleClas
 from tenacity import retry, stop_after_attempt, wait_random
 
+from log import PROJECT_ROOT
+
 
 @retry(stop=stop_after_attempt(3), wait=wait_random(1, 3), reraise=True,
        after=lambda x: logging.warning("获取图片失败！"))
@@ -247,3 +251,35 @@ def combined(img1, img2):
     combined_img[:height1, :width1] = img1
     combined_img[:height2, width1:width1 + width2] = img2
     return combined_img
+
+
+def parse_img_url(url):
+    """
+    解析图片url
+    :param url: 图片url
+    :return: 图片名称和图片后缀
+    """
+    url = url.split('?')[0]
+    return os.path.basename(url).rsplit('.', 1)
+
+
+@retry(stop=stop_after_attempt(3), wait=wait_random(1, 3), reraise=True,
+       after=lambda x: logging.warning('保存图片失败！'))
+def save_to_local(img_url, save_path=None):
+    """
+    保存图片到本地
+    :param img_url: 图片url
+    :param save_path: 本地保存地址，精确到文件名
+    :return: 本地保存地址
+    """
+    response = requests.get(img_url)
+    response.raise_for_status()  # 检查响应状态码是否正常
+
+    if save_path is None:
+        img_name, img_ext = parse_img_url(img_url)
+        save_path = os.path.join(PROJECT_ROOT, 'tmp_img', img_name + '.' + img_ext)
+
+    with open(save_path, 'wb') as file:
+        file.write(response.content)
+
+    return save_path