279 lines
12 KiB
Python
279 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""RF-DETR 整图批量检测野鸡,生成候选标注(初标)
|
|
|
|
- 整图等比缩放至最长边 INPUT_SIZE(700)送检测,**不做任何裁剪**,坐标按比例映射回原图
|
|
- 只保留 COCO bird 类;同目标重复框用 NMS 去重
|
|
- 藏身点:图像分析(植被密度/地形边缘/光线)生成 class 1 疑似藏匿位置
|
|
- 输出: YOLO txt (class 0 = 确认野鸡, class 1 = 疑似(低置信检测+藏身点))
|
|
-> datasets/labels/pheasant_rfdetr/
|
|
- 分析JSON -> datasets/analysis/pheasant_rfdetr/
|
|
- 预览图 -> datasets/previews/pheasant_rfdetr/
|
|
"""
|
|
import argparse
|
|
import base64
|
|
import glob
|
|
import json
|
|
import os
|
|
import urllib.request
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
SRC_IMAGES = os.path.join(BASE_DIR, "datasets/images/pheasant")
|
|
OUT_LABELS = os.path.join(BASE_DIR, "datasets/labels/pheasant_rfdetr")
|
|
OUT_JSON = os.path.join(BASE_DIR, "datasets/analysis/pheasant_rfdetr")
|
|
OUT_PREVIEW = os.path.join(BASE_DIR, "datasets/previews/pheasant_rfdetr")
|
|
|
|
ENDPOINT = "http://192.168.3.210:18080/v1/detection"
|
|
MODEL = "rfdetr-xlarge"
|
|
TIMEOUT = 300
|
|
THRESHOLD = 0.08
|
|
CONF_CONFIRMED = 0.2 # 置信度高于此为明确(class 0 红框),否则为疑似(class 1 蓝框)
|
|
# RF-DETR 置信度分布:强目标 0.3-0.9,0.05-0.14 多为噪声尾;
|
|
# THRESHOLD=0.08 砍掉纯噪声,CONF_CONFIRMED=0.2 只把强目标标红,宁多勿漏
|
|
INPUT_SIZE = 700 # 整图等比缩放后的最长边,不裁剪
|
|
NMS_IOU = 0.4
|
|
WORKERS = 4
|
|
KEEP_CLASSES = set() # 空 = 保留全部类别(RF-DETR 对小目标类别不稳定,定位一致即可)
|
|
|
|
|
|
IMG_MIME = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png"}
|
|
|
|
|
|
def detect(buf, mime):
|
|
body = json.dumps({"model": MODEL, "image": f"data:{mime};base64,{base64.b64encode(buf.tobytes()).decode()}",
|
|
"threshold": THRESHOLD}).encode()
|
|
opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
|
|
req = urllib.request.Request(ENDPOINT, data=body,
|
|
headers={"Content-Type": "application/json"}, method="POST")
|
|
with opener.open(req, timeout=TIMEOUT) as r:
|
|
resp = json.loads(r.read().decode())
|
|
return resp.get("detections", [])
|
|
|
|
|
|
def nms(boxes, iou_thr):
|
|
"""boxes: [(x, y, w, h, conf)], 按置信度降序去重"""
|
|
boxes = sorted(boxes, key=lambda b: -b[4])
|
|
keep = []
|
|
for b in boxes:
|
|
overlap = False
|
|
for k in keep:
|
|
ix = min(b[0] + b[2], k[0] + k[2]) - max(b[0], k[0])
|
|
iy = min(b[1] + b[3], k[1] + k[3]) - max(b[1], k[1])
|
|
if ix <= 0 or iy <= 0:
|
|
continue
|
|
inter = ix * iy
|
|
union = b[2] * b[3] + k[2] * k[3] - inter
|
|
if inter / union > iou_thr:
|
|
overlap = True
|
|
break
|
|
if not overlap:
|
|
keep.append(b)
|
|
return keep
|
|
|
|
|
|
def find_hiding_spots(img, H, W, n_spots=3):
|
|
"""图像分析疑似藏身点:植被密度/地形边缘(田埂、林田交界、沟渠)/光线调制。
|
|
|
|
野鸡习性:地面活动,藏身于茂密低矮遮蔽物,避开开阔裸地;
|
|
强光场景偏好阴影遮蔽,弱光场景偏好相对亮且有遮蔽的区域。
|
|
返回 [(x, y, w, h)] 像素坐标,框面积约占画面 2%~8%。
|
|
"""
|
|
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
|
s, v = hsv[..., 1], hsv[..., 2]
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
|
|
b = img[..., 0].astype(np.float32) / 255.0
|
|
g = img[..., 1].astype(np.float32) / 255.0
|
|
r = img[..., 2].astype(np.float32) / 255.0
|
|
veg = np.clip(g - np.maximum(r, b), 0, 1)
|
|
veg[s / 255.0 < 0.12] = 0
|
|
tex = np.abs(gray - cv2.GaussianBlur(gray, (0, 0), 3))
|
|
base = veg + tex
|
|
grad = cv2.magnitude(cv2.Sobel(base, cv2.CV_32F, 1, 0, ksize=3),
|
|
cv2.Sobel(base, cv2.CV_32F, 0, 1, ksize=3))
|
|
light = v / 255.0
|
|
mean_light = float(gray.mean())
|
|
if mean_light > 0.55:
|
|
light_w = 1.0 - np.abs(light - 0.25)
|
|
else:
|
|
light_w = 1.0 - np.abs(light - 0.6)
|
|
cover = (veg * 0.55 + tex * 0.30 + grad * 0.15) * (0.4 + 0.6 * light_w)
|
|
cover[light < 0.12] = 0
|
|
cover[light > 0.93] *= 0.15
|
|
cover = cv2.GaussianBlur(cover, (0, 0), 7)
|
|
windows = []
|
|
for area_frac in (0.02, 0.04, 0.07):
|
|
for aspect in (1.0, 1.6, 2.5):
|
|
bw = int(round((area_frac * aspect * H * W) ** 0.5))
|
|
bh = int(round((area_frac * H * W / aspect) ** 0.5))
|
|
bw, bh = min(bw, W), min(bh, H)
|
|
windows.append((cv2.boxFilter(cover, -1, (bw, bh)), bw, bh))
|
|
stack = np.stack([w[0] for w in windows])
|
|
best = stack.max(axis=0)
|
|
best_idx = stack.argmax(axis=0)
|
|
floor = max(float(best.max()) * 0.35, 0.02)
|
|
spots = []
|
|
for _ in range(n_spots):
|
|
idx = np.unravel_index(int(np.argmax(best)), best.shape)
|
|
y, x = idx
|
|
if best[y, x] < floor:
|
|
break
|
|
_, bw, bh = windows[int(best_idx[y, x])]
|
|
x0 = max(0, min(x - bw // 2, W - bw))
|
|
y0 = max(0, min(y - bh // 2, H - bh))
|
|
spots.append((x0, y0, bw, bh))
|
|
mx, my = max(bw, 64), max(bh, 64)
|
|
best[max(0, y - my):min(H, y + my), max(0, x - mx):min(W, x + mx)] = -1
|
|
return [(int(x), int(y), int(w), int(h)) for x, y, w, h in spots]
|
|
|
|
|
|
def suppress_overlap(lines):
|
|
"""class 1(疑似)框与任一 class 0(确认)框重叠时删除,只保留 class 0"""
|
|
cls0 = [l for l in lines if l.startswith("0 ")]
|
|
if not cls0:
|
|
return lines
|
|
boxes0 = [tuple(map(float, l.split())) for l in cls0]
|
|
keep1 = []
|
|
for l in lines:
|
|
if l.startswith("0 "):
|
|
continue
|
|
c, cx, cy, bw, bh = map(float, l.split())
|
|
x1, y1, x2, y2 = cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2
|
|
overlap = False
|
|
for kc, kx, ky, kw, kh in boxes0:
|
|
ix = min(x2, kx + kw / 2) - max(x1, kx - kw / 2)
|
|
iy = min(y2, ky + kh / 2) - max(y1, ky - kh / 2)
|
|
if ix > 0 and iy > 0:
|
|
overlap = True
|
|
break
|
|
if not overlap:
|
|
keep1.append(l)
|
|
return cls0 + keep1
|
|
|
|
|
|
def process(name):
|
|
path = os.path.join(SRC_IMAGES, name)
|
|
img = cv2.imread(path)
|
|
H, W = img.shape[:2]
|
|
ext = os.path.splitext(path)[1].lower()
|
|
mime = IMG_MIME.get(ext, "image/jpeg")
|
|
jpeg_params = [cv2.IMWRITE_JPEG_QUALITY, 100] if ext in (".jpg", ".jpeg") else []
|
|
# 整图等比缩放至最长边 INPUT_SIZE,不裁剪;检测坐标按比例映射回原图
|
|
scale = min(1.0, INPUT_SIZE / max(H, W))
|
|
if scale < 1.0:
|
|
small = cv2.resize(img, (round(W * scale), round(H * scale)),
|
|
interpolation=cv2.INTER_AREA)
|
|
else:
|
|
small = img
|
|
ok, buf = cv2.imencode(ext, small, jpeg_params)
|
|
boxes = []
|
|
for d in detect(buf, mime):
|
|
if not KEEP_CLASSES or d.get("class_name") in KEEP_CLASSES:
|
|
boxes.append((d["x"] / scale, d["y"] / scale,
|
|
d["width"] / scale, d["height"] / scale, d["confidence"]))
|
|
# 野鸡目标 20-100px(近大远小);排除大物体误检(田地纹理)与细条误检
|
|
boxes = [b for b in nms(boxes, NMS_IOU)
|
|
if 4 <= b[3] <= H * 0.06 and 4 <= b[2] <= W * 0.1
|
|
and 0.3 <= b[2] / b[3] <= 3.0]
|
|
hiding = find_hiding_spots(img, H, W)
|
|
return name, W, H, boxes, hiding
|
|
|
|
|
|
def reclassify():
|
|
"""不重新检测:按当前 CONF_CONFIRMED 从分析 JSON 重出 txt 标签与预览"""
|
|
total_boxes = total_empty = 0
|
|
for p in sorted(glob.glob(os.path.join(OUT_JSON, "*.json"))):
|
|
j = json.load(open(p, encoding="utf-8"))
|
|
W, H = j["size"]
|
|
stem = os.path.splitext(os.path.basename(p))[0]
|
|
lines = []
|
|
for b in j["candidates"]:
|
|
x, y, w, h, conf = b["x"], b["y"], b["width"], b["height"], b["confidence"]
|
|
cx, cy, bw, bh = (x + w / 2) / W, (y + h / 2) / H, w / W, h / H
|
|
cls = 0 if conf >= CONF_CONFIRMED else 1
|
|
lines.append(f"{cls} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
|
|
for b in j.get("hiding_spots", []):
|
|
x, y, w, h = b["x"], b["y"], b["width"], b["height"]
|
|
cx, cy, bw, bh = (x + w / 2) / W, (y + h / 2) / H, w / W, h / H
|
|
lines.append(f"1 {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
|
|
lines = suppress_overlap(lines)
|
|
with open(os.path.join(OUT_LABELS, f"{stem}.txt"), "w", encoding="utf-8") as f:
|
|
f.write("\n".join(lines) + ("\n" if lines else ""))
|
|
draw_preview(os.path.join(SRC_IMAGES, j["image"]), lines,
|
|
os.path.join(OUT_PREVIEW, f"{stem}.jpg"))
|
|
total_boxes += len(lines)
|
|
if not lines:
|
|
total_empty += 1
|
|
print(f"重分类完成: {total_boxes} 框, 空图 {total_empty} 张", flush=True)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="RF-DETR 切块放大批量检测野鸡候选框")
|
|
ap.add_argument("--names", type=str, default="", help="只处理指定编号,逗号分隔")
|
|
ap.add_argument("--reclassify", action="store_true",
|
|
help="不重新检测,按当前 CONF_CONFIRMED 从分析JSON重出标签与预览")
|
|
args = ap.parse_args()
|
|
|
|
if args.reclassify:
|
|
reclassify()
|
|
return
|
|
|
|
for d in (OUT_LABELS, OUT_JSON, OUT_PREVIEW):
|
|
os.makedirs(d, exist_ok=True)
|
|
names = sorted(p.name for p in os.scandir(SRC_IMAGES)
|
|
if p.name.startswith("pheasant_") and p.name.endswith(".jpg"))
|
|
if args.names:
|
|
want = {f"pheasant_{n}.jpg" for n in args.names.split(",")}
|
|
names = [n for n in names if n in want]
|
|
|
|
total_boxes = total_empty = 0
|
|
with ThreadPoolExecutor(max_workers=WORKERS) as pool:
|
|
futures = [pool.submit(process, n) for n in names]
|
|
for fut in as_completed(futures):
|
|
name, W, H, boxes, hiding = fut.result()
|
|
stem = os.path.splitext(name)[0]
|
|
lines = []
|
|
for x, y, w, h, conf in boxes:
|
|
cx, cy, bw, bh = (x + w / 2) / W, (y + h / 2) / H, w / W, h / H
|
|
cls = 0 if conf >= CONF_CONFIRMED else 1
|
|
lines.append(f"{cls} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
|
|
for x, y, w, h in hiding:
|
|
cx, cy, bw, bh = (x + w / 2) / W, (y + h / 2) / H, w / W, h / H
|
|
lines.append(f"1 {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
|
|
lines = suppress_overlap(lines)
|
|
with open(os.path.join(OUT_JSON, f"{stem}.json"), "w", encoding="utf-8") as f:
|
|
json.dump({"image": name, "size": [W, H],
|
|
"candidates": [{"x": b[0], "y": b[1], "width": b[2],
|
|
"height": b[3], "confidence": b[4]} for b in boxes],
|
|
"hiding_spots": [{"x": b[0], "y": b[1], "width": b[2],
|
|
"height": b[3]} for b in hiding]},
|
|
f, ensure_ascii=False, indent=1)
|
|
with open(os.path.join(OUT_LABELS, f"{stem}.txt"), "w", encoding="utf-8") as f:
|
|
f.write("\n".join(lines) + ("\n" if lines else ""))
|
|
draw_preview(os.path.join(SRC_IMAGES, name), lines,
|
|
os.path.join(OUT_PREVIEW, f"{stem}.jpg"))
|
|
total_boxes += len(lines)
|
|
if not lines:
|
|
total_empty += 1
|
|
print(f"{name}: 候选{len(lines)}", flush=True)
|
|
|
|
print(f"完成: {len(names)} 张, 候选框 {total_boxes}, 无候选 {total_empty} 张", flush=True)
|
|
|
|
|
|
def draw_preview(img_path, lines, out_path):
|
|
img = cv2.imread(img_path)
|
|
h, w = img.shape[:2]
|
|
for line in lines:
|
|
cls, cx, cy, bw, bh = map(float, line.split())
|
|
x1, y1 = int((cx - bw / 2) * w), int((cy - bh / 2) * h)
|
|
x2, y2 = int((cx + bw / 2) * w), int((cy + bh / 2) * h)
|
|
color = (0, 0, 255) if cls == 0 else (255, 0, 0)
|
|
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
|
|
cv2.imwrite(out_path, img)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|