66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
import cv2
|
||
import os
|
||
import shutil
|
||
from ultralytics import YOLO
|
||
|
||
# ====================== 配置 ======================
|
||
MODEL_PATH = 'point.pt'
|
||
IMAGE_SOURCE_DIR = '/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/ailaipoint'
|
||
|
||
OUTPUT_ROOT = '/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/ailaipoint/train_split'
|
||
OUTPUT_DIR_0 = os.path.join(OUTPUT_ROOT, '0') # 无目标 / conf=0
|
||
OUTPUT_DIR_1 = os.path.join(OUTPUT_ROOT, '1') # 0 < conf < 0.5
|
||
OUTPUT_DIR_2 = os.path.join(OUTPUT_ROOT, '2') # conf >= 0.5
|
||
|
||
for d in [OUTPUT_DIR_0, OUTPUT_DIR_1, OUTPUT_DIR_2]:
|
||
os.makedirs(d, exist_ok=True)
|
||
|
||
IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'}
|
||
|
||
# ====================== 主程序 ======================
|
||
if __name__ == "__main__":
|
||
print("🚀 bbox 置信度分桶（移动原图，含无目标图像）")
|
||
|
||
model = YOLO(MODEL_PATH)
|
||
|
||
image_files = [
|
||
f for f in os.listdir(IMAGE_SOURCE_DIR)
|
||
if os.path.splitext(f.lower())[1] in IMG_EXTENSIONS
|
||
]
|
||
|
||
print(f"📸 找到图片 {len(image_files)} 张")
|
||
|
||
for img_name in image_files:
|
||
src_path = os.path.join(IMAGE_SOURCE_DIR, img_name)
|
||
|
||
img = cv2.imread(src_path)
|
||
if img is None:
|
||
continue
|
||
|
||
results = model(img, verbose=False)
|
||
|
||
# ====================== 关键修复点 ======================
|
||
if not results or results[0].boxes is None or len(results[0].boxes.conf) == 0:
|
||
# 没有任何检测框 → 当作 conf = 0
|
||
bbox_conf = 0.0
|
||
else:
|
||
# 有检测框 → 取第一个（或最大 conf）
|
||
bbox_conf = float(results[0].boxes.conf[0].cpu().item())
|
||
|
||
# ====================== 分桶 ======================
|
||
if bbox_conf == 0:
|
||
dst_dir = OUTPUT_DIR_0
|
||
elif bbox_conf < 0.5:
|
||
dst_dir = OUTPUT_DIR_1
|
||
else:
|
||
dst_dir = OUTPUT_DIR_2
|
||
|
||
dst_path = os.path.join(dst_dir, img_name)
|
||
|
||
# ====================== 移动文件 ======================
|
||
shutil.move(src_path, dst_path)
|
||
|
||
print(f"{img_name} -> conf={bbox_conf:.3f} -> {os.path.basename(dst_dir)}")
|
||
|
||
print("✅ 完成（含无目标图片）")
|