301 lines
9.7 KiB
Python
301 lines
9.7 KiB
Python
"""从 Wikimedia Commons 下载训练图片（自由版权，可离线使用）。
|
||
|
||
策略：优先用物种分类目录（图片内容精确），再用全文搜索补充场景/姿态/光线多样性。
|
||
输出: datasets/images/<class>/<index>.jpg
|
||
串行 + 失败重试 + pHash 视觉去重。
|
||
"""
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from threading import Lock
|
||
|
||
import imagehash
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
API = "https://commons.wikimedia.org/w/api.php"
|
||
OUT = os.path.join(os.path.dirname(__file__), "datasets", "images")
|
||
|
||
# 更多图片来源
|
||
ADDITIONAL_SOURCES = {
|
||
"inaturalist": "https://api.inaturalist.org/v1/observations",
|
||
"flickr": "https://api.flickr.com/services/rest/",
|
||
}
|
||
|
||
SOURCES = {
|
||
"pheasant": {
|
||
"categories": ["Phasianus colchicus"],
|
||
"queries": [
|
||
# 户外真实场景（重点！）
|
||
"pheasant in field",
|
||
"pheasant in grassland",
|
||
"pheasant in meadow",
|
||
"pheasant in farmland",
|
||
"pheasant in countryside",
|
||
"pheasant in wild",
|
||
# 部分遮挡场景（重点！实际使用场景）
|
||
"pheasant hiding in grass",
|
||
"pheasant hiding in bushes",
|
||
"pheasant partially hidden",
|
||
"pheasant behind vegetation",
|
||
"pheasant peeking through grass",
|
||
"pheasant concealed in foliage",
|
||
"pheasant camouflaged",
|
||
"pheasant blending in",
|
||
# 不同距离和角度
|
||
"pheasant distant view",
|
||
"pheasant far away",
|
||
"pheasant small in frame",
|
||
"pheasant side view",
|
||
"pheasant back view",
|
||
"pheasant from behind",
|
||
# 不同姿态和行为
|
||
"pheasant walking in grass",
|
||
"pheasant foraging",
|
||
"pheasant feeding",
|
||
"pheasant running",
|
||
"pheasant flying low",
|
||
# 光线条件（户外真实光线）
|
||
"pheasant natural light",
|
||
"pheasant daylight",
|
||
"pheasant shade",
|
||
"pheasant shadow",
|
||
"pheasant backlit",
|
||
"pheasant overcast",
|
||
# 季节环境
|
||
"pheasant in autumn",
|
||
"pheasant in winter",
|
||
"pheasant in spring",
|
||
"pheasant in summer",
|
||
"pheasant in dry grass",
|
||
"pheasant in green grass",
|
||
"pheasant in snow",
|
||
"pheasant in mud",
|
||
],
|
||
},
|
||
}
|
||
PER_CLASS = 80
|
||
MIN_PIXEL = 320
|
||
RETRY = 2
|
||
WORKERS = 1
|
||
SLEEP_S = 2.0
|
||
RATE_LIMIT_WAIT_S = 60
|
||
|
||
|
||
def api_request(params: dict) -> dict:
|
||
url = API + "?" + urllib.parse.urlencode(params)
|
||
for attempt in range(RETRY + 1):
|
||
try:
|
||
req = urllib.request.Request(url, headers={"User-Agent": "observer-training/1.0"})
|
||
with urllib.request.urlopen(req, timeout=30) as r:
|
||
return json.loads(r.read().decode("utf-8"))
|
||
except Exception:
|
||
if attempt == RETRY:
|
||
raise
|
||
time.sleep(2 * (attempt + 1))
|
||
|
||
|
||
def collect_pages(pages: dict, results: list) -> None:
|
||
for p in pages.values():
|
||
info = (p.get("imageinfo") or [{}])[0]
|
||
thumb = info.get("thumburl")
|
||
if not thumb:
|
||
continue
|
||
w, h = info.get("width", 0), info.get("height", 0)
|
||
if min(w, h) < MIN_PIXEL:
|
||
continue
|
||
results.append({"url": thumb, "w": w, "h": h, "title": p.get("title", "")})
|
||
|
||
|
||
def category_images(category: str, limit: int) -> list[dict]:
|
||
results = []
|
||
params = {
|
||
"action": "query",
|
||
"generator": "categorymembers",
|
||
"gcmtitle": f"Category:{category}",
|
||
"gcmtype": "file",
|
||
"gcmlimit": "50",
|
||
"prop": "imageinfo",
|
||
"iiprop": "url|size",
|
||
"iiurlwidth": "640",
|
||
"format": "json",
|
||
}
|
||
while len(results) < limit:
|
||
data = api_request(params)
|
||
pages = (data.get("query", {}) or {}).get("pages", {})
|
||
collect_pages(pages, results)
|
||
cont = (data.get("continue") or {}).get("gcmcontinue")
|
||
if not cont:
|
||
break
|
||
params["gcmcontinue"] = cont
|
||
return results[:limit]
|
||
|
||
|
||
def search_images(query: str, limit: int) -> list[dict]:
|
||
results = []
|
||
params = {
|
||
"action": "query",
|
||
"generator": "search",
|
||
"gsrsearch": f"filetype:bitmap {query}",
|
||
"gsrnamespace": "6",
|
||
"gsrlimit": str(limit),
|
||
"prop": "imageinfo",
|
||
"iiprop": "url|size",
|
||
"iiurlwidth": "640",
|
||
"format": "json",
|
||
}
|
||
data = api_request(params)
|
||
pages = (data.get("query", {}) or {}).get("pages", {})
|
||
collect_pages(pages, results)
|
||
return results
|
||
|
||
|
||
def is_text_heavy(image_path: str) -> bool:
|
||
"""检测图片是否包含大量文字或为图表/标志"""
|
||
try:
|
||
img = Image.open(image_path).convert("L") # 灰度
|
||
arr = np.array(img, dtype=np.float32)
|
||
|
||
# 1. 检测边缘密度（文字产生大量边缘）
|
||
# 简单边缘检测：计算像素梯度
|
||
dx = np.abs(np.diff(arr, axis=1))
|
||
dy = np.abs(np.diff(arr, axis=0))
|
||
edge_density = (np.mean(dx) + np.mean(dy)) / 2
|
||
|
||
# 2. 检测颜色方差（图表通常颜色单一）
|
||
# 转回RGB检查
|
||
img_rgb = Image.open(image_path).convert("RGB")
|
||
arr_rgb = np.array(img_rgb, dtype=np.float32)
|
||
color_variance = np.std(arr_rgb)
|
||
|
||
# 3. 检测对比度（文字通常有高对比度的边缘）
|
||
contrast = np.std(arr)
|
||
|
||
# 判断逻辑：
|
||
# - 高边缘密度 + 低颜色方差 = 可能是图表/标志
|
||
# - 高对比度 + 高边缘密度 = 可能是文字图片
|
||
if edge_density > 30 and color_variance < 50:
|
||
return True
|
||
if edge_density > 40 and contrast > 80:
|
||
return True
|
||
|
||
return False
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def download(url: str, path: str) -> bool:
|
||
for attempt in range(RETRY + 1):
|
||
try:
|
||
req = urllib.request.Request(url, headers={"User-Agent": "observer-training/1.0"})
|
||
with urllib.request.urlopen(req, timeout=60) as r:
|
||
data = r.read()
|
||
if "image/" not in r.headers.get("Content-Type", ""):
|
||
return False
|
||
if len(data) < 10_000:
|
||
return False
|
||
with open(path, "wb") as f:
|
||
f.write(data)
|
||
return True
|
||
except urllib.error.HTTPError as e:
|
||
if e.code == 429:
|
||
# 限流：冷却后重试
|
||
print(f" ⏳ 限流(429)，冷却 {RATE_LIMIT_WAIT_S}s", flush=True)
|
||
time.sleep(RATE_LIMIT_WAIT_S)
|
||
continue
|
||
if attempt == RETRY:
|
||
return False
|
||
time.sleep(2 * (attempt + 1))
|
||
except Exception:
|
||
if attempt == RETRY:
|
||
return False
|
||
time.sleep(2 * (attempt + 1))
|
||
|
||
|
||
def main():
|
||
for cls, src in SOURCES.items():
|
||
cls_dir = os.path.join(OUT, cls)
|
||
os.makedirs(cls_dir, exist_ok=True)
|
||
existing = len([f for f in os.listdir(cls_dir) if f.endswith(".jpg")])
|
||
if existing >= PER_CLASS:
|
||
print(f"[skip] {cls}: 已有 {existing} 张", flush=True)
|
||
continue
|
||
|
||
# 重建去重池（续跑兼容）
|
||
seen_hashes = []
|
||
for f in os.listdir(cls_dir):
|
||
if f.endswith(".jpg"):
|
||
try:
|
||
seen_hashes.append(imagehash.phash(Image.open(os.path.join(cls_dir, f))))
|
||
except Exception:
|
||
pass
|
||
|
||
candidates = []
|
||
seen_titles = set()
|
||
for cat in src["categories"]:
|
||
print(f"[category] {cls} <- {cat}", flush=True)
|
||
try:
|
||
for item in category_images(cat, 80):
|
||
if item["title"] in seen_titles:
|
||
continue
|
||
seen_titles.add(item["title"])
|
||
candidates.append(item)
|
||
except Exception as e:
|
||
print(f" ! category 失败: {e}", flush=True)
|
||
for query in src["queries"]:
|
||
if len(candidates) >= PER_CLASS * 2:
|
||
break
|
||
print(f"[fetch] {cls} <- \"{query}\"", flush=True)
|
||
try:
|
||
for item in search_images(query, 80):
|
||
if item["title"] in seen_titles:
|
||
continue
|
||
seen_titles.add(item["title"])
|
||
candidates.append(item)
|
||
except Exception as e:
|
||
print(f" ! 搜索失败: {e}", flush=True)
|
||
print(f"[download] {cls}: 候选 {len(candidates)} 张", flush=True)
|
||
|
||
saved = existing
|
||
fail = 0
|
||
dup = 0
|
||
for item in candidates:
|
||
if saved >= PER_CLASS:
|
||
break
|
||
saved += 1
|
||
path = os.path.join(cls_dir, f"{saved:03d}.jpg")
|
||
if not download(item["url"], path):
|
||
fail += 1
|
||
saved -= 1
|
||
continue
|
||
|
||
# 过滤文字/图表类图片
|
||
if is_text_heavy(path):
|
||
os.remove(path)
|
||
saved -= 1
|
||
continue
|
||
|
||
try:
|
||
h = imagehash.phash(Image.open(path))
|
||
if any(h - other <= 8 for other in seen_hashes):
|
||
os.remove(path)
|
||
dup += 1
|
||
saved -= 1
|
||
continue
|
||
seen_hashes.append(h)
|
||
except Exception:
|
||
pass
|
||
if saved % 10 == 0:
|
||
print(f" + {cls}: {saved}", flush=True)
|
||
time.sleep(SLEEP_S)
|
||
print(f"[done] {cls}: {saved} 张（失败 {fail}，重复 {dup}）", flush=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|