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