#!/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())