red-future/observer
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

1

Browse Source
This commit is contained in:
admin
2026-08-20 18:00:47 +08:00
parent b9934c996d
commit 45d55028d7
17 changed files with 500 additions and 1520 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/main/assets/labels.txt
+1 -3
View File
@@ -1,4 +1,2 @@
pheasant
hare
dove
fish
cover
app/src/main/java/com/example/observer/detection/TFLiteDetector.kt
+2 -3
View File
@@ -32,7 +32,7 @@ class TFLiteDetector private constructor(
ByteBuffer.allocateDirect(1 * INPUT_SIZE * INPUT_SIZE * 3 * 4)
.order(ByteOrder.nativeOrder())
private val outputFloats = FloatArray((4 + NUM_CLASSES) * NUM_ANCHORS)
private val outputFloats = FloatArray((4 + labels.size) * NUM_ANCHORS)
override fun detect(bitmap: Bitmap): List<DetectionResult> {
preprocess(bitmap)
@@ -68,7 +68,7 @@ class TFLiteDetector private constructor(
val h = outputFloats[3 * NUM_ANCHORS + a]
var bestCls = 0
var bestScore = 0f
for (c in 0 until NUM_CLASSES) {
for (c in 0 until labels.size) {
val s = outputFloats[(4 + c) * NUM_ANCHORS + a]
if (s > bestScore) {
bestScore = s
@@ -96,7 +96,6 @@ class TFLiteDetector private constructor(
/** 动物类最低保留分数：低于此分不输出（低分候选由运动检测提升显示） */
const val MIN_SCORE = 0.20f
private const val NUM_CLASSES = 4
private const val NUM_ANCHORS = 8400
private const val IOU_THRESHOLD = 0.45f
private const val MAX_DETECTIONS = 20
app/src/main/java/com/example/observer/distance/DistanceEstimator.kt
+15 -10
View File
@@ -4,10 +4,11 @@ import android.content.Context
import android.hardware.camera2.CameraCharacteristics
import android.hardware.camera2.CameraManager
import kotlin.math.roundToInt
import kotlin.math.tan
/**
* 单目距离估计（针孔模型）：距离 = 焦距px × 参考体型 / 框高px。
* 误差预期 ±30%（5~50m）；鱼类目标受折射影响、生境区域按植被高度估算，均仅供参考。
* 误差预期 ±30%（5~50m）；生境区域按植被高度估算，均仅供参考。
*/
class DistanceEstimator(context: Context) {
@@ -19,27 +20,24 @@ class DistanceEstimator(context: Context) {
// 参考体型（米）
private val speciesSizeM = mapOf(
"pheasant" to 0.45f, // 身高
"hare" to 0.45f, // 身长
"dove" to 0.30f, // 体长
"fish" to 1.00f, // 典型可见体长（误差大）
"cover" to 0.50f, // 植被高度（水面区域误差大）
"cover" to 0.50f, // 植被高度（误差大）
)
fun estimate(
label: String,
boxHeightNorm: Float,
imageHeightPx: Int,
visibleHeightPx: Int,
cameraId: String?,
): Float? {
val realH = speciesSizeM[label] ?: return null
val boxH = boxHeightNorm * imageHeightPx
val boxH = boxHeightNorm * visibleHeightPx
if (boxH < 8f) return null // 过小目标不估算
val focalPx = focalPx(imageHeightPx, cameraId)
val focalPx = focalPx(visibleHeightPx, cameraId)
if (focalPx <= 0f) return null
return (focalPx * realH / boxH).roundToInt().toFloat()
}
/** focal_px = focal_mm × (imageHeightPx / sensorHeightMm) */
/** focal_px = focal_mm × (imageHeightPx / sensorHeightMm)；内参缺失时用视场角推算 */
private fun focalPx(imageHeightPx: Int, cameraId: String?): Float {
val key = "$cameraId:$imageHeightPx"
focalPxCache[key]?.let { return it }
@@ -52,7 +50,7 @@ class DistanceEstimator(context: Context) {
if (focalMm != null && sensor != null) {
focalMm * imageHeightPx / sensor.height
} else {
-1f
fovFallback(imageHeightPx, c)
}
} catch (e: Exception) {
-1f
@@ -60,4 +58,11 @@ class DistanceEstimator(context: Context) {
focalPxCache[key] = value
return value
}
/** focal_px = imageHeightPx / (2·tan(fovV/2)) */
private fun fovFallback(imageHeightPx: Int, c: CameraCharacteristics): Float {
val fovV = c.get(CameraCharacteristics.LENS_INFO_AVAILABLE_VERTICAL_VIEW_ANGLES)
?.firstOrNull() ?: return -1f
return (imageHeightPx / (2.0 * tan(Math.toRadians(fovV / 2.0)))).toFloat()
}
}
app/src/main/java/com/example/observer/overlay/DetectionOverlay.kt
-6
View File
@@ -18,17 +18,11 @@ import kotlin.math.roundToInt
private val speciesColors = mapOf(
"pheasant" to Color(0xFFE53935),
"hare" to Color(0xFF1E88E5),
"dove" to Color(0xFF8E24AA),
"fish" to Color(0xFF00ACC1),
"cover" to Color(0xFFFDD835),
)
private val speciesLabels = mapOf(
"pheasant" to "野鸡",
"hare" to "野兔",
"dove" to "斑鸠",
"fish" to "鱼",
"cover" to "疑似区域",
)
app/src/main/java/com/example/observer/ui/camera/CameraViewModel.kt
+4 -1
View File
@@ -108,7 +108,10 @@ class CameraViewModel(
val boosted = r.score < confThreshold &&
motionRegions.any { MotionAggregator.centerInRegion(r, it) }
val distance = if (showDistance) {
distanceEstimator.estimate(r.label, r.height, imageHeightPx, cameraId)
// 模型 CENTER_CROP 到方形输入, 归一化框高对应原图较短边(最大内接正方形)
distanceEstimator.estimate(
r.label, r.height, minOf(imageWidthPx, imageHeightPx), cameraId,
)
} else {
null
}
docs/01-技术方案.md
+4 -2
View File
@@ -191,6 +191,7 @@ flowchart LR
| 类别 ID | 英文标签 | 中文名 | 涵盖范围 | 参考体型（距离估计用） |
| --- | --- | --- | --- | --- |
| 0 | pheasant | 野鸡 | 环颈雉等雉类 | ≈ 0.45m（身高） |
| 1 | cover | 生境区域 | 草丛 / 灌木 / 水面等疑似生境（黄色虚线框预警，F08） | ≈ 0.5m（参考植被高度） |
### 6.2 模型选型对比
@@ -209,7 +210,8 @@ flowchart LR
| --- | --- |
| 数据量 | 1000~2000 张（首版可 500+ 起步，滚动补充） |
| 多样性 | 覆盖不同季节、晨昏/正午/逆光、远近距离、姿态、遮挡、背景（草丛/农田/林地/雪地） |
| 标注 | Roboflow 或 labelImg，YOLO 格式（class, cx, cy, w, h） |
| 标注 | 本地 LocalAI qwen3.8-9b 多模态自动标注（192.168.3.210:18080），先过滤再标注，YOLO 格式（class, cx, cy, w, h） |
| 生境标注 | 同一模型标注可疑度最高的 3 个具体藏身点（cover 类，黄色框，约占画面 2%~15%），综合植被密度 / 地形 / 光线判断，NMS 去重 + 尺寸过滤 |
| 数据增强 | Mosaic、MixUp、HSV 扰动、随机翻转、随机缩放裁剪 |
| 数据划分 | train 80% / val 10% / test 10% |
| 负样本 | 补充无目标场景图，控制误检 |
@@ -245,7 +247,7 @@ flowchart LR
| 检测阈值（默认） | confidence ≥ 0.40，IoU-NMS = 0.45 |
| 单帧推理延迟 | ≤ 80ms（中端机，320 输入，GPU） |
| 模型体积 | fp16 ≤ 8MB，int8 ≤ 4MB |
| 识别类别数 | 1 类（野鸡） |
| 识别类别数 | 2 类（野鸡 + 生境区域 cover） |
| 距离标注 | "约 X m"，5~50m 误差 ≤ ±30%，计算开销可忽略 |
---
docs/02-项目功能文档.md
+1 -1
View File
@@ -55,7 +55,7 @@ Observer
| --- | --- | --- | --- |
| F01 | 相机实时预览 | 全屏取景，默认后置，支持自动对焦、双指缩放、点击对焦、前后摄切换 | P0 |
| F02 | 目标实时检测 | 对野鸡实时检测框选 | P0 |
| F03 | 识别结果展示与距离标注 | 检测框 + 类别名 + 置信度百分比 + 大致距离（约 X m）；四类目标使用不同颜色边框 | P0 |
| F03 | 识别结果展示与距离标注 | 检测框 + 类别名 + 置信度百分比 + 大致距离（约 X m）；野鸡红色实线框，生境区域黄色虚线框 | P0 |
| F04 | 检测提醒 | 检测到目标时震动 / 提示音；支持开关与 10s 防重复 | P1 |
| F05 | 设置 | 置信度阈值（0.2~0.7）、检测频率（连续 / 标准 / 省电）、提醒开关、低光增强开关 | P1 |
| F06 | 低光增强 | 低光环境下自动亮度 / 对比度增强，提高识别率 | P2 |
training/auto_label.py
+394 -178
View File
@@ -1,202 +1,418 @@
"""用 LocalAI qwen3.5-9b 多模态模型自动标注图片。
#!/usr/bin/env python3
"""qwen3.8-9b (LocalAI 192.168.3.210:18080) 多模态自动标注 → YOLO 格式
用法：python auto_label.py
输入：datasets/images/<class>/*.jpg
输出：datasets/dataset/{train,val}/{images,labels} (YOLO 格式) + data.yaml
类别: pheasant(0) hare(1) dove(2) fish(3) cover(4, 生境区域)
标注两路:
- 动物框(红): 描述门控(2 采样, 均明确"无鸟"才判负) + 定位共识(3 采样, 中位数),
样本分歧大(两两 IoU < 0.5) 的图像进复核清单
- 生境框(黄, cover): 植被覆盖地带识别 + NMS 去重
经验(2026-08-20 实测):
- 模型对"有无野鸡"的布尔判断有"有"偏置, 全图 3/3 答"有", 不可用
- 描述模式能区分大部分正负样本(负样本会明确说"没有鸟类或动物"), 但个别图多次描述矛盾
- 模型对负样本图像会稳定幻觉出小框, 所以不能只用定位采样判断有无
- 定位提示词强调"完整包住/从头到尾"比"贴合身体"更准(原提示词框偏小偏右)
- 隐藏目标(只露头/尾巴)会被描述门控误判负, 判负后需定位保险(42 图实测)
- 模型可能多采样一致地错(系统性偏差, 41/44 图实测), 需双提示词交叉验证
- 框面积 <1% 多为局部误检, 但真目标(只露头)也可能很小, 一律进复核
用法:
venv/bin/python auto_label.py --input datasets/images --output datasets/labels
venv/bin/python auto_label.py --input datasets/images --output datasets/labels --habitat
venv/bin/python auto_label.py --input datasets/images --output datasets/labels --animal
venv/bin/python auto_label.py --input datasets/images --output datasets/labels --dry-run
"""
import os
import json
import re
import argparse
import base64
import shutil
import io
import json
import os
import re
import sys
import time
import urllib.request
from PIL import Image, ImageDraw
from pathlib import Path
# LocalAI 服务器配置
LOCAL_AI_URL = "http://192.168.3.210:18080/v1/chat/completions"
MODEL_NAME = "qwen3.5-9b"
from PIL import Image
# 标注配置
KEEP_PER_CLASS = 100
MIN_CONF = 0.5 # 自动标注最低置信度
MIN_BOX = 0.03 # 框面积占比下限
API_URL = os.environ.get("OBSERVER_AI_URL", "http://192.168.3.210:18080")
MODEL = os.environ.get("OBSERVER_AI_MODEL", "qwen3.8-9b")
MAX_SIDE = 800 # 预处理尺寸, 对齐 LocalAI 聊天界面(100K 上下文下约 600~700 tokens/图)
JPEG_QUALITY = 85
TIMEOUT = 300
MAX_RETRIES = 4
RETRY_BACKOFF = 3 # 秒, 指数退避
DESC_SAMPLES = 2 # 描述门控采样数
BOX_SAMPLES = 3 # 定位采样数
CROSS_SAMPLES = 2 # 交叉提示词采样数(双提示词交叉验证)
CONSENSUS_IOU = 0.5 # 定位共识: 低于此值=采样分歧大, 进复核
CONSENSUS_OK_IOU = 0.7 # 定位共识: 高于此值才算可信, 中间段=勉强一致, 进复核
CROSS_IOU = 0.5 # 双提示词共识框 IoU 低于此值=交叉不一致, 进复核
CENTER_AGREE = 0.08 # 小框对 IoU 敏感, 中心距离不超过此值也视为一致
MIN_ANIMAL_AREA = 0.002 # 动物框面积低于 0.2%=极小疑似误检, 进复核
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")
# 项目范围: 仅野鸡 + 生境区域 cover(2026-08-20 确认, 不含野兔/斑鸠/鱼)
CLASSES = ["pheasant", "cover"]
CLASS_CN = {"pheasant": "野鸡（环颈雉）"}
CLASSES = ["pheasant"]
PROMPTS = {
"pheasant": "pheasant",
}
PHEASANT_FEATURES = ("野鸡（环颈雉）识别特征（雄性个体）：黑脑袋、红色脸颊、白色颈环、细长尾羽；"
"体型似鸡，站立或行走姿态")
# 标注提示词模板
ANNOTATION_PROMPT = """你是一个目标检测助手。请检测图片中【{target}】动物。
DESC_PROMPT = ("请客观描述这张图片的内容（50字以内）：画面里有什么？"
"是否有任何鸟类或动物？如果看到鸟类或动物，请明确说出来。"
f"注意：{PHEASANT_FEATURES}。"
"野鸡可能藏在草丛灌木中，只露出头部或尾巴，这样也算看到野鸡，要明确说出来。")
输出格式（严格只输出这一行，不要其他文字）：
pheasant 0.92 0.1 0.2 0.3 0.4
ANIMAL_PROMPT_TMPL = ("图片中有{cn}。注意识别特征：{feat}。"
"给出完整包住{cn}的边界框，尽量贴合，不要切掉身体任何部分（含尾巴）。"
"如果{cn}被遮挡、只露出部分（头部/尾巴/局部身体），也要框住可见部分，"
"不要因为没有全身就漏标。"
"每只{cn}一个框。如果图片中没有{cn}，boxes 输出空数组。"
"只输出JSON: {{\"boxes\": [[ymin, xmin, ymax, xmax], ...]}}。"
"坐标必须是0到1之间的归一化小数，禁止输出像素坐标")
字段说明：
- 第一个数字：置信度（0-1）
- 后四个数字：边界框坐标 x1 y1 x2 y2（归一化 0-1）
# 交叉验证提示词: 结构与主提示词不同, 用于拆穿"多采样一致的错"(系统性偏差)
ANIMAL_PROMPT2_TMPL = ("画面中可能有{cn}。先在脑海中定位：{cn}的头、身体、尾巴各在什么位置？"
"然后给出完整包住{cn}的边界框；如果只露出部分，框住可见部分。"
"只输出JSON: {{\"boxes\": [[ymin, xmin, ymax, xmax], ...]}}，没有则空数组。"
"坐标0~1，禁止像素坐标")
如果没找到目标，只输出：NOT_FOUND
"""
HABITAT_PROMPT = """你是野生动物观察辅助工具。请找出野鸡最可能藏身或出现的具体位置。
规则：
- 只标注【具体的可疑位置】（如浓密草丛、植被边缘、沟渠边、倒木旁），不要框大片地带
- 每个位置一个小框，框住该可疑处即可，框的面积适中（约占画面 2%~15%）
- 按可疑度从高到低排列，最多 3 个
- 综合植被密度、地形、光线判断：植被浓密、能藏身、光线被遮挡处优先
- 不要把裸露地面、道路、天空框进去
- 只输出JSON: {"regions": [[ymin, xmin, ymax, xmax], ...]}，没有则 {"regions": []}"""
NO_BIRD_RE = re.compile(r"(没有|未发现|未看到|没有任何|看不到|不见).{0,10}(鸟类|鸟|动物|野鸡|雉)")
opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
def encode_image_to_base64(path: str) -> str:
"""将图片路径转换为 base64 数据"""
with open(path, "rb") as f:
return base64.b64encode(f.read()).decode("utf-8")
def preprocess_image(path: Path) -> tuple[bytes, int, int]:
"""返回 (jpeg字节, 缩放后宽, 缩放后高)"""
with Image.open(path) as im:
im = im.convert("RGB")
w, h = im.size
if max(w, h) > MAX_SIDE:
scale = MAX_SIDE / max(w, h)
w, h = round(w * scale), round(h * scale)
im = im.resize((w, h), Image.LANCZOS)
buf = io.BytesIO()
im.save(buf, format="JPEG", quality=JPEG_QUALITY)
return buf.getvalue(), w, h
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,
def call_vision(jpeg: bytes, prompt: str, max_tokens: int = 400,
temperature: float = 0.1) -> str:
b64 = base64.b64encode(jpeg).decode()
payload = {
"model": MODEL,
"messages": [{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{"type": "image_url",
"image_url": {"url": "data:image/jpeg;base64," + b64}},
],
}],
"temperature": temperature,
"max_tokens": max_tokens,
}
req = urllib.request.Request(
f"{API_URL}/v1/chat/completions",
data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"},
)
last_err = None
for attempt in range(MAX_RETRIES):
try:
with opener.open(req, timeout=TIMEOUT) as resp:
data = json.load(resp)
return data["choices"][0]["message"]["content"]
except Exception as e:
last_err = e
time.sleep(RETRY_BACKOFF * (2 ** attempt))
raise RuntimeError(f"请求失败: {last_err}")
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}
def parse_json_box_list(content: str, key: str, disp_w: int, disp_h: int):
"""解析 {"key": [[ymin,xmin,ymax,xmax],...]} 返回 [(ymin,xmin,ymax,xmax)]
模型偶发输出像素坐标(相对缩放后输入图), 自动按实际尺寸归一化
"""
text = re.sub(r"^```(?:json)?\s*|\s*```$", "", content.strip(), flags=re.MULTILINE)
m = re.search(r"\{.*\}", text, re.DOTALL)
if not m:
raise ValueError(f"响应中无 JSON: {content[:200]}")
data = json.loads(m.group(0))
boxes = []
for b in data.get(key, []):
ymin, xmin, ymax, xmax = (float(v) for v in b)
if max(ymin, xmin, ymax, xmax) > 1:
# 模型偶发输出像素坐标(相对缩放后输入图), 超出部分裁剪到边界
xmin, xmax = xmin / disp_w, xmax / disp_w
ymin, ymax = ymin / disp_h, ymax / disp_h
xmin, xmax = min(max(xmin, 0), 1), min(max(xmax, 0), 1)
ymin, ymax = min(max(ymin, 0), 1), min(max(ymax, 0), 1)
# 模型偶发输出坐标序颠倒的框
if xmin > xmax:
xmin, xmax = xmax, xmin
if ymin > ymax:
ymin, ymax = ymax, ymin
if xmin >= xmax or ymin >= ymax:
continue # 裁剪后退化(整体越界)的框直接丢弃
boxes.append((ymin, xmin, ymax, xmax))
return boxes
def iou(a, b):
ymin1, xmin1, ymax1, xmax1 = a
ymin2, xmin2, ymax2, xmax2 = b
iw = min(xmax1, xmax2) - max(xmin1, xmin2)
ih = min(ymax1, ymax2) - max(ymin1, ymin2)
if iw <= 0 or ih <= 0:
return 0.0
inter = iw * ih
union = (xmax1 - xmin1) * (ymax1 - ymin1) + (xmax2 - xmin2) * (ymax2 - ymin2) - inter
return inter / union if union > 0 else 0.0
def nms(boxes, thr=0.5):
"""按面积降序贪心去重(用于生境多框)"""
kept = []
for b in sorted(boxes, key=lambda x: (x[2] - x[0]) * (x[3] - x[1]), reverse=True):
if all(iou(b, k) < thr for k in kept):
kept.append(b)
return kept
def boxes_agree(a, b, iou_thr=CROSS_IOU, center_thr=CENTER_AGREE):
"""小框 IoU 敏感(同位置小框 IoU 可能很低), 中心距离足够近也视为一致"""
if iou(a, b) >= iou_thr:
return True
ca = ((a[1] + a[3]) / 2, (a[0] + a[2]) / 2)
cb = ((b[1] + b[3]) / 2, (b[0] + b[2]) / 2)
return max(abs(ca[0] - cb[0]), abs(ca[1] - cb[1])) <= center_thr
def median_box(boxes):
"""按坐标分量的中位数合成框(定位共识)"""
return tuple(
sorted(v)[len(v) // 2] for v in zip(*boxes)
)
def consensus_box(nonempty: list[list]) -> tuple[float, tuple]:
"""投票法共识: 选出其他采样中同意率最高的框, 返回 (同意率, 框)
同意 = boxes_agree(IoU 达标或中心距离足够近), 对小框友好(IoU 均值法会误伤小框)
"""
if len(nonempty) == 1:
return 0.0, nonempty[0][0]
best = None
for b in nonempty[0]:
agree = [any(boxes_agree(b, o) for o in s) for s in nonempty[1:]]
score = sum(agree) / len(agree)
if best is None or score > best[0]:
best = (score, b)
return best
def annotate_animal(jpeg: bytes, cls: str, disp_w: int, disp_h: int) -> tuple[bool, list, str, str]:
"""返回 (检出, 框列表[(ymin,xmin,ymax,xmax)], 状态 ok|review, 复核原因)"""
cn = CLASS_CN[cls]
prompt = ANIMAL_PROMPT_TMPL.format(cn=cn, feat=PHEASANT_FEATURES)
prompt2 = ANIMAL_PROMPT2_TMPL.format(cn=cn)
# 1. 描述门控: 2 采样, 均明确无鸟才判负
descs = [call_vision(jpeg, DESC_PROMPT, max_tokens=150) for _ in range(DESC_SAMPLES)]
no_birds = [bool(NO_BIRD_RE.search(d)) for d in descs]
gate_blocked = False
if all(no_birds):
# 门控保险: 判负后仍各跑 1 次主/交叉定位, 任一有框则门控不可信
ins1 = parse_json_box_list(
call_vision(jpeg, prompt, temperature=0.3), "boxes", disp_w, disp_h)
ins2 = parse_json_box_list(
call_vision(jpeg, prompt2, temperature=0.3), "boxes", disp_w, disp_h)
if not ins1 and not ins2:
return False, [], "ok", ""
gate_blocked = True
# 2. 定位: 主提示词 3 采样 + 交叉提示词 2 采样
samples = [parse_json_box_list(call_vision(jpeg, prompt, temperature=0.3),
"boxes", disp_w, disp_h) for _ in range(BOX_SAMPLES)]
cross = [parse_json_box_list(call_vision(jpeg, prompt2, temperature=0.3),
"boxes", disp_w, disp_h) for _ in range(CROSS_SAMPLES)]
nonempty = [s for s in samples if s]
if not nonempty:
cn2 = [s for s in cross if s]
if cn2 and (gate_blocked or len(cn2) >= CROSS_SAMPLES):
_, b2 = consensus_box(cn2)
return True, [b2], "review", "主提示词未检出但交叉提示词有框"
return False, [], "ok", ""
score, box = consensus_box(nonempty)
if len(nonempty) < BOX_SAMPLES:
return True, [box], "review", "定位采样检出不一致"
# 3. 双提示词交叉验证: 拆穿多采样一致的系统性偏差
cn2 = [s for s in cross if s]
if cn2:
_, b2 = consensus_box(cn2)
if not boxes_agree(box, b2):
return True, [box], "review", "双提示词交叉不一致"
# 4. 共识分级: <0.5 分歧大, 0.5~0.7 勉强一致, 均进复核
if score < CONSENSUS_IOU:
return True, [box], "review", "定位采样分歧大"
if score < CONSENSUS_OK_IOU:
return True, [box], "review", "定位采样勉强一致"
# 5. 小框复核: 面积 <1% 疑似局部误检
if (box[2] - box[0]) * (box[3] - box[1]) < MIN_ANIMAL_AREA:
return True, [box], "review", "框过小,疑似局部或误检"
if gate_blocked:
return True, [box], "review", "门控判负但定位有框"
return True, [box], "ok", ""
MAX_HABITAT_BOXES = 3 # 生境只保留可疑度最高的 3 个位置
MIN_HABITAT_AREA = 0.01 # 太小(碎点, <1% 面积)对训练无意义
MAX_HABITAT_AREA = 0.25 # 太大(大面积地带)不是"可疑点"
MAX_HABITAT_SIDE = 0.7 # 全宽/全高条带排除
def annotate_habitat(jpeg: bytes, disp_w: int, disp_h: int) -> list:
content = call_vision(jpeg, HABITAT_PROMPT, max_tokens=400)
boxes = parse_json_box_list(content, "regions", disp_w, disp_h)
boxes = nms(boxes, 0.5) # 去重(模型偶发输出重复框)
filtered = []
for ymin, xmin, ymax, xmax in boxes:
w, h = xmax - xmin, ymax - ymin
area = w * h
if area < MIN_HABITAT_AREA or area > MAX_HABITAT_AREA:
continue
if w > MAX_HABITAT_SIDE or h > MAX_HABITAT_SIDE:
continue
filtered.append((ymin, xmin, ymax, xmax))
return filtered[:MAX_HABITAT_BOXES]
def to_yolo(class_id: int, boxes) -> str:
lines = []
for ymin, xmin, ymax, xmax in boxes:
cx = min(max((xmin + xmax) / 2, 0), 1)
cy = min(max((ymin + ymax) / 2, 0), 1)
w = min(max(xmax - xmin, 0), 1)
h = min(max(ymax - ymin, 0), 1)
lines.append(f"{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}")
return "\n".join(lines)
def main():
ap = argparse.ArgumentParser(description="qwen 自动标注 → YOLO(动物+生境)")
ap.add_argument("--input", default="datasets/images")
ap.add_argument("--output", default="datasets/labels")
ap.add_argument("--review-file", default="datasets/review.txt")
ap.add_argument("--mode", choices=["both", "animal", "habitat"], default="both")
ap.add_argument("--dry-run", action="store_true")
ap.add_argument("--limit", type=int, default=0)
ap.add_argument("--force", action="store_true")
args = ap.parse_args()
img_dir = Path(args.input)
out_dir = Path(args.output)
review_file = Path(args.review_file)
if not img_dir.is_dir():
sys.exit(f"输入目录不存在: {img_dir}")
class_ids = {name: i for i, name in enumerate(CLASSES)}
images = sorted(p for p in img_dir.rglob("*")
if p.suffix.lower() in (".jpg", ".jpeg", ".png", ".webp", ".bmp"))
todo = []
for p in images:
rel = p.relative_to(img_dir)
cls = rel.parts[0] if len(rel.parts) > 1 else ""
if cls not in class_ids or cls == "cover":
print(f"跳过: 未知类别目录 {rel}", file=sys.stderr)
continue
label_path = out_dir / rel.with_suffix(".txt")
if not args.force and label_path.exists():
existing = label_path.read_text().splitlines()
if args.mode == "both":
continue # 完整模式,已有结果即跳过
if args.mode == "animal" and any(
not l.startswith(f"{class_ids['cover']} ") for l in existing if l):
continue # 已有动物标注
if args.mode == "habitat" and any(
l.startswith(f"{class_ids['cover']} ") for l in existing if l):
continue # 已有生境标注
todo.append((p, cls, label_path))
if args.limit > 0:
todo = todo[:args.limit]
if not todo:
print("没有待标注的图片。")
return
print(f"待标注 {len(todo)} 张 [模式: {args.mode}]")
if args.dry_run:
return
ok = no_target = fail = review = 0
t0 = time.time()
for i, (img_path, cls, label_path) in enumerate(todo, 1):
rel = img_path.relative_to(img_dir)
label_path.parent.mkdir(parents=True, exist_ok=True)
try:
jpeg, disp_w, disp_h = preprocess_image(img_path)
cover_id = class_ids["cover"]
status = "ok"
# 单模式重跑时保留另一类已有标注
existing = (label_path.read_text().splitlines()
if label_path.exists() else [])
animal_lines, cover_lines = [], []
if args.mode in ("both", "animal"):
found, boxes, status, reason = annotate_animal(jpeg, cls, disp_w, disp_h)
animal_lines = to_yolo(class_ids[cls], boxes).splitlines()
if args.mode in ("both", "habitat"):
hboxes = annotate_habitat(jpeg, disp_w, disp_h)
cover_lines = to_yolo(cover_id, hboxes).splitlines()
if args.mode == "animal":
cover_lines = [l for l in existing
if l.startswith(f"{cover_id} ")]
if args.mode == "habitat":
animal_lines = [l for l in existing
if not l.startswith(f"{cover_id} ")]
lines = animal_lines + cover_lines
label_path.write_text("\n".join(lines))
if status == "review":
review += 1
with review_file.open("a") as f:
f.write(f"{rel}\t{cls}\t{reason}\n")
tag = f"复核[{reason[:10]}]"
elif lines:
ok += 1
tag = "OK"
else:
return {"success": False, "error": f"HTTP {resp.status}"}
except Exception as e:
return {"success": False, "error": str(e)}
no_target += 1
tag = "无目标"
print(f"[{i}/{len(todo)}] {rel}: {tag} ({len(lines)} 行)")
except Exception as e:
fail += 1
with review_file.open("a") as f:
f.write(f"{rel}\t{cls}\t失败: {e}\n")
print(f"[{i}/{len(todo)}] {rel}: 失败 - {e}", file=sys.stderr)
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)
print(f"\n完成: 有目标 {ok} 张, 无目标 {no_target} 张, 需复核 {review} 张, "
f"失败 {fail} 张, 耗时 {time.time() - t0:.0f}s")
if review:
print(f"复核清单: {review_file}")
if __name__ == "__main__":
sys.exit(main())
main()
training/auto_label_for_training.py
-120
View File
@@ -1,120 +0,0 @@
#!/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()
training/auto_label_v2.py
-234
View File
@@ -1,234 +0,0 @@
#!/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())
training/clip_generate_labels.py
-207
View File
@@ -1,207 +0,0 @@
#!/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()
training/download_data.py
-300
View File
@@ -1,300 +0,0 @@
"""从 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())
training/filter_images.py
-112
View File
@@ -1,112 +0,0 @@
"""用 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())
training/filter_images_clip.py
-130
View File
@@ -1,130 +0,0 @@
"""用本地 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())
training/train_yolov8n.py
-59
View File
@@ -1,59 +0,0 @@
#!/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()
training/visualize_labels.py
+79
View File
@@ -0,0 +1,79 @@
#!/usr/bin/env python3
"""标注预览图生成: 动物红色实线框 + 生境 cover 黄色虚线框
用法:
venv/bin/python visualize_labels.py --input datasets/images --labels datasets/labels --output datasets/pheasant_label
"""
import argparse
from pathlib import Path
from PIL import Image, ImageDraw
CLASSES = ["pheasant", "cover"]
LABELS = {"pheasant": "野鸡", "cover": "疑似区域"}
def draw_dashed(draw, box, outline, width, dash=12, gap=8):
x1, y1, x2, y2 = box
for (ax, ay, bx, by) in [(x1, y1, x2, y1), (x2, y1, x2, y2),
(x2, y2, x1, y2), (x1, y2, x1, y1)]:
length = max(abs(bx - ax), abs(by - ay))
steps = max(int(length / (dash + gap)), 1)
for i in range(steps):
s = i / steps
e = min((i * (dash + gap) + dash) / length, 1.0)
if e <= s:
continue
draw.line([ax + (bx - ax) * s, ay + (by - ay) * s,
ax + (bx - ax) * e, ay + (by - ay) * e],
fill=outline, width=width)
def main():
ap = argparse.ArgumentParser(description="标注预览图生成")
ap.add_argument("--input", default="datasets/images")
ap.add_argument("--labels", default="datasets/labels")
ap.add_argument("--output", default="datasets/pheasant_label")
args = ap.parse_args()
img_dir = Path(args.input)
lab_dir = Path(args.labels)
out_dir = Path(args.output)
out_dir.mkdir(parents=True, exist_ok=True)
count = 0
for img_path in sorted(img_dir.rglob("*")):
if img_path.suffix.lower() not in (".jpg", ".jpeg", ".png", ".webp", ".bmp"):
continue
rel = img_path.relative_to(img_dir)
lab_path = lab_dir / rel.with_suffix(".txt")
if not lab_path.exists():
continue
im = Image.open(img_path).convert("RGB")
W, H = im.size
d = ImageDraw.Draw(im)
for line in lab_path.read_text().splitlines():
parts = line.split()
if not parts:
continue
cid, cx, cy, w, h = map(float, parts)
box = ((cx - w / 2) * W, (cy - h / 2) * H,
(cx + w / 2) * W, (cy + h / 2) * H)
if int(cid) == 1: # cover 生境: 黄色虚线
draw_dashed(d, box, (255, 200, 0), width=5)
d.text((box[0] + 6, max(box[1] - 28, 4)), "疑似区域",
fill=(255, 200, 0))
else: # 动物: 红色实线
d.rectangle(box, outline=(255, 0, 0), width=6)
label = LABELS.get(CLASSES[int(cid)], CLASSES[int(cid)])
d.text((box[0] + 6, max(box[1] - 28, 4)), label,
fill=(255, 0, 0))
out = out_dir / (img_path.stem + ".jpg")
im.save(out, quality=92)
count += 1
print(out.name)
print(f"共生成 {count} 张")
if __name__ == "__main__":
main()
training/zero_shot_detection.py
-154
View File
@@ -1,154 +0,0 @@
#!/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()