235 lines
9.4 KiB
Python
235 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
||
"""野鸡图片全自动标注 v2。
|
||
|
||
方案：Grounding DINO（开放词汇检测，多提示词）→ NMS 去重 →
|
||
CLIP 裁剪验证（剔除误检）→ 仍无框时 CLIP 滑动窗口兜底。
|
||
输出 YOLO 格式标注 + 可视化预览，覆盖 datasets/images/pheasant 下全部图片。
|
||
"""
|
||
import os
|
||
import sys
|
||
import torch
|
||
import numpy as np
|
||
from pathlib import Path
|
||
from PIL import Image, ImageDraw
|
||
from torchvision.ops import nms
|
||
|
||
from transformers import (
|
||
GroundingDinoProcessor,
|
||
GroundingDinoForObjectDetection,
|
||
CLIPProcessor,
|
||
CLIPModel,
|
||
)
|
||
|
||
BASE = Path(__file__).parent
|
||
IMG_DIR = BASE / "datasets" / "images" / "pheasant"
|
||
OUT_DIR = BASE / "datasets" / "yolo_format"
|
||
PREVIEW_DIR = BASE / "datasets" / "preview"
|
||
|
||
DINO_PROMPTS = [
|
||
"pheasant",
|
||
"ring-necked pheasant",
|
||
"wild pheasant",
|
||
"common pheasant",
|
||
"bird",
|
||
]
|
||
DINO_BOX_THRESHOLD = 0.13
|
||
DINO_TEXT_THRESHOLD = 0.15
|
||
NMS_IOU = 0.45
|
||
KEEP_THRESHOLD = 0.16 # 低于此分的框必须通过 CLIP 验证才保留
|
||
MIN_BOX_AREA = 0.002 # 框面积占比下限（排除过小误检）
|
||
|
||
CLIP_POSITIVE = ["a photo of a pheasant", "a photo of a wild bird"]
|
||
CLIP_NEGATIVE = [
|
||
"a photo of grass and leaves",
|
||
"a photo of rocks and soil",
|
||
"a photo of a landscape",
|
||
"a photo of a fence",
|
||
"a photo of trees",
|
||
]
|
||
CLIP_VERIFY_THRESHOLD = 0.40 # softmax(正例) >= 此值才算验证通过
|
||
FALLBACK_WINDOW_SIZES = [0.5, 0.35, 0.25] # 滑动窗口占图宽比例
|
||
|
||
|
||
class AutoLabeler:
|
||
def __init__(self):
|
||
print("加载 Grounding DINO...", flush=True)
|
||
self.proc = GroundingDinoProcessor.from_pretrained("IDEA-Research/grounding-dino-tiny")
|
||
self.dino = GroundingDinoForObjectDetection.from_pretrained("IDEA-Research/grounding-dino-tiny")
|
||
self.dino.eval()
|
||
|
||
print("加载 CLIP...", flush=True)
|
||
self.clip_proc = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
|
||
self.clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
|
||
self.clip.eval()
|
||
|
||
pos_tokens = self.clip_proc(text=CLIP_POSITIVE + CLIP_NEGATIVE, return_tensors="pt", padding=True)
|
||
with torch.no_grad():
|
||
out = self.clip.get_text_features(**pos_tokens)
|
||
if hasattr(out, "text_embeds"):
|
||
feats = out.text_embeds
|
||
elif hasattr(out, "pooler_output"):
|
||
feats = out.pooler_output
|
||
else:
|
||
feats = out
|
||
feats = feats / feats.norm(dim=-1, keepdim=True)
|
||
self.pos_feats = feats[: len(CLIP_POSITIVE)]
|
||
self.neg_feats = feats[len(CLIP_POSITIVE):]
|
||
|
||
def clip_score(self, region: Image.Image) -> float:
|
||
"""返回裁剪区域是野鸡的概率（softmax 归一化，0~1）"""
|
||
inputs = self.clip_proc(images=region, return_tensors="pt")
|
||
with torch.no_grad():
|
||
out = self.clip.get_image_features(**inputs)
|
||
if hasattr(out, "pooler_output"):
|
||
img_feat = out.pooler_output
|
||
else:
|
||
img_feat = out
|
||
img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)
|
||
sim_pos = img_feat @ self.pos_feats.T # (1, P)
|
||
sim_neg = img_feat @ self.neg_feats.T # (1, N)
|
||
logits = torch.cat([sim_pos * 100, sim_neg * 100], dim=1)
|
||
prob = logits.softmax(dim=1)[0, : len(CLIP_POSITIVE)].max().item()
|
||
return prob
|
||
|
||
def dino_detect(self, img: Image.Image):
|
||
"""Grounding DINO 检测，返回 [x1,y1,x2,y2] 绝对坐标列表"""
|
||
inputs = self.proc(images=img, text=DINO_PROMPTS, return_tensors="pt")
|
||
with torch.no_grad():
|
||
out = self.dino(**inputs)
|
||
res = self.proc.post_process_grounded_object_detection(
|
||
out, input_ids=inputs["input_ids"],
|
||
threshold=DINO_BOX_THRESHOLD, text_threshold=DINO_TEXT_THRESHOLD,
|
||
target_sizes=[img.size[::-1]],
|
||
)[0]
|
||
w, h = img.size
|
||
boxes = res["boxes"].tolist()
|
||
scores = res["scores"].tolist()
|
||
valid = []
|
||
for (x1, y1, x2, y2), s in zip(boxes, scores):
|
||
bw, bh = (x2 - x1) / w, (y2 - y1) / h
|
||
if bw * bh < MIN_BOX_AREA:
|
||
continue
|
||
valid.append((x1, y1, x2, y2, s))
|
||
return valid
|
||
|
||
def clip_fallback(self, img: Image.Image):
|
||
"""滑动窗口找最像野鸡的区域（兜底），返回 [x1,y1,x2,y2] 或 None"""
|
||
w, h = img.size
|
||
base = 768
|
||
scale = base / max(w, h)
|
||
small = img.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS)
|
||
sw, sh = small.size
|
||
best = None
|
||
best_score = 0.0
|
||
for f in FALLBACK_WINDOW_SIZES:
|
||
ws, hs = int(sw * f), int(sh * f)
|
||
stride = max(int(ws * 0.5), 8)
|
||
for y in range(0, max(sh - hs, 1), stride):
|
||
for x in range(0, max(sw - ws, 1), stride):
|
||
region = small.crop((x, y, x + ws, y + hs))
|
||
s = self.clip_score(region)
|
||
if s > best_score:
|
||
best_score = s
|
||
best = (x, y, x + ws, y + hs)
|
||
if best is None:
|
||
return None, best_score
|
||
x1, y1, x2, y2 = [v / scale for v in best]
|
||
return (x1, y1, x2, y2), best_score
|
||
|
||
def label_image(self, img: Image.Image, name: str):
|
||
"""返回 (boxes, source, debug)，boxes 为 [x1,y1,x2,y2] 绝对坐标列表"""
|
||
w, h = img.size
|
||
|
||
cands = self.dino_detect(img)
|
||
if cands:
|
||
boxes = torch.tensor([[c[0], c[1], c[2], c[3]] for c in cands], dtype=torch.float32)
|
||
scores = torch.tensor([c[4] for c in cands], dtype=torch.float32)
|
||
keep = nms(boxes, scores, NMS_IOU)
|
||
kept = [(cands[i][0], cands[i][1], cands[i][2], cands[i][3], cands[i][4]) for i in keep]
|
||
|
||
# CLIP 验证每个框
|
||
final = []
|
||
for x1, y1, x2, y2, s in kept:
|
||
pad_x, pad_y = (x2 - x1) * 0.2, (y2 - y1) * 0.2
|
||
region = img.crop((max(x1 - pad_x, 0), max(y1 - pad_y, 0),
|
||
min(x2 + pad_x, w), min(y2 + pad_y, h)))
|
||
prob = self.clip_score(region)
|
||
final.append((x1, y1, x2, y2, s, prob))
|
||
best_clip = max(p for _, _, _, _, _, p in final)
|
||
if best_clip < CLIP_VERIFY_THRESHOLD:
|
||
print(f" {name}: 整图框分均过低，改用 CLIP 兜底 (best_clip={best_clip:.2f})", flush=True)
|
||
box, score = self.clip_fallback(img)
|
||
if box is not None:
|
||
x1, y1, x2, y2 = box
|
||
return [(x1, y1, x2, y2, score, f"clip({score:.2f})")], "clip"
|
||
return [], "none"
|
||
kept_final = [f for f in final if f[5] >= 0.20]
|
||
if not kept_final:
|
||
kept_final = [max(final, key=lambda t: t[5])]
|
||
for x1, y1, x2, y2, s, p in final:
|
||
if (x1, y1, x2, y2, s, p) not in kept_final:
|
||
print(f" {name}: 剔除低分框 clip={p:.2f} [{x1:.0f},{y1:.0f},{x2:.0f},{y2:.0f}]", flush=True)
|
||
return [(x1, y1, x2, y2, s, f"dino(clip={p:.2f})")
|
||
for x1, y1, x2, y2, s, p in kept_final], "dino"
|
||
else:
|
||
print(f" {name}: DINO 无候选，使用 CLIP 兜底", flush=True)
|
||
|
||
# 兜底：CLIP 滑动窗口
|
||
box, score = self.clip_fallback(img)
|
||
if box is not None:
|
||
x1, y1, x2, y2 = box
|
||
return [(x1, y1, x2, y2, score, f"clip({score:.2f})")], "clip"
|
||
return [], "none"
|
||
|
||
|
||
def main():
|
||
(OUT_DIR / "images" / "train").mkdir(parents=True, exist_ok=True)
|
||
(OUT_DIR / "labels" / "train").mkdir(parents=True, exist_ok=True)
|
||
PREVIEW_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
||
labeler = AutoLabeler()
|
||
|
||
files = sorted(f for f in IMG_DIR.iterdir() if f.suffix.lower() in {".png", ".jpg", ".jpeg"})
|
||
print(f"共 {len(files)} 张图片", flush=True)
|
||
|
||
summary = []
|
||
for p in files:
|
||
img = Image.open(p).convert("RGB")
|
||
w, h = img.size
|
||
boxes, src = labeler.label_image(img, p.name)
|
||
summary.append((p.name, len(boxes), src))
|
||
|
||
label_path = OUT_DIR / "labels" / "train" / (p.stem + ".txt")
|
||
img.save(OUT_DIR / "images" / "train" / p.name)
|
||
if boxes:
|
||
with open(label_path, "w") as f:
|
||
for x1, y1, x2, y2, s, _ in boxes:
|
||
cx, cy = (x1 + x2) / 2 / w, (y1 + y2) / 2 / h
|
||
bw, bh = (x2 - x1) / w, (y2 - y1) / h
|
||
f.write(f"0 {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}\n")
|
||
|
||
# 预览图
|
||
draw = ImageDraw.Draw(img)
|
||
for x1, y1, x2, y2, s, tag in boxes:
|
||
draw.rectangle([x1, y1, x2, y2], outline="#E53935", width=4)
|
||
draw.text((x1 + 4, max(y1 - 18, 0)), f"pheasant {s:.2f} [{tag}]", fill="#E53935")
|
||
img.save(PREVIEW_DIR / (p.stem + ".jpg"), "JPEG", quality=85)
|
||
|
||
status = "✓" if boxes else "✗"
|
||
print(f"{status} {p.name}: {len(boxes)} 框 ({src})", flush=True)
|
||
|
||
# 生成 data.yaml
|
||
with open(OUT_DIR / "data.yaml", "w") as f:
|
||
f.write(f"path: {OUT_DIR}\n")
|
||
f.write("train: images/train\nval: images/train\n")
|
||
f.write("nc: 1\nnames: ['pheasant']\n")
|
||
|
||
labeled = sum(1 for _, n, _ in summary if n > 0)
|
||
print(f"\n完成: {labeled}/{len(summary)} 张已标注", flush=True)
|
||
for name, n, src in summary:
|
||
print(f" {name}: {n} 框 ({src})", flush=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|