155 lines
4.7 KiB
Python
155 lines
4.7 KiB
Python
#!/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()
|