huangxin/zjsh_yolov11
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

最新推送

Browse Source
This commit is contained in:
琉璃月光
2026-03-10 13:58:21 +08:00
parent 032479f558
commit eb16eeada3
97 changed files with 16865 additions and 670 deletions
Download Patch File Download Diff File Expand all files Collapse all files

132
angle_base_obb/angle_caculate_file.py
View File

@ -6,19 +6,28 @@ from ultralytics import YOLO
IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}
def process_obb_images(model_path, image_dir, output_dir="./inference_results", conf_thresh=0.15, imgsz=640):
"""
批量处理图像的 OBB 推理，计算每张图像检测目标的主方向和夹角。
def draw_direction(img, cx, cy, angle_deg, length=80, color=(0, 255, 0)):
"""画主方向箭头"""
theta = np.radians(angle_deg)
x2 = int(cx + length * np.cos(theta))
y2 = int(cy + length * np.sin(theta))
cv2.arrowedLine(
img,
(int(cx), int(cy)),
(x2, y2),
color,
2,
tipLength=0.2
)
输入：
model_path: YOLO 权重路径
image_dir: 图像文件夹路径
output_dir: 输出结果保存路径
conf_thresh: 置信度阈值
imgsz: 输入图像大小
输出：
results_dict: {image_filename: {'angles_deg': [...], 'pairwise_angles_deg': [...]}}
"""
def process_obb_images(
model_path,
image_dir,
output_dir="./inference_results",
conf_thresh=0.15,
imgsz=640
):
os.makedirs(output_dir, exist_ok=True)
results_dict = {}
@ -26,8 +35,11 @@ def process_obb_images(model_path, image_dir, output_dir="./inference_results",
model = YOLO(model_path)
print("✅ 模型加载完成")
# 获取图像文件
image_files = [f for f in os.listdir(image_dir) if os.path.splitext(f.lower())[1] in IMG_EXTENSIONS]
image_files = [
f for f in os.listdir(image_dir)
if os.path.splitext(f.lower())[1] in IMG_EXTENSIONS
]
if not image_files:
print(f"❌ 未找到图像文件：{image_dir}")
return results_dict
@ -40,63 +52,101 @@ def process_obb_images(model_path, image_dir, output_dir="./inference_results",
img = cv2.imread(img_path)
if img is None:
print(f"❌ 跳过：无法读取图像 {img_path}")
print("❌ 读取失败，跳过")
continue
# 推理 OBB
results = model(img, save=False, imgsz=imgsz, conf=conf_thresh, mode='obb')
# ---------- OBB 推理 ----------
results = model(img, save=False, imgsz=imgsz, conf=conf_thresh, mode="obb")
result = results[0]
annotated_img = result.plot()
# 保存可视化
save_path = os.path.join(output_dir, "detected_" + img_filename)
cv2.imwrite(save_path, annotated_img)
print(f"✅ 推理结果已保存至: {save_path}")
# 提取旋转角
boxes = result.obb
angles_deg = []
centers = []
if boxes is None or len(boxes) == 0:
print("❌ 该图像中未检测到任何目标")
print("❌ 未检测到目标")
else:
for i, box in enumerate(boxes):
cls = int(box.cls.cpu().numpy()[0])
conf = box.conf.cpu().numpy()[0]
cx, cy, w, h, r_rad = box.xywhr.cpu().numpy()[0]
direction = r_rad if w >= h else r_rad + np.pi / 2
direction = direction % np.pi
angle_deg = np.degrees(direction)
angles_deg.append(angle_deg)
print(f" Box {i + 1}: Class={cls}, Conf={conf:.3f}, 主方向={angle_deg:.2f}°")
# 两两夹角
angles_deg.append(angle_deg)
centers.append((int(cx), int(cy)))
print(f" Box {i + 1} 主方向: {angle_deg:.2f} deg")
# 主方向可视化
draw_direction(annotated_img, cx, cy, angle_deg)
cv2.circle(annotated_img, (int(cx), int(cy)), 4, (0, 0, 255), -1)
# ---------- 两两夹角 ----------
pairwise_angles_deg = []
if len(angles_deg) >= 2:
for i in range(len(angles_deg)):
for j in range(i + 1, len(angles_deg)):
diff_rad = abs(np.radians(angles_deg[i]) - np.radians(angles_deg[j]))
diff_rad = abs(
np.radians(angles_deg[i]) -
np.radians(angles_deg[j])
)
min_diff_rad = min(diff_rad, np.pi - diff_rad)
pairwise_angles_deg.append(np.degrees(min_diff_rad))
print(f" Box {i + 1} 与 Box {j + 1} 夹角: {np.degrees(min_diff_rad):.2f}°")
angle_ij = np.degrees(min_diff_rad)
pairwise_angles_deg.append(angle_ij)
print(
f" Box {i + 1} 与 Box {j + 1} 夹角: {angle_ij:.2f} deg"
)
# ---------- 右上角粗体 angle ----------
if pairwise_angles_deg:
max_angle = max(pairwise_angles_deg)
h, w = annotated_img.shape[:2]
text = f"angle: {max_angle:.1f} deg"
# 粗体效果（多次叠加）
for dx, dy in [(0, 0), (1, 0), (0, 1), (1, 1)]:
cv2.putText(
annotated_img,
text,
(w - 300 + dx, 40 + dy),
cv2.FONT_HERSHEY_SIMPLEX,
1.2,
(0, 0, 255),
3
)
# ---------- 保存 ----------
save_path = os.path.join(output_dir, "detected_" + img_filename)
cv2.imwrite(save_path, annotated_img)
print(f"✅ 保存完成: {save_path}")
# 保存每张图像结果
results_dict[img_filename] = {
"angles_deg": angles_deg,
"pairwise_angles_deg": pairwise_angles_deg
}
print("\n所有图像处理完成！")
print("\n🎉 全部处理完成")
return results_dict
# ------------------- 测试调用 -------------------
# ------------------- 主入口 -------------------
if __name__ == "__main__":
MODEL_PATH = r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb_new3/weights/best.pt'
IMAGE_SOURCE_DIR = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb/val"
MODEL_PATH = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb_new3/weights/best.pt"
IMAGE_SOURCE_DIR = r"/home/hx/yolo/angle_base_obb/test_image"
OUTPUT_DIR = "./inference_results"
results = process_obb_images(MODEL_PATH, IMAGE_SOURCE_DIR, OUTPUT_DIR)
results = process_obb_images(
MODEL_PATH,
IMAGE_SOURCE_DIR,
OUTPUT_DIR
)
for img_name, info in results.items():
print(f"\n {img_name}:")
print(f"主方向角度列表: {info['angles_deg']}")
print(f"两两夹角列表: {info['pairwise_angles_deg']}")
print(f"\n{img_name}")
print("主方向角:", info["angles_deg"])
print("两两夹角:", info["pairwise_angles_deg"])