import os
import re

import requests
from bs4 import BeautifulSoup
from PIL import Image
from io import BytesIO


headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
}

animal_names = ["猫", "狗", "cat", "dog"]

save_dir = "D:/project_space/01/005/datasets/001/images"
if not os.path.exists(save_dir):
    os.makedirs(save_dir)


def download_images(fish_name, num_images=5):
    search_url = f"https://www.google.com/search?hl=zh-CN&tbm=isch&q={fish_name}"

    response = requests.get(search_url, headers=headers)
    soup = BeautifulSoup(response.text, "html.parser")
    img_tags = soup.find_all("img")

    img_urls = []

    for img_tag in img_tags:
        img_url = img_tag.get("src")
        if img_url and img_url.startswith("http"):
            img_urls.append(img_url)
        if len(img_urls) >= num_images:
            break

    for idx, img_url in enumerate(img_urls):
        try:
            img_data = requests.get(img_url).content
            img_name = os.path.join(save_dir, f"{fish_name}_{idx + 1}.jpg")
            with open(img_name, "wb") as f:
                f.write(img_data)
            print(f"下载成功：{img_name}")
        except Exception as e:
            print(f"下载失败：{e}")


for animal_name in animal_names:
    download_images(animal_name, 200)

print("所有图片下载完成！")