Files
zjsh_yolov11/output_masks/mask_rotation.py

64 lines
1.9 KiB
Python
Raw Normal View History

2025-08-13 12:53:33 +08:00
import cv2
import numpy as np
import os
# 输入:已保存的黑白掩码图像路径
mask_path = 'composite_mask.png'
# 输出:可视化结果保存路径
output_dir = 'output_masks'
os.makedirs(output_dir, exist_ok=True)
result_image_path = os.path.join(output_dir, 'mask_with_rotated_rects.jpg')
# 读取掩码图像(灰度图)
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
if mask is None:
raise FileNotFoundError(f"❌ 无法读取掩码图像: {mask_path}")
# 二值化(确保是 0 和 255
_, binary = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
# 查找轮廓
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# 创建三通道图像用于可视化(白色背景,黑色轮廓)
vis_image = np.ones((*mask.shape, 3), dtype=np.uint8) * 255 # 白底
# 存储所有旋转矩形结果
rotated_rects = []
for i, cnt in enumerate(contours):
if len(cnt) < 2:
continue # 至少需要两个点才能拟合
# 计算最小外接旋
rect = cv2.minAreaRect(cnt) # (center, (width, height), angle)
rotated_rects.append(rect)
# 获取四个角
box = cv2.boxPoints(rect)
box = np.float16(box)
# 在可视化图像上绘制旋转矩形(绿色)
cv2.drawContours(vis_image, [box], 0, (0, 180, 0), 2)
# 可选:标注序号
cx, cy = int(rect[0][0]), int(rect[0][1])
cv2.putText(vis_image, str(i), (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)
# 打印信息
center = rect[0]
size = rect[1]
angle = rect[2]
print(f"物体 {i}:")
print(f" 中心点: ({center[0]:.1f}, {center[1]:.1f})")
print(f" 尺寸: 宽={size[0]:.1f}, 高={size[1]:.1f}")
print(f" 旋转角度: {angle:.2f}°")
print(f" 四个角点: {box.tolist()}")
print("-" * 40)
# 保存可视化结果
cv2.imwrite(result_image_path, vis_image)
print(f"✅ 旋转矩形可视化已保存至: {result_image_path}")