初始化 observer 项目:纯代码,不含权重与训练数据
This commit is contained in:
@@ -0,0 +1,202 @@
|
||||
"""用 LocalAI qwen3.5-9b 多模态模型自动标注图片。
|
||||
|
||||
用法：python auto_label.py
|
||||
输入：datasets/images/<class>/*.jpg
|
||||
输出：datasets/dataset/{train,val}/{images,labels} (YOLO 格式) + data.yaml
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import re
|
||||
import base64
|
||||
import shutil
|
||||
import sys
|
||||
import urllib.request
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
# LocalAI 服务器配置
|
||||
LOCAL_AI_URL = "http://192.168.3.210:18080/v1/chat/completions"
|
||||
MODEL_NAME = "qwen3.5-9b"
|
||||
|
||||
# 标注配置
|
||||
KEEP_PER_CLASS = 100
|
||||
MIN_CONF = 0.5 # 自动标注最低置信度
|
||||
MIN_BOX = 0.03 # 框面积占比下限
|
||||
|
||||
BASE = os.path.dirname(__file__)
|
||||
IMG_DIR = os.path.join(BASE, "datasets", "images")
|
||||
OUT_DIR = os.path.join(BASE, "datasets", "dataset")
|
||||
PREVIEW_DIR = os.path.join(BASE, "datasets", "preview")
|
||||
|
||||
CLASSES = ["pheasant"]
|
||||
PROMPTS = {
|
||||
"pheasant": "pheasant",
|
||||
}
|
||||
|
||||
# 标注提示词模板
|
||||
ANNOTATION_PROMPT = """你是一个目标检测助手。请检测图片中【{target}】动物。
|
||||
|
||||
输出格式（严格只输出这一行，不要其他文字）：
|
||||
pheasant 0.92 0.1 0.2 0.3 0.4
|
||||
|
||||
字段说明：
|
||||
- 第一个数字：置信度（0-1）
|
||||
- 后四个数字：边界框坐标 x1 y1 x2 y2（归一化 0-1）
|
||||
|
||||
如果没找到目标，只输出：NOT_FOUND
|
||||
"""
|
||||
|
||||
|
||||
def encode_image_to_base64(path: str) -> str:
|
||||
"""将图片路径转换为 base64 数据"""
|
||||
with open(path, "rb") as f:
|
||||
return base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
|
||||
def call_localai(image_path: str, prompt: str) -> dict:
|
||||
"""调用 LocalAI qwen3.5-9b 进行标注"""
|
||||
image_b64 = encode_image_to_base64(image_path)
|
||||
|
||||
data = {
|
||||
"model": MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
"temperature": 0.1,
|
||||
}
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(LOCAL_AI_URL, data=json.dumps(data).encode("utf-8")) as resp:
|
||||
if resp.status == 200:
|
||||
result = json.loads(resp.read().decode("utf-8"))
|
||||
return {"success": True, "data": result}
|
||||
else:
|
||||
return {"success": False, "error": f"HTTP {resp.status}"}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
|
||||
def parse_annotation(response: dict) -> list:
|
||||
"""解析 LocalAI qwen3.5-9b 返回的标注结果（文本格式）"""
|
||||
if not response.get("success"):
|
||||
return []
|
||||
|
||||
content = response.get("choices", [{}])[0].get("message", {}).get("content", "")
|
||||
|
||||
# 如果没找到目标
|
||||
if "NOT_FOUND" in content or "未找到" in content:
|
||||
return []
|
||||
|
||||
# 解析文本格式：类别 置信度 x1 y1 x2 y2
|
||||
# 例如：pheasant 0.92 0.1 0.2 0.3 0.4
|
||||
parts = content.strip().split()
|
||||
if len(parts) < 5:
|
||||
return []
|
||||
|
||||
cls_name = parts[0].lower()
|
||||
try:
|
||||
confidence = float(parts[1])
|
||||
x1, y1, x2, y2 = [float(v) for v in parts[2:6]]
|
||||
except ValueError:
|
||||
return []
|
||||
|
||||
# 确保坐标顺序
|
||||
if x1 > x2 or y1 > y2:
|
||||
x1, x2 = x2, x1
|
||||
y1, y2 = y2, y1
|
||||
|
||||
return [{
|
||||
"class": cls_name,
|
||||
"confidence": confidence,
|
||||
"bbox": [x1, y1, x2, y2],
|
||||
}]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
for split in ("train", "val"):
|
||||
for sub in ("images", "labels"):
|
||||
os.makedirs(os.path.join(OUT_DIR, split, sub), exist_ok=True)
|
||||
os.makedirs(PREVIEW_DIR, exist_ok=True)
|
||||
|
||||
for cls in CLASSES:
|
||||
src_dir = os.path.join(IMG_DIR, cls)
|
||||
files = sorted(f for f in os.listdir(src_dir) if f.endswith(".jpg"))
|
||||
print(f"[label] {cls}: {len(files)} 张", flush=True)
|
||||
|
||||
kept = []
|
||||
for f in files:
|
||||
path = os.path.join(src_dir, f)
|
||||
image = Image.open(path).convert("RGB")
|
||||
w, h = image.size
|
||||
|
||||
# 调用 LocalAI 进行标注
|
||||
prompt = ANNOTATION_PROMPT.format(target=cls)
|
||||
result = call_localai(path, prompt)
|
||||
annotations = parse_annotation(result)
|
||||
|
||||
# 取置信度最高的标注
|
||||
best = None
|
||||
best_score = 0
|
||||
for ann in annotations:
|
||||
if ann.get("class") != cls:
|
||||
continue
|
||||
score = ann.get("confidence", 0)
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best = ann
|
||||
|
||||
if best is None:
|
||||
print(f" {f}: ✗ (未找到目标)", flush=True)
|
||||
continue
|
||||
|
||||
if best_score < MIN_CONF:
|
||||
print(f" {f}: ✗ (置信度 {best_score:.2f} < {MIN_CONF})", flush=True)
|
||||
continue
|
||||
|
||||
x1, y1, x2, y2 = best.get("bbox", [0, 0, 1, 1])
|
||||
kept.append((path, (best_score, x1, y1, x2, y2)))
|
||||
print(f" {f}: ✓ ({best_score:.2f})", flush=True)
|
||||
|
||||
kept.sort(key=lambda t: t[1][0], reverse=True)
|
||||
kept = kept[:KEEP_PER_CLASS]
|
||||
avg = sum(k[1][0] for k in kept) / max(len(kept), 1)
|
||||
print(f" → 保留 {len(kept)} 张（平均置信度 {avg:.2f}）", flush=True)
|
||||
|
||||
for i, (path, (score, x1, y1, x2, y2)) in enumerate(kept):
|
||||
split = "train" if i % 10 else "val"
|
||||
dst_img = os.path.join(OUT_DIR, split, "images", f"{cls}_{i:03d}.jpg")
|
||||
shutil.copy(path, dst_img)
|
||||
|
||||
# 计算 YOLO 格式坐标
|
||||
w, h = Image.open(dst_img).size
|
||||
cx = (x1 + x2) / 2 / w
|
||||
cy = (y1 + y2) / 2 / h
|
||||
bw = (x2 - x1) / w
|
||||
bh = (y2 - y1) / h
|
||||
|
||||
label_file = os.path.join(OUT_DIR, split, "labels", f"{cls}_{i:03d}.txt")
|
||||
with open(label_file, "w") as f:
|
||||
f.write(f"{CLASSES.index(cls)} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}\n")
|
||||
|
||||
# 生成预览图
|
||||
img = Image.open(dst_img).convert("RGB")
|
||||
draw = ImageDraw.Draw(img)
|
||||
draw.rectangle([x1, y1, x2, y2], outline="#E53935", width=3)
|
||||
draw.text((x1, max(y1 - 14, 0)), f"{cls} {score:.0%}", fill="#E53935")
|
||||
img.save(os.path.join(PREVIEW_DIR, f"{cls}_{i:03d}.jpg"), "JPEG", quality=85)
|
||||
|
||||
# 生成 data.yaml
|
||||
with open(os.path.join(OUT_DIR, "data.yaml"), "w") as f:
|
||||
f.write(f"path: {OUT_DIR}\n")
|
||||
f.write("train: train/images\nval: val/images\n")
|
||||
f.write(f"names: {CLASSES}\n")
|
||||
print("[done] 数据集就绪:", OUT_DIR, flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,120 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
为训练图片自动生成标注
|
||||
使用 YOLO 模型检测野鸡并生成 YOLO 格式标注
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import os
|
||||
from pathlib import Path
|
||||
from ultralytics import YOLO
|
||||
|
||||
# 配置
|
||||
MODEL_PATH = Path(__file__).parent / "yolov8s-world.pt"
|
||||
IMAGES_DIR = Path(__file__).parent / "datasets" / "images" / "pheasant"
|
||||
OUTPUT_DIR = Path(__file__).parent / "datasets" / "yolo_format"
|
||||
|
||||
def create_yolo_structure():
|
||||
"""创建 YOLO 格式目录结构"""
|
||||
(OUTPUT_DIR / "images" / "train").mkdir(parents=True, exist_ok=True)
|
||||
(OUTPUT_DIR / "labels" / "train").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def convert_to_yolo_format(bbox, img_width, img_height):
|
||||
"""将边界框转换为 YOLO 格式 (cx, cy, w, h)"""
|
||||
x1, y1, x2, y2 = bbox
|
||||
cx = (x1 + x2) / 2 / img_width
|
||||
cy = (y1 + y2) / 2 / img_height
|
||||
w = (x2 - x1) / img_width
|
||||
h = (y2 - y1) / img_height
|
||||
return cx, cy, w, h
|
||||
|
||||
def main():
|
||||
# 加载模型
|
||||
print(f"加载模型: {MODEL_PATH}")
|
||||
model = YOLO(str(MODEL_PATH))
|
||||
|
||||
# 创建目录结构
|
||||
create_yolo_structure()
|
||||
|
||||
# 获取所有图片
|
||||
image_extensions = {".jpg", ".jpeg", ".png", ".bmp"}
|
||||
image_files = [
|
||||
f for f in IMAGES_DIR.iterdir()
|
||||
if f.suffix.lower() in image_extensions
|
||||
]
|
||||
|
||||
print(f"找到 {len(image_files)} 张图片")
|
||||
|
||||
# 处理每张图片
|
||||
labeled_count = 0
|
||||
for img_path in image_files:
|
||||
print(f"处理: {img_path.name}")
|
||||
|
||||
# 读取图片
|
||||
image = cv2.imread(str(img_path))
|
||||
if image is None:
|
||||
print(f" 跳过: 无法读取 {img_path.name}")
|
||||
continue
|
||||
|
||||
img_height, img_width = image.shape[:2]
|
||||
|
||||
# 推理
|
||||
results = model(image, conf=0.3, iou=0.45)
|
||||
|
||||
# 收集标注
|
||||
labels = []
|
||||
for result in results:
|
||||
boxes = result.boxes
|
||||
if boxes is None or len(boxes) == 0:
|
||||
continue
|
||||
|
||||
for box in boxes:
|
||||
class_name = result.names.get(int(box.cls[0]), "")
|
||||
|
||||
# 只保留 bird 或 pheasant 类别
|
||||
if class_name in ["bird", "pheasant"]:
|
||||
class_id = 0 # 只有一个类别：野鸡
|
||||
bbox = list(map(int, box.xyxy[0].tolist()))
|
||||
conf = float(box.conf[0])
|
||||
|
||||
# 转换为 YOLO 格式
|
||||
cx, cy, w, h = convert_to_yolo_format(bbox, img_width, img_height)
|
||||
labels.append(f"{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")
|
||||
|
||||
# 保存标注文件
|
||||
if len(labels) > 0:
|
||||
labeled_count += 1
|
||||
# 复制图片
|
||||
dst_img_path = OUTPUT_DIR / "images" / "train" / img_path.name
|
||||
cv2.imwrite(str(dst_img_path), image)
|
||||
|
||||
# 保存标注
|
||||
label_path = OUTPUT_DIR / "labels" / "train" / (img_path.stem + ".txt")
|
||||
with open(label_path, "w") as f:
|
||||
f.write("\n".join(labels))
|
||||
|
||||
print(f" ✓ 标注了 {len(labels)} 个目标")
|
||||
else:
|
||||
print(f" - 未检测到野鸡")
|
||||
|
||||
print(f"\n完成!")
|
||||
print(f" 总图片数: {len(image_files)}")
|
||||
print(f" 有效标注: {labeled_count}")
|
||||
print(f" 输出目录: {OUTPUT_DIR}")
|
||||
|
||||
# 创建 data.yaml
|
||||
data_yaml = OUTPUT_DIR / "data.yaml"
|
||||
with open(data_yaml, "w") as f:
|
||||
f.write(f"""# Observer 数据集配置 - 只识别野鸡
|
||||
path: {OUTPUT_DIR}
|
||||
train: images/train
|
||||
val: images/train
|
||||
|
||||
# 类别
|
||||
nc: 1
|
||||
names: ['pheasant']
|
||||
""")
|
||||
print(f" 数据配置: {data_yaml}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,234 @@
|
||||
#!/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())
|
||||
@@ -0,0 +1,207 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
使用 CLIP 模型的滑动窗口方法检测野鸡位置
|
||||
生成 YOLO 格式的标注文件
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import clip
|
||||
from PIL import Image
|
||||
from pathlib import Path
|
||||
from torchvision.ops import nms
|
||||
|
||||
# 配置
|
||||
IMAGES_DIR = Path(__file__).parent / "datasets" / "images" / "pheasant"
|
||||
OUTPUT_DIR = Path(__file__).parent / "datasets" / "yolo_format"
|
||||
|
||||
# 检测配置
|
||||
WINDOW_SIZES = [128, 256, 512] # 滑动窗口大小（增大）
|
||||
STRIDE_RATIO = 0.7 # 窗口滑动步长比例（增大）
|
||||
CONFIDENCE_THRESHOLD = 0.5 # 置信度阈值（提高）
|
||||
NMS_THRESHOLD = 0.5 # NMS 阈值（提高）
|
||||
|
||||
# 野鸡的文本描述
|
||||
PHEASANT_PROMPTS = [
|
||||
"a photo of a pheasant in the wild",
|
||||
"a wild pheasant in natural habitat",
|
||||
"a bird with colorful feathers in grass",
|
||||
]
|
||||
|
||||
class CLIPDetector:
|
||||
def __init__(self):
|
||||
print("加载 CLIP 模型...")
|
||||
self.model, self.preprocess = clip.load("ViT-B/32", device="cpu")
|
||||
self.model.eval()
|
||||
print("CLIP 模型加载完成")
|
||||
|
||||
# 预计算文本特征
|
||||
print("预计算文本特征...")
|
||||
text_tokens = clip.tokenize(PHEASANT_PROMPTS).to("cpu")
|
||||
with torch.no_grad():
|
||||
self.text_features = self.model.encode_text(text_tokens)
|
||||
self.text_features /= self.text_features.norm(dim=-1, keepdim=True)
|
||||
|
||||
def classify_region(self, image_region):
|
||||
"""对图像区域进行分类"""
|
||||
# 转换为 PIL Image
|
||||
if isinstance(image_region, np.ndarray):
|
||||
image_region = Image.fromarray(cv2.cvtColor(image_region, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 预处理
|
||||
image_input = self.preprocess(image_region).unsqueeze(0).to("cpu")
|
||||
|
||||
# 计算图像特征
|
||||
with torch.no_grad():
|
||||
image_features = self.model.encode_image(image_input)
|
||||
image_features /= image_features.norm(dim=-1, keepdim=True)
|
||||
|
||||
# 计算相似度
|
||||
similarity = (100.0 * image_features @ self.text_features.T).softmax(dim=-1)
|
||||
|
||||
# 返回最高分
|
||||
return float(similarity[0].max())
|
||||
|
||||
def detect(self, image):
|
||||
"""检测图片中的野鸡"""
|
||||
h, w = image.shape[:2]
|
||||
detections = []
|
||||
|
||||
# 多尺度滑动窗口
|
||||
for window_size in WINDOW_SIZES:
|
||||
stride = int(window_size * STRIDE_RATIO)
|
||||
|
||||
# 滑动窗口
|
||||
for y in range(0, h - window_size, stride):
|
||||
for x in range(0, w - window_size, stride):
|
||||
# 提取窗口区域
|
||||
region = image[y:y+window_size, x:x+window_size]
|
||||
|
||||
# 分类
|
||||
score = self.classify_region(region)
|
||||
|
||||
# 如果置信度足够高，保存检测结果
|
||||
if score > CONFIDENCE_THRESHOLD:
|
||||
detections.append({
|
||||
'bbox': [x, y, x + window_size, y + window_size],
|
||||
'score': score,
|
||||
})
|
||||
|
||||
# NMS 去重
|
||||
if len(detections) > 0:
|
||||
detections = self.nms(detections)
|
||||
|
||||
return detections
|
||||
|
||||
def nms(self, detections):
|
||||
"""非极大值抑制"""
|
||||
if len(detections) == 0:
|
||||
return []
|
||||
|
||||
# 转换为 torch 格式
|
||||
boxes = torch.tensor([d['bbox'] for d in detections], dtype=torch.float32)
|
||||
scores = torch.tensor([d['score'] for d in detections], dtype=torch.float32)
|
||||
|
||||
# 应用 NMS
|
||||
keep_indices = nms(boxes, scores, NMS_THRESHOLD)
|
||||
|
||||
# 保留 NMS 后的检测结果
|
||||
filtered_detections = [detections[i] for i in keep_indices]
|
||||
|
||||
return filtered_detections
|
||||
|
||||
def convert_to_yolo_format(bbox, img_width, img_height):
|
||||
"""将边界框转换为 YOLO 格式 (cx, cy, w, h)"""
|
||||
x1, y1, x2, y2 = bbox
|
||||
cx = (x1 + x2) / 2 / img_width
|
||||
cy = (y1 + y2) / 2 / img_height
|
||||
w = (x2 - x1) / img_width
|
||||
h = (y2 - y1) / img_height
|
||||
return cx, cy, w, h
|
||||
|
||||
def main():
|
||||
# 初始化检测器
|
||||
detector = CLIPDetector()
|
||||
|
||||
# 创建输出目录
|
||||
(OUTPUT_DIR / "images" / "train").mkdir(parents=True, exist_ok=True)
|
||||
(OUTPUT_DIR / "labels" / "train").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 获取所有图片
|
||||
image_extensions = {".jpg", ".jpeg", ".png", ".bmp"}
|
||||
image_files = [
|
||||
f for f in IMAGES_DIR.iterdir()
|
||||
if f.suffix.lower() in image_extensions
|
||||
]
|
||||
|
||||
print(f"找到 {len(image_files)} 张图片")
|
||||
|
||||
# 处理每张图片
|
||||
labeled_count = 0
|
||||
for img_path in image_files:
|
||||
print(f"\n处理: {img_path.name}")
|
||||
|
||||
# 读取图片
|
||||
image = cv2.imread(str(img_path))
|
||||
if image is None:
|
||||
print(f" 跳过: 无法读取 {img_path.name}")
|
||||
continue
|
||||
|
||||
img_height, img_width = image.shape[:2]
|
||||
|
||||
# 检测野鸡
|
||||
detections = detector.detect(image)
|
||||
|
||||
# 生成标注
|
||||
labels = []
|
||||
for det in detections:
|
||||
bbox = det['bbox']
|
||||
score = det['score']
|
||||
|
||||
# 转换为 YOLO 格式
|
||||
cx, cy, w, h = convert_to_yolo_format(bbox, img_width, img_height)
|
||||
labels.append(f"0 {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")
|
||||
|
||||
# 在图片上绘制检测框（用于可视化）
|
||||
x1, y1, x2, y2 = bbox
|
||||
cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
|
||||
cv2.putText(image, f"pheasant: {score:.2f}", (int(x1), int(y1) - 10),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
|
||||
|
||||
# 保存标注文件
|
||||
if len(labels) > 0:
|
||||
labeled_count += 1
|
||||
# 保存带标注的图片（用于可视化）
|
||||
cv2.imwrite(str(OUTPUT_DIR / "images" / "train" / img_path.name), image)
|
||||
|
||||
# 保存标注
|
||||
label_path = OUTPUT_DIR / "labels" / "train" / (img_path.stem + ".txt")
|
||||
with open(label_path, "w") as f:
|
||||
f.writelines(labels)
|
||||
|
||||
print(f" ✓ 标注了 {len(labels)} 个目标")
|
||||
else:
|
||||
print(f" - 未检测到野鸡")
|
||||
|
||||
print(f"\n完成!")
|
||||
print(f" 总图片数: {len(image_files)}")
|
||||
print(f" 有效标注: {labeled_count}")
|
||||
print(f" 输出目录: {OUTPUT_DIR}")
|
||||
|
||||
# 创建 data.yaml
|
||||
data_yaml = OUTPUT_DIR / "data.yaml"
|
||||
with open(data_yaml, "w") as f:
|
||||
f.write(f"""# Observer 数据集配置 - 只识别野鸡
|
||||
path: {OUTPUT_DIR}
|
||||
train: images/train
|
||||
val: images/train
|
||||
|
||||
# 类别
|
||||
nc: 1
|
||||
names: ['pheasant']
|
||||
""")
|
||||
print(f" 数据配置: {data_yaml}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,300 @@
|
||||
"""从 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())
|
||||
@@ -0,0 +1,112 @@
|
||||
"""用 LocalAI qwen3.5-9b 过滤不含活体动物的图片。
|
||||
|
||||
用法：python filter_images.py
|
||||
输入：datasets/images/<class>/*.jpg
|
||||
输出：删除不含目标动物的图片
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import base64
|
||||
import sys
|
||||
import urllib.request
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
|
||||
# LocalAI 服务器配置
|
||||
LOCAL_AI_URL = "http://192.168.3.210:18080/v1/chat/completions"
|
||||
MODEL_NAME = "qwen3.5-9b"
|
||||
|
||||
# 过滤配置
|
||||
BASE = os.path.dirname(__file__)
|
||||
IMG_DIR = os.path.join(BASE, "datasets", "images")
|
||||
CLASSES = ["pheasant"]
|
||||
MAX_IMAGE_SIZE = 400 # 缩放到最大边长
|
||||
|
||||
# 过滤提示词
|
||||
FILTER_PROMPT = """你是一个图片质量检查助手。请判断这张图片是否包含【{target}】的活体动物照片。
|
||||
|
||||
判断标准：
|
||||
✓ 保留：真实动物照片（活体、自然姿态、野外或自然环境）
|
||||
✗ 删除：标本照片、插画、图表、文字图片、logo、标志、空场景、纯风景
|
||||
|
||||
只回答：KEEP 或 DELETE"""
|
||||
|
||||
|
||||
def encode_image_to_base64(path: str) -> str:
|
||||
"""读取图片并压缩后转为 base64"""
|
||||
img = Image.open(path).convert("RGB") # 转为 RGB 避免 RGBA 问题
|
||||
# 缩放图片以减少 API 负载
|
||||
img.thumbnail((MAX_IMAGE_SIZE, MAX_IMAGE_SIZE), Image.Resampling.LANCZOS)
|
||||
buffer = BytesIO()
|
||||
img.save(buffer, format="JPEG", quality=70)
|
||||
return base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def call_localai(image_path: str, prompt: str) -> str:
|
||||
"""调用 LocalAI 检测图片是否包含活体动物"""
|
||||
image_b64 = encode_image_to_base64(image_path)
|
||||
|
||||
data = {
|
||||
"model": MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
"temperature": 0.1,
|
||||
}
|
||||
|
||||
for attempt in range(3):
|
||||
try:
|
||||
with urllib.request.urlopen(LOCAL_AI_URL, data=json.dumps(data).encode("utf-8"), timeout=60) as resp:
|
||||
if resp.status == 200:
|
||||
result = json.loads(resp.read().decode("utf-8"))
|
||||
return result.get("choices", [{}])[0].get("message", {}).get("content", "").strip().upper()
|
||||
except Exception as e:
|
||||
if attempt < 2:
|
||||
import time
|
||||
time.sleep(2 * (attempt + 1))
|
||||
else:
|
||||
print(f" ! API 错误: {e}", flush=True)
|
||||
return "ERROR"
|
||||
|
||||
|
||||
def main():
|
||||
total_deleted = 0
|
||||
|
||||
for cls in CLASSES:
|
||||
cls_dir = os.path.join(IMG_DIR, cls)
|
||||
if not os.path.exists(cls_dir):
|
||||
continue
|
||||
|
||||
files = sorted(f for f in os.listdir(cls_dir) if f.endswith(".jpg"))
|
||||
print(f"\n[filter] {cls}: {len(files)} 张", flush=True)
|
||||
|
||||
deleted = 0
|
||||
kept = 0
|
||||
|
||||
for f in files:
|
||||
path = os.path.join(cls_dir, f)
|
||||
prompt = FILTER_PROMPT.format(target=cls)
|
||||
result = call_localai(path, prompt)
|
||||
|
||||
if "DELETE" in result:
|
||||
os.remove(path)
|
||||
deleted += 1
|
||||
print(f" {f}: DELETE", flush=True)
|
||||
else:
|
||||
kept += 1
|
||||
print(f" {f}: KEEP", flush=True)
|
||||
|
||||
print(f" → 保留 {kept} 张，删除 {deleted} 张", flush=True)
|
||||
total_deleted += deleted
|
||||
|
||||
print(f"\n[done] 共删除 {total_deleted} 张图片", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,130 @@
|
||||
"""用本地 CLIP 模型过滤不含活体动物的图片。
|
||||
|
||||
用法：python filter_images_clip.py
|
||||
输入：datasets/images/<class>/*.jpg
|
||||
输出：删除不含目标动物的图片
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
import clip
|
||||
from PIL import Image
|
||||
|
||||
# 过滤配置
|
||||
BASE = os.path.dirname(__file__)
|
||||
IMG_DIR = os.path.join(BASE, "datasets", "images")
|
||||
CLASSES = ["pheasant"]
|
||||
MODEL_PATH = os.path.join(BASE, "..", "weights", "clip", "ViT-B-32.pt")
|
||||
|
||||
# 判断阈值 - 相似度低于此值的图片将被删除
|
||||
SIMILARITY_THRESHOLD = 0.10
|
||||
|
||||
# 每个类别的正向和负向描述
|
||||
PROMPTS = {
|
||||
"pheasant": {
|
||||
"positive": [
|
||||
"a photo of a pheasant in the wild",
|
||||
"a wild pheasant in natural outdoor habitat",
|
||||
"a pheasant walking in grass or field",
|
||||
"a pheasant hiding in bushes",
|
||||
"a pheasant in natural environment",
|
||||
],
|
||||
"negative": [
|
||||
"a taxidermy pheasant",
|
||||
"an illustration of a pheasant",
|
||||
"a drawing of a bird",
|
||||
"a painting of a bird",
|
||||
"a cartoon of a bird",
|
||||
"a person holding a bird",
|
||||
"a person catching a bird",
|
||||
"a bird in a cage",
|
||||
"a bird indoors",
|
||||
"a bird in a house",
|
||||
"a bird in a zoo",
|
||||
"a bird in captivity",
|
||||
"a logo or icon",
|
||||
"text or writing",
|
||||
"a landscape without animals",
|
||||
"a statue or sculpture",
|
||||
"a stuffed animal",
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
# 加载 CLIP 模型
|
||||
print("[init] 加载 CLIP 模型...", flush=True)
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model, preprocess = clip.load("ViT-B/32", device=device, download_root=os.path.join(BASE, "..", "weights"))
|
||||
|
||||
# 缓存文本特征
|
||||
text_features_cache = {}
|
||||
for cls in CLASSES:
|
||||
if cls not in PROMPTS:
|
||||
continue
|
||||
|
||||
pos_texts = clip.tokenize(PROMPTS[cls]["positive"]).to(device)
|
||||
neg_texts = clip.tokenize(PROMPTS[cls]["negative"]).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pos_features = model.encode_text(pos_texts)
|
||||
neg_features = model.encode_text(neg_texts)
|
||||
# 取正向描述的平均特征
|
||||
pos_features = pos_features.mean(dim=0, keepdim=True)
|
||||
pos_features /= pos_features.norm(dim=-1, keepdim=True)
|
||||
neg_features = neg_features.mean(dim=0, keepdim=True)
|
||||
neg_features /= neg_features.norm(dim=-1, keepdim=True)
|
||||
|
||||
text_features_cache[cls] = (pos_features, neg_features)
|
||||
|
||||
total_deleted = 0
|
||||
|
||||
for cls in CLASSES:
|
||||
cls_dir = os.path.join(IMG_DIR, cls)
|
||||
if not os.path.exists(cls_dir):
|
||||
continue
|
||||
|
||||
files = sorted(f for f in os.listdir(cls_dir) if f.endswith(".jpg"))
|
||||
print(f"\n[filter] {cls}: {len(files)} 张", flush=True)
|
||||
|
||||
deleted = 0
|
||||
kept = 0
|
||||
pos_features, neg_features = text_features_cache[cls]
|
||||
|
||||
for f in files:
|
||||
path = os.path.join(cls_dir, f)
|
||||
try:
|
||||
image = preprocess(Image.open(path).convert("RGB")).unsqueeze(0).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
image_features = model.encode_image(image)
|
||||
image_features /= image_features.norm(dim=-1, keepdim=True)
|
||||
|
||||
# 计算与正向和负向描述的相似度
|
||||
pos_similarity = (image_features @ pos_features.T).item()
|
||||
neg_similarity = (image_features @ neg_features.T).item()
|
||||
|
||||
# 综合得分：正向相似度 - 负向相似度
|
||||
score = pos_similarity - neg_similarity
|
||||
|
||||
if score < SIMILARITY_THRESHOLD:
|
||||
os.remove(path)
|
||||
deleted += 1
|
||||
print(f" {f}: DELETE (score={score:.3f})", flush=True)
|
||||
else:
|
||||
kept += 1
|
||||
print(f" {f}: KEEP (score={score:.3f})", flush=True)
|
||||
|
||||
except Exception as e:
|
||||
print(f" {f}: ERROR ({e})", flush=True)
|
||||
kept += 1 # 出错时保留
|
||||
|
||||
print(f" → 保留 {kept} 张，删除 {deleted} 张", flush=True)
|
||||
total_deleted += deleted
|
||||
|
||||
print(f"\n[done] 共删除 {total_deleted} 张图片", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,59 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
训练 YOLOv8n 模型 - 只识别野鸡（pheasant）
|
||||
使用现有的训练数据进行迁移学习
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from ultralytics import YOLO
|
||||
|
||||
# 配置
|
||||
DATA_DIR = Path(__file__).parent / "datasets"
|
||||
MODEL_NAME = "yolov8n.pt" # 预训练模型
|
||||
OUTPUT_DIR = Path(__file__).parent / "runs"
|
||||
|
||||
def main():
|
||||
# 创建数据集配置文件（只包含野鸡类别）
|
||||
data_yaml = DATA_DIR / "data.yaml"
|
||||
print("创建数据集配置文件（只训练野鸡类别）...")
|
||||
with open(data_yaml, "w") as f:
|
||||
f.write(f"""# Observer 数据集配置 - 只识别野鸡
|
||||
path: {DATA_DIR}
|
||||
train: images/pheasant
|
||||
val: images/pheasant
|
||||
|
||||
# 类别
|
||||
nc: 1
|
||||
names: ['pheasant']
|
||||
""")
|
||||
|
||||
# 加载预训练模型
|
||||
print(f"加载预训练模型: {MODEL_NAME}")
|
||||
model = YOLO(MODEL_NAME)
|
||||
|
||||
# 训练模型
|
||||
print("开始训练...")
|
||||
results = model.train(
|
||||
data=str(data_yaml),
|
||||
epochs=50,
|
||||
imgsz=640,
|
||||
batch=16,
|
||||
name="observer_yolov8n",
|
||||
patience=20,
|
||||
save=True,
|
||||
plots=True
|
||||
)
|
||||
|
||||
print(f"\n训练完成!")
|
||||
print(f"最佳模型保存在: {OUTPUT_DIR / 'observer_yolov8n' / 'weights' / 'best.pt'}")
|
||||
|
||||
# 导出为 TFLite 格式
|
||||
print("\n导出为 TFLite 格式...")
|
||||
best_model_path = OUTPUT_DIR / "observer_yolov8n" / "weights" / "best.pt"
|
||||
if best_model_path.exists():
|
||||
best_model = YOLO(str(best_model_path))
|
||||
best_model.export(format="tflite", imgsz=320)
|
||||
print(f"TFLite 模型导出完成")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,154 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
使用 Hugging Face transformers 进行零样本目标检测
|
||||
使用 OWL-ViT 模型检测野鸡
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from transformers import OwlViTProcessor, OwlViTForObjectDetection
|
||||
from PIL import Image
|
||||
|
||||
# 配置
|
||||
IMAGES_DIR = Path(__file__).parent / "datasets" / "images" / "pheasant"
|
||||
OUTPUT_DIR = Path(__file__).parent / "datasets" / "yolo_format"
|
||||
|
||||
# 检测配置
|
||||
CONFIDENCE_THRESHOLD = 0.1 # 置信度阈值
|
||||
|
||||
# 文本描述
|
||||
TEXT_PROMPTS = ["pheasant", "wild bird", "bird in grass"]
|
||||
|
||||
class ZeroShotDetector:
|
||||
def __init__(self):
|
||||
print("加载 OWL-ViT 模型...")
|
||||
self.processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
|
||||
self.model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
|
||||
self.model.eval()
|
||||
print("OWL-ViT 模型加载完成")
|
||||
|
||||
def detect(self, image_path):
|
||||
"""检测图片中的野鸡"""
|
||||
# 读取图片
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
|
||||
# 准备输入
|
||||
inputs = self.processor(text=TEXT_PROMPTS, images=image, return_tensors="pt")
|
||||
|
||||
# 推理
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
|
||||
# 获取结果
|
||||
target_sizes = torch.tensor([image.size[::-1]]) # [height, width]
|
||||
results = self.processor.post_process_grounded_object_detection(
|
||||
outputs, threshold=CONFIDENCE_THRESHOLD, target_sizes=target_sizes
|
||||
)[0]
|
||||
|
||||
# 解析结果
|
||||
detections = []
|
||||
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
|
||||
box = box.tolist()
|
||||
detections.append({
|
||||
'bbox': box, # [x1, y1, x2, y2]
|
||||
'score': score.item(),
|
||||
'label': TEXT_PROMPTS[label],
|
||||
})
|
||||
|
||||
return detections
|
||||
|
||||
def convert_to_yolo_format(bbox, img_width, img_height):
|
||||
"""将边界框转换为 YOLO 格式 (cx, cy, w, h)"""
|
||||
x1, y1, x2, y2 = bbox
|
||||
cx = (x1 + x2) / 2 / img_width
|
||||
cy = (y1 + y2) / 2 / img_height
|
||||
w = (x2 - x1) / img_width
|
||||
h = (y2 - y1) / img_height
|
||||
return cx, cy, w, h
|
||||
|
||||
def main():
|
||||
# 初始化检测器
|
||||
detector = ZeroShotDetector()
|
||||
|
||||
# 创建输出目录
|
||||
(OUTPUT_DIR / "images" / "train").mkdir(parents=True, exist_ok=True)
|
||||
(OUTPUT_DIR / "labels" / "train").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 获取所有图片
|
||||
image_extensions = {".jpg", ".jpeg", ".png", ".bmp"}
|
||||
image_files = [
|
||||
f for f in IMAGES_DIR.iterdir()
|
||||
if f.suffix.lower() in image_extensions
|
||||
]
|
||||
|
||||
print(f"找到 {len(image_files)} 张图片")
|
||||
|
||||
# 处理每张图片
|
||||
labeled_count = 0
|
||||
for img_path in image_files:
|
||||
print(f"\n处理: {img_path.name}")
|
||||
|
||||
# 读取图片
|
||||
image = cv2.imread(str(img_path))
|
||||
if image is None:
|
||||
print(f" 跳过: 无法读取 {img_path.name}")
|
||||
continue
|
||||
|
||||
img_height, img_width = image.shape[:2]
|
||||
|
||||
# 检测野鸡
|
||||
detections = detector.detect(img_path)
|
||||
|
||||
# 生成标注
|
||||
labels = []
|
||||
for det in detections:
|
||||
bbox = det['bbox']
|
||||
score = det['score']
|
||||
|
||||
# 转换为 YOLO 格式
|
||||
cx, cy, w, h = convert_to_yolo_format(bbox, img_width, img_height)
|
||||
labels.append(f"0 {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")
|
||||
|
||||
# 在图片上绘制检测框（用于可视化）
|
||||
x1, y1, x2, y2 = bbox
|
||||
cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
|
||||
cv2.putText(image, f"pheasant: {score:.2f}", (int(x1), int(y1) - 10),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
|
||||
|
||||
# 保存标注文件
|
||||
if len(labels) > 0:
|
||||
labeled_count += 1
|
||||
# 保存带标注的图片（用于可视化）
|
||||
cv2.imwrite(str(OUTPUT_DIR / "images" / "train" / img_path.name), image)
|
||||
|
||||
# 保存标注
|
||||
label_path = OUTPUT_DIR / "labels" / "train" / (img_path.stem + ".txt")
|
||||
with open(label_path, "w") as f:
|
||||
f.writelines(labels)
|
||||
|
||||
print(f" ✓ 标注了 {len(labels)} 个目标")
|
||||
else:
|
||||
print(f" - 未检测到野鸡")
|
||||
|
||||
print(f"\n完成!")
|
||||
print(f" 总图片数: {len(image_files)}")
|
||||
print(f" 有效标注: {labeled_count}")
|
||||
print(f" 输出目录: {OUTPUT_DIR}")
|
||||
|
||||
# 创建 data.yaml
|
||||
data_yaml = OUTPUT_DIR / "data.yaml"
|
||||
with open(data_yaml, "w") as f:
|
||||
f.write(f"""# Observer 数据集配置 - 只识别野鸡
|
||||
path: {OUTPUT_DIR}
|
||||
train: images/train
|
||||
val: images/train
|
||||
|
||||
# 类别
|
||||
nc: 1
|
||||
names: ['pheasant']
|
||||
""")
|
||||
print(f" 数据配置: {data_yaml}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user