utils.py

import os
import requests
from tqdm import tqdm


def download_file(url, save_path):
    if os.path.exists(save_path):
        print("目标已存在，无需下载")
        return

    create_dir(os.path.dirname(save_path))

    response = requests.get(url, stream=True)
    total_size = int(response.headers.get('content-length', 0))

    # 使用 tqdm 创建一个进度条
    progress_bar = tqdm(total=total_size, unit='B', unit_scale=True)

    with open(save_path, 'wb') as file:
        for data in response.iter_content(chunk_size=1024):
            file.write(data)
            progress_bar.update(len(data))

    progress_bar.close()

    if total_size != 0 and progress_bar.n != total_size:
        os.remove(save_path)
        print("下载失败，重试中...")
        download_file(url, save_path)
    else:
        print("下载完成")


def create_dir(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)


def get_jpg_files(folder_path):
    """
    获取指定文件夹下所有 .jpg 文件的路径。

    Parameters:
    - folder_path (str): 文件夹路径。

    Returns:
    - list: 包含所有 .jpg 文件路径的列表。
    """
    all_files = os.listdir(folder_path)
    jpg_files = [os.path.join(folder_path, file)
                 for file in all_files if file.lower().endswith('.jpg')]
    return jpg_files