diff --git a/ailai_obb/angle.py b/ailai_obb/angle.py new file mode 100644 index 0000000..1ff66c4 --- /dev/null +++ b/ailai_obb/angle.py @@ -0,0 +1,105 @@ +from ultralytics import YOLO +import cv2 +import os +import numpy as np + +def get_best_obb_angle(image_path, weight_path, return_degree=False): + """ + 输入: + image_path: 图像路径 + weight_path: YOLO权重路径 + return_degree: 是否返回角度单位为度,默认 False(返回弧度) + 输出: + 置信度最高目标的旋转角 + 如果未检测到目标返回 None + """ + # 读取图像 + img = cv2.imread(image_path) + if img is None: + print(f"❌ 无法读取图像:{image_path}") + return None + + # 加载模型并预测 + model = YOLO(weight_path) + results = model(img, save=False, imgsz=640, conf=0.15, mode='obb') + result = results[0] + + boxes = result.obb + if not boxes: + print("⚠️ 未检测到目标。") + return None + + # 取置信度最高框的旋转角 + best_box = max(boxes, key=lambda x: x.conf.cpu().numpy()[0]) + r = best_box.xywhr.cpu().numpy()[0][4] # 弧度 + + if return_degree: + return np.degrees(r) + else: + return r + + +def save_obb_visual(image_path, weight_path, save_path): + """ + 输入: + image_path: 图像路径 + weight_path: YOLO权重路径 + save_path: 保存带角度标注图像路径 + 功能: + 检测 OBB 并标注置信度最高框旋转角度,保存图片 + """ + img = cv2.imread(image_path) + if img is None: + print(f"❌ 无法读取图像:{image_path}") + return + + model = YOLO(weight_path) + results = model(img, save=False, imgsz=640, conf=0.15, mode='obb') + result = results[0] + + boxes = result.obb + if not boxes: + print("⚠️ 未检测到目标。") + return + + best_box = max(boxes, key=lambda x: x.conf.cpu().numpy()[0]) + cx, cy, w, h, r = best_box.xywhr.cpu().numpy()[0] + angle_deg = np.degrees(r) + + # 绘制 OBB + annotated_img = img.copy() + rect = ((cx, cy), (w, h), angle_deg) + box_pts = cv2.boxPoints(rect).astype(int) + cv2.polylines(annotated_img, [box_pts], isClosed=True, color=(0, 255, 0), thickness=2) + + # 标注角度 + text = f"{angle_deg:.1f}°" + font_scale = max(0.5, min(w, h)/100) + thickness = 2 + text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness) + text_x = int(cx - text_size[0]/2) + text_y = int(cy + text_size[1]/2) + cv2.putText(annotated_img, text, (text_x, text_y), + cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 255), thickness) + + # 保存 + os.makedirs(os.path.dirname(save_path), exist_ok=True) + cv2.imwrite(save_path, annotated_img) + print(f"✅ 检测结果已保存至: {save_path}") + + +# =============================== +# 示例调用 +# =============================== +if __name__ == "__main__": + weight = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb3/weights/best.pt" + image = r"/home/hx/yolo/output_masks/2.jpg" + save_path = "./inference_results/best_detected_2.jpg" + + angle_rad = get_best_obb_angle(image, weight) + print(f"旋转角(弧度):{angle_rad:.4f}") + + angle_deg = get_best_obb_angle(image, weight, return_degree=True) + print(f"旋转角(度):{angle_deg:.2f}°") + + save_obb_visual(image, weight, save_path) diff --git a/ailai_obb/bag_bushu.py b/ailai_obb/bag_bushu.py new file mode 100644 index 0000000..a57228d --- /dev/null +++ b/ailai_obb/bag_bushu.py @@ -0,0 +1,171 @@ +import cv2 +import numpy as np +import math +from shapely.geometry import Polygon +from rknnlite.api import RKNNLite +import os + +CLASSES = ['clamp'] +nmsThresh = 0.4 +objectThresh = 0.5 + +# ------------------- 工具函数 ------------------- +def letterbox_resize(image, size, bg_color=114): + target_width, target_height = size + image_height, image_width, _ = image.shape + scale = min(target_width / image_width, target_height / image_height) + new_width, new_height = int(image_width * scale), int(image_height * scale) + image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA) + canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color + offset_x, offset_y = (target_width - new_width) // 2, (target_height - new_height) // 2 + canvas[offset_y:offset_y + new_height, offset_x:offset_x + new_width] = image + return canvas, scale, offset_x, offset_y + +class DetectBox: + def __init__(self, classId, score, xmin, ymin, xmax, ymax, angle): + self.classId = classId + self.score = score + self.xmin = xmin + self.ymin = ymin + self.xmax = xmax + self.ymax = ymax + self.angle = angle + +def rotate_rectangle(x1, y1, x2, y2, a): + cx, cy = (x1 + x2) / 2, (y1 + y2) / 2 + x1_new = int((x1 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx) + y1_new = int((x1 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy) + x2_new = int((x2 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx) + y2_new = int((x2 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy) + x3_new = int((x1 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx) + y3_new = int((x1 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy) + x4_new = int((x2 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx) + y4_new = int((x2 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy) + return [(x1_new, y1_new), (x3_new, y3_new), (x2_new, y2_new), (x4_new, y4_new)] + +def intersection(g, p): + g = Polygon(np.array(g).reshape(-1,2)) + p = Polygon(np.array(p).reshape(-1,2)) + if not g.is_valid or not p.is_valid: + return 0 + inter = g.intersection(p).area + union = g.area + p.area - inter + return 0 if union == 0 else inter / union + +def NMS(detectResult): + predBoxs = [] + sort_detectboxs = sorted(detectResult, key=lambda x: x.score, reverse=True) + for i in range(len(sort_detectboxs)): + if sort_detectboxs[i].classId == -1: + continue + p1 = rotate_rectangle(sort_detectboxs[i].xmin, sort_detectboxs[i].ymin, + sort_detectboxs[i].xmax, sort_detectboxs[i].ymax, + sort_detectboxs[i].angle) + predBoxs.append(sort_detectboxs[i]) + for j in range(i + 1, len(sort_detectboxs)): + if sort_detectboxs[j].classId == sort_detectboxs[i].classId: + p2 = rotate_rectangle(sort_detectboxs[j].xmin, sort_detectboxs[j].ymin, + sort_detectboxs[j].xmax, sort_detectboxs[j].ymax, + sort_detectboxs[j].angle) + if intersection(p1, p2) > nmsThresh: + sort_detectboxs[j].classId = -1 + return predBoxs + +def sigmoid(x): + return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x))) + +def softmax(x, axis=-1): + exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) + return exp_x / np.sum(exp_x, axis=axis, keepdims=True) + +def process(out, model_w, model_h, stride, angle_feature, index, scale_w=1, scale_h=1): + class_num = len(CLASSES) + angle_feature = angle_feature.reshape(-1) + xywh = out[:, :64, :] + conf = sigmoid(out[:, 64:, :]).reshape(-1) + boxes = [] + for ik in range(model_h * model_w * class_num): + if conf[ik] > objectThresh: + w = ik % model_w + h = (ik % (model_w * model_h)) // model_w + c = ik // (model_w * model_h) + # 解析xywh + xywh_ = xywh[0, :, (h * model_w) + w].reshape(1, 4, 16, 1) + data = np.arange(16).reshape(1, 1, 16, 1) + xywh_ = softmax(xywh_, 2) + xywh_ = np.sum(xywh_ * data, axis=2).reshape(-1) + xywh_add = xywh_[:2] + xywh_[2:] + xywh_sub = (xywh_[2:] - xywh_[:2]) / 2 + # 安全取角度 + angle_idx = min(index + (h * model_w) + w, len(angle_feature) - 1) + angle = (angle_feature[angle_idx] - 0.25) * math.pi + cos_a, sin_a = math.cos(angle), math.sin(angle) + xy = xywh_sub[0] * cos_a - xywh_sub[1] * sin_a, xywh_sub[0] * sin_a + xywh_sub[1] * cos_a + xywh1 = np.array([xy[0] + w + 0.5, xy[1] + h + 0.5, xywh_add[0], xywh_add[1]]) + xywh1 *= stride + xmin = (xywh1[0] - xywh1[2]/2) * scale_w + ymin = (xywh1[1] - xywh1[3]/2) * scale_h + xmax = (xywh1[0] + xywh1[2]/2) * scale_w + ymax = (xywh1[1] + xywh1[3]/2) * scale_h + boxes.append(DetectBox(c, conf[ik], xmin, ymin, xmax, ymax, angle)) + return boxes + +# ------------------- 主函数 ------------------- +def detect_boxes_angle_rknn(model_path, image_path, save_path=None): + img = cv2.imread(image_path) + if img is None: + print(f"❌ 无法读取图像: {image_path}") + return None, None + + img_resized, scale, offset_x, offset_y = letterbox_resize(img, (640, 640)) + infer_img = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB) + infer_img = np.expand_dims(infer_img, 0) + + rknn_lite = RKNNLite(verbose=False) + rknn_lite.load_rknn(model_path) + rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0) + + results = rknn_lite.inference([infer_img]) + detect_boxes = [] + for x in results[:-1]: + index, stride = 0, 0 + if x.shape[2] == 20: + stride, index = 32, 20*4*20*4 + 20*2*20*2 + elif x.shape[2] == 40: + stride, index = 16, 20*4*20*4 + elif x.shape[2] == 80: + stride, index = 8, 0 + feature = x.reshape(1, 65, -1) + detect_boxes += process(feature, x.shape[3], x.shape[2], stride, results[-1], index) + + detect_boxes = NMS(detect_boxes) + + # 输出每个检测框角度 + for i, box in enumerate(detect_boxes): + print(f"框 {i+1}: angle = {box.angle:.4f} rad ({np.degrees(box.angle):.2f}°)") + if save_path: + xmin = int((box.xmin - offset_x)/scale) + ymin = int((box.ymin - offset_y)/scale) + xmax = int((box.xmax - offset_x)/scale) + ymax = int((box.ymax - offset_y)/scale) + points = rotate_rectangle(xmin, ymin, xmax, ymax, box.angle) + cv2.polylines(img, [np.array(points, np.int32)], True, (0, 255, 0), 1) + cv2.putText(img, f"{np.degrees(box.angle):.1f}°", (xmin, ymin-5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1) + + if save_path: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + cv2.imwrite(save_path, img) + print(f"✅ 带角度的检测结果已保存到 {save_path}") + + rknn_lite.release() + return detect_boxes, img + +# ------------------- 使用示例 ------------------- +if __name__ == "__main__": + model_path = "obb.rknn" + image_path = "2.jpg" + save_path = "./inference_results/boxes_with_angle.jpg" + os.makedirs(os.path.dirname(save_path), exist_ok=True) + detect_boxes_angle_rknn(model_path, image_path, save_path) + diff --git a/ailai_obb/bushu_angle.py b/ailai_obb/bushu_angle.py new file mode 100644 index 0000000..891ee62 --- /dev/null +++ b/ailai_obb/bushu_angle.py @@ -0,0 +1,197 @@ + +import cv2 +import numpy as np +import math +from shapely.geometry import Polygon +from rknnlite.api import RKNNLite +import os + +# ------------------- 配置 ------------------- +CLASSES = ['clamp'] +nmsThresh = 0.4 +objectThresh = 0.5 + +# ------------------- 全局原图尺寸 ------------------- +ORIG_W = 2560 # 原图宽 +ORIG_H = 1440 # 原图高 + +# ------------------- 工具函数 ------------------- +def letterbox_resize(image, size, bg_color=114): + target_width, target_height = size + image_height, image_width, _ = image.shape + scale = min(target_width / image_width, target_height / image_height) + new_width, new_height = int(image_width * scale), int(image_height * scale) + image_resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA) + canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color + offset_x, offset_y = (target_width - new_width) // 2, (target_height - new_height) // 2 + canvas[offset_y:offset_y + new_height, offset_x:offset_x + new_width] = image_resized + return canvas, scale, offset_x, offset_y + +class DetectBox: + def __init__(self, classId, score, xmin, ymin, xmax, ymax, angle): + self.classId = classId + self.score = score + self.xmin = xmin + self.ymin = ymin + self.xmax = xmax + self.ymax = ymax + self.angle = angle + +def rotate_rectangle(x1, y1, x2, y2, a): + cx, cy = (x1 + x2) / 2, (y1 + y2) / 2 + x1_new = int((x1 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx) + y1_new = int((x1 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy) + x2_new = int((x2 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx) + y2_new = int((x2 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy) + x3_new = int((x1 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx) + y3_new = int((x1 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy) + x4_new = int((x2 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx) + y4_new = int((x2 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy) + return [(x1_new, y1_new), (x3_new, y3_new), (x2_new, y2_new), (x4_new, y4_new)] + +def intersection(g, p): + g = Polygon(np.array(g).reshape(-1,2)) + p = Polygon(np.array(p).reshape(-1,2)) + if not g.is_valid or not p.is_valid: + return 0 + inter = g.intersection(p).area + union = g.area + p.area - inter + return 0 if union == 0 else inter / union + +def NMS(detectResult): + predBoxs = [] + sort_detectboxs = sorted(detectResult, key=lambda x: x.score, reverse=True) + for i in range(len(sort_detectboxs)): + if sort_detectboxs[i].classId == -1: + continue + p1 = rotate_rectangle(sort_detectboxs[i].xmin, sort_detectboxs[i].ymin, + sort_detectboxs[i].xmax, sort_detectboxs[i].ymax, + sort_detectboxs[i].angle) + predBoxs.append(sort_detectboxs[i]) + for j in range(i + 1, len(sort_detectboxs)): + if sort_detectboxs[j].classId == sort_detectboxs[i].classId: + p2 = rotate_rectangle(sort_detectboxs[j].xmin, sort_detectboxs[j].ymin, + sort_detectboxs[j].xmax, sort_detectboxs[j].ymax, + sort_detectboxs[j].angle) + if intersection(p1, p2) > nmsThresh: + sort_detectboxs[j].classId = -1 + return predBoxs + +def sigmoid(x): + return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x))) + +def softmax(x, axis=-1): + exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) + return exp_x / np.sum(exp_x, axis=axis, keepdims=True) + +# ------------------- 关键修改:process函数加入scale ------------------- +def process(out, model_w, model_h, stride, angle_feature, index, scale=1.0, offset_x=0, offset_y=0): + class_num = len(CLASSES) + angle_feature = angle_feature.reshape(-1) + xywh = out[:, :64, :] + conf = sigmoid(out[:, 64:, :]).reshape(-1) + boxes = [] + for ik in range(model_h * model_w * class_num): + if conf[ik] > objectThresh: + w = ik % model_w + h = (ik % (model_w * model_h)) // model_w + c = ik // (model_w * model_h) + # 解析xywh + xywh_ = xywh[0, :, (h * model_w) + w].reshape(1, 4, 16, 1) + xywh_ = softmax(xywh_, 2) + data = np.arange(16).reshape(1, 1, 16, 1) + xywh_ = np.sum(xywh_ * data, axis=2).reshape(-1) + xywh_add = xywh_[:2] + xywh_[2:] + xywh_sub = (xywh_[2:] - xywh_[:2]) / 2 + # 取角度 + angle_idx = min(index + (h * model_w) + w, len(angle_feature) - 1) + angle = (angle_feature[angle_idx] - 0.25) * math.pi + cos_a, sin_a = math.cos(angle), math.sin(angle) + xy = xywh_sub[0] * cos_a - xywh_sub[1] * sin_a, xywh_sub[0] * sin_a + xywh_sub[1] * cos_a + xywh1 = np.array([xy[0] + w + 0.5, xy[1] + h + 0.5, xywh_add[0], xywh_add[1]]) + xywh1 *= stride + # 映射回原图坐标 + xmin = (xywh1[0] - xywh1[2]/2 - offset_x) / scale + ymin = (xywh1[1] - xywh1[3]/2 - offset_y) / scale + xmax = (xywh1[0] + xywh1[2]/2 - offset_x) / scale + ymax = (xywh1[1] + xywh1[3]/2 - offset_y) / scale + boxes.append(DetectBox(c, conf[ik], xmin, ymin, xmax, ymax, angle)) + return boxes + +# ------------------- 新可调用函数 ------------------- +def detect_boxes_rknn(model_path, image_path): + img = cv2.imread(image_path) + if img is None: + print(f"❌ 无法读取图像: {image_path}") + return None, None + + img_resized, scale, offset_x, offset_y = letterbox_resize(img, (640, 640)) + infer_img = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB) + infer_img = np.expand_dims(infer_img, 0) + + rknn_lite = RKNNLite(verbose=False) + rknn_lite.load_rknn(model_path) + rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0) + + results = rknn_lite.inference([infer_img]) + detect_boxes = [] + for x in results[:-1]: + index, stride = 0, 0 + if x.shape[2] == 20: + stride, index = 32, 20*4*20*4 + 20*2*20*2 + elif x.shape[2] == 40: + stride, index = 16, 20*4*20*4 + elif x.shape[2] == 80: + stride, index = 8, 0 + feature = x.reshape(1, 65, -1) + detect_boxes += process(feature, x.shape[3], x.shape[2], stride, results[-1], index, + scale=scale, offset_x=offset_x, offset_y=offset_y) + + detect_boxes = NMS(detect_boxes) + rknn_lite.release() + return detect_boxes, img + +# ------------------- 绘制与辅助函数 ------------------- +def get_angles(detect_boxes): + return [box.angle for box in detect_boxes] + +def draw_boxes(img, detect_boxes, save_path=None): + for box in detect_boxes: + points = rotate_rectangle(box.xmin, box.ymin, box.xmax, box.ymax, box.angle) + cv2.polylines(img, [np.array(points, np.int32)], True, (0, 255, 0), 1) + cv2.putText(img, f"{np.degrees(box.angle):.1f}°", (int(box.xmin), int(box.ymin)-5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1) + if save_path: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + cv2.imwrite(save_path, img) + print(f"✅ 带角度的检测结果已保存到 {save_path}") + return img + +def visualize_top_box(img, detect_boxes, save_path=None): + if not detect_boxes: + return img + top_box = max(detect_boxes, key=lambda x: x.score) + points = rotate_rectangle(top_box.xmin, top_box.ymin, top_box.xmax, top_box.ymax, top_box.angle) + cv2.polylines(img, [np.array(points, np.int32)], True, (0, 255, 0), 2) + cv2.putText(img, f"{np.degrees(top_box.angle):.1f}°", (int(top_box.xmin), int(top_box.ymin)-5), + cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,0,255), 2) + if save_path: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + cv2.imwrite(save_path, img) + return img + +# ------------------- 使用示例 ------------------- +if __name__ == "__main__": + model_path = "obb.rknn" + image_path = "2.jpg" + + detect_boxes, img = detect_boxes_rknn(model_path, image_path) + angles = get_angles(detect_boxes) + for i, angle in enumerate(angles): + print(f"框 {i+1}: angle = {angle:.4f} rad ({np.degrees(angle):.2f}°)") + + save_path_all = "./inference_results/boxes_all.jpg" + draw_boxes(img.copy(), detect_boxes, save_path_all) + + save_path_top = "./inference_results/top_box.jpg" + visualize_top_box(img.copy(), detect_boxes, save_path_top) diff --git a/angle_base_obb/anger_caculate.py b/angle_base_obb/anger_caculate.py deleted file mode 100644 index 1aebd67..0000000 --- a/angle_base_obb/anger_caculate.py +++ /dev/null @@ -1,85 +0,0 @@ -from ultralytics import YOLO -import cv2 -import numpy as np -import os - -# 1. 加载模型 -model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb2/weights/best.pt') - -# 2. 读取图像 -img_path = r"/home/hx/yolo/output_masks/2.jpg" -img = cv2.imread(img_path) - -if img is None: - print(f"❌ 错误:无法读取图像!请检查路径:{img_path}") - exit(1) - -# 3. 预测(OBB 模式) -results = model( - img, - save=False, - imgsz=640, - conf=0.15, - mode='obb' -) - -# 4. 获取结果并绘制 -result = results[0] -annotated_img = result.plot() - -# 5. 保存结果 -output_dir = "./inference_results" -os.makedirs(output_dir, exist_ok=True) -filename = os.path.basename(img_path) -save_path = os.path.join(output_dir, "detected_" + filename) -cv2.imwrite(save_path, annotated_img) -print(f"✅ 推理结果已保存至: {save_path}") - -# 6. 提取旋转框并计算 **两个框之间的夹角** -boxes = result.obb -if boxes is None or len(boxes) == 0: - print("❌ No objects detected.") -else: - print(f"✅ Detected {len(boxes)} object(s):") - directions = [] # 存储每个框的主方向(弧度),归一化到 [0, π) - - for i, box in enumerate(boxes): - cls = int(box.cls.cpu().numpy()[0]) - conf = box.conf.cpu().numpy()[0] - xywhr = box.xywhr.cpu().numpy()[0] # [cx, cy, w, h, r] - cx, cy, w, h, r_rad = xywhr - - # 确定主方向(长边方向) - if w >= h: - direction = r_rad # 长边方向就是 r - else: - direction = r_rad + np.pi / 2 # 长边方向是 r + 90° - - # 归一化到 [0, π) - direction = direction % np.pi - - directions.append(direction) - angle_deg = np.degrees(direction) - print(f" Box {i+1}: Class: {cls}, Confidence: {conf:.3f}, 主方向: {angle_deg:.2f}°") - - # ✅ 计算任意两个框之间的夹角(最小夹角,0° ~ 90°) - if len(directions) >= 2: - print("\n🔍 计算两个旋转框之间的夹角(主方向夹角):") - for i in range(len(directions)): - for j in range(i + 1, len(directions)): - dir1 = directions[i] - dir2 = directions[j] - - # 计算方向差(取最小夹角,考虑周期性) - diff = abs(dir1 - dir2) - diff = min(diff, np.pi - diff) # 最小夹角(0 ~ π/2) - diff_deg = np.degrees(diff) - - print(f" Box {i+1} 与 Box {j+1} 之间的夹角: {diff_deg:.2f}°") - else: - print("⚠️ 检测到少于两个目标,无法计算夹角。") - -# 7. 显示图像 -cv2.imshow("YOLO OBB Prediction", annotated_img) -cv2.waitKey(0) -cv2.destroyAllWindows() \ No newline at end of file diff --git a/angle_base_obb/anger_caculate_file.py b/angle_base_obb/anger_caculate_file.py deleted file mode 100644 index d3a712a..0000000 --- a/angle_base_obb/anger_caculate_file.py +++ /dev/null @@ -1,103 +0,0 @@ -from ultralytics import YOLO -import cv2 -import numpy as np -import os - -# ================== 配置参数 ================== -MODEL_PATH = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb4/weights/best.pt" -IMAGE_SOURCE_DIR = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb2/test" # 图像文件夹路径 -OUTPUT_DIR = "./inference_results" # 输出结果保存路径 -IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'} - -# 创建输出目录 -os.makedirs(OUTPUT_DIR, exist_ok=True) - -# 1. 加载模型 -print("🔄 加载 YOLO 模型...") -model = YOLO(MODEL_PATH) -print("✅ 模型加载完成") - -# 获取所有图像文件 -image_files = [ - f for f in os.listdir(IMAGE_SOURCE_DIR) - if os.path.splitext(f.lower())[1] in IMG_EXTENSIONS -] - -if not image_files: - print(f"❌ 错误:在路径中未找到图像文件:{IMAGE_SOURCE_DIR}") - exit(1) - -print(f"📁 发现 {len(image_files)} 张图像待处理") - -# ================== 批量处理每张图像 ================== -for img_filename in image_files: - img_path = os.path.join(IMAGE_SOURCE_DIR, img_filename) - print(f"\n🖼️ 正在处理:{img_filename}") - - # 读取图像 - img = cv2.imread(img_path) - if img is None: - print(f"❌ 跳过:无法读取图像 {img_path}") - continue - - # 推理(OBB 模式) - results = model( - img, - save=False, - imgsz=640, - conf=0.15, - 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 - directions = [] # 存储每个框的主方向(弧度),归一化到 [0, π) - - if boxes is None or len(boxes) == 0: - print("❌ 该图像中未检测到任何目标") - else: - print(f"✅ 检测到 {len(boxes)} 个目标:") - for i, box in enumerate(boxes): - cls = int(box.cls.cpu().numpy()[0]) - conf = box.conf.cpu().numpy()[0] - xywhr = box.xywhr.cpu().numpy()[0] # [cx, cy, w, h, r] - cx, cy, w, h, r_rad = xywhr - - # 确定主方向(长边方向) - if w >= h: - direction = r_rad # 长边方向 - else: - direction = r_rad + np.pi / 2 # 长边是宽的方向 - - # 归一化到 [0, π) - direction = direction % np.pi - directions.append(direction) - - angle_deg = np.degrees(direction) - print(f" Box {i+1}: Class: {cls}, Confidence: {conf:.3f}, 主方向: {angle_deg:.2f}°") - - # 计算两两之间的夹角(最小夹角,0°~90°) - if len(directions) >= 2: - print("\n🔍 计算各框之间的夹角(主方向最小夹角):") - for i in range(len(directions)): - for j in range(i + 1, len(directions)): - dir1 = directions[i] - dir2 = directions[j] - - diff = abs(dir1 - dir2) - min_diff_rad = min(diff, np.pi - diff) # 最小夹角(考虑周期性) - min_diff_deg = np.degrees(min_diff_rad) - - print(f" Box {i+1} 与 Box {j+1} 之间夹角: {min_diff_deg:.2f}°") - else: - print("⚠️ 检测到少于两个目标,无法计算夹角。") - -print("\n🎉 所有图像处理完成!") \ No newline at end of file diff --git a/angle_base_obb/angle_caculate.py b/angle_base_obb/angle_caculate.py new file mode 100644 index 0000000..dc7d9aa --- /dev/null +++ b/angle_base_obb/angle_caculate.py @@ -0,0 +1,76 @@ +import cv2 +import os +import numpy as np +from ultralytics import YOLO + +def predict_obb_best_angle(model_path, image_path, save_path=None): + """ + 输入: + model_path: YOLO 权重路径 + image_path: 图片路径 + save_path: 可选,保存带标注图像 + 输出: + angle_deg: 置信度最高两个框的主方向夹角(度),如果检测少于两个目标返回 None + annotated_img: 可视化图像 + """ + # 1. 加载模型 + model = YOLO(model_path) + + # 2. 读取图像 + img = cv2.imread(image_path) + if img is None: + print(f"无法读取图像: {image_path}") + return None, None + + # 3. 推理 OBB + results = model(img, save=False, imgsz=640, conf=0.5, mode='obb') + result = results[0] + + # 4. 可视化 + annotated_img = result.plot() + if save_path: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + cv2.imwrite(save_path, annotated_img) + print(f"推理结果已保存至: {save_path}") + + # 5. 提取旋转角度和置信度 + boxes = result.obb + if boxes is None or len(boxes) < 2: + print("检测到少于两个目标,无法计算夹角。") + return None, annotated_img + + box_info = [] + for box in boxes: + 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 + box_info.append((conf, direction)) + + # 6. 取置信度最高两个框 + box_info = sorted(box_info, key=lambda x: x[0], reverse=True) + dir1, dir2 = box_info[0][1], box_info[1][1] + + # 7. 计算夹角(最小夹角,0~90°) + diff = abs(dir1 - dir2) + diff = min(diff, np.pi - diff) + angle_deg = np.degrees(diff) + + print(f"置信度最高两个框主方向夹角: {angle_deg:.2f}°") + return angle_deg, annotated_img + + +# ------------------- 测试 ------------------- +if __name__ == "__main__": + weight_path = r'best.pt' + image_path = r"./test_image/3.jpg" + save_path = "./inference_results/detected_3.jpg" + + #angle_deg, annotated_img = predict_obb_best_angle(weight_path, image_path, save_path) + angle_deg,_ = predict_obb_best_angle(weight_path, image_path, save_path) + annotated_img = None + print(angle_deg) + if annotated_img is not None: + cv2.imshow("YOLO OBB Prediction", annotated_img) + cv2.waitKey(0) + cv2.destroyAllWindows() diff --git a/angle_base_obb/angle_caculate_file.py b/angle_base_obb/angle_caculate_file.py new file mode 100644 index 0000000..6ba93ec --- /dev/null +++ b/angle_base_obb/angle_caculate_file.py @@ -0,0 +1,102 @@ +import cv2 +import os +import numpy as np +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 推理,计算每张图像检测目标的主方向和夹角。 + + 输入: + model_path: YOLO 权重路径 + image_dir: 图像文件夹路径 + output_dir: 输出结果保存路径 + conf_thresh: 置信度阈值 + imgsz: 输入图像大小 + 输出: + results_dict: {image_filename: {'angles_deg': [...], 'pairwise_angles_deg': [...]}} + """ + os.makedirs(output_dir, exist_ok=True) + results_dict = {} + + print("加载 YOLO 模型...") + 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] + if not image_files: + print(f"❌ 未找到图像文件:{image_dir}") + return results_dict + + print(f"发现 {len(image_files)} 张图像待处理") + + for img_filename in image_files: + img_path = os.path.join(image_dir, img_filename) + print(f"\n正在处理:{img_filename}") + + img = cv2.imread(img_path) + if img is None: + print(f"❌ 跳过:无法读取图像 {img_path}") + continue + + # 推理 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 = [] + if boxes is None or len(boxes) == 0: + 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}°") + + # 两两夹角 + 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])) + 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}°") + + # 保存每张图像结果 + results_dict[img_filename] = { + "angles_deg": angles_deg, + "pairwise_angles_deg": pairwise_angles_deg + } + + print("\n所有图像处理完成!") + return results_dict + + +# ------------------- 测试调用 ------------------- +if __name__ == "__main__": + MODEL_PATH = r'best.pt' + IMAGE_SOURCE_DIR = r"./test_image" + OUTPUT_DIR = "./inference_results" + + 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']}") diff --git a/angle_base_obb/best.pt b/angle_base_obb/best.pt new file mode 100644 index 0000000..7b5b32d Binary files /dev/null and b/angle_base_obb/best.pt differ diff --git a/angle_base_obb/test_image/1.jpg b/angle_base_obb/test_image/1.jpg new file mode 100644 index 0000000..2882cd5 Binary files /dev/null and b/angle_base_obb/test_image/1.jpg differ diff --git a/angle_base_obb/test_image/2.jpg b/angle_base_obb/test_image/2.jpg new file mode 100644 index 0000000..dcd5369 Binary files /dev/null and b/angle_base_obb/test_image/2.jpg differ diff --git a/angle_base_obb/test_image/3.jpg b/angle_base_obb/test_image/3.jpg new file mode 100644 index 0000000..8df7feb Binary files /dev/null and b/angle_base_obb/test_image/3.jpg differ diff --git a/yemian/resize/rtest.py b/yemian/resize/rtest.py index e656978..afca979 100644 --- a/yemian/resize/rtest.py +++ b/yemian/resize/rtest.py @@ -7,7 +7,7 @@ from ultralytics import YOLO from pathlib import Path # ====================== 配置参数 ====================== -MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/seg_r/exp/weights/best.pt" +MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/seg_r/exp2/weights/best.pt" #SOURCE_IMG_DIR = "/home/hx/yolo/output_masks" # 原始输入图像目录 SOURCE_IMG_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/f6" # 原始输入图像目录 OUTPUT_DIR = "/home/hx/yolo/output_masks2" # 推理输出根目录 diff --git a/zhuangtai_class_cls/remain_tuili.py b/zhuangtai_class_cls/remain_tuili.py new file mode 100644 index 0000000..d23840e --- /dev/null +++ b/zhuangtai_class_cls/remain_tuili.py @@ -0,0 +1,143 @@ +import os +from pathlib import Path +import cv2 +import numpy as np +from ultralytics import YOLO + +# --------------------------- +# 类别映射 +# --------------------------- +CLASS_NAMES = { + 0: "未堆料", + 1: "小堆料", + 2: "大堆料", + 3: "未浇筑满", + 4: "浇筑满" +} + +# --------------------------- +# 加载 ROI 列表 +# --------------------------- +def load_global_rois(txt_path): + rois = [] + if not os.path.exists(txt_path): + print(f"❌ ROI 文件不存在: {txt_path}") + return rois + with open(txt_path, 'r') as f: + for line in f: + s = line.strip() + if s: + try: + x, y, w, h = map(int, s.split(',')) + rois.append((x, y, w, h)) + except Exception as e: + print(f"⚠️ 无法解析 ROI 行 '{s}': {e}") + return rois + +# --------------------------- +# 裁剪并 resize ROI +# --------------------------- +def crop_and_resize(img, rois, target_size=640): + crops = [] + h_img, w_img = img.shape[:2] + for i, (x, y, w, h) in enumerate(rois): + if x < 0 or y < 0 or x + w > w_img or y + h > h_img: + continue + roi = img[y:y+h, x:x+w] + roi_resized = cv2.resize(roi, (target_size, target_size), interpolation=cv2.INTER_AREA) + crops.append((roi_resized, i)) + return crops + +# --------------------------- +# class1/class2 加权判断 +# --------------------------- +def weighted_small_large(pred_probs, threshold=0.4, w1=0.3, w2=0.7): + p1 = float(pred_probs[1]) + p2 = float(pred_probs[2]) + total = p1 + p2 + if total > 0: + score = (w1 * p1 + w2 * p2) / total + else: + score = 0.0 + final_class = "大堆料" if score >= threshold else "小堆料" + return final_class, score, p1, p2 + +# --------------------------- +# 单张图片推理函数 +# --------------------------- +def classify_image_weighted(image, model, threshold=0.5): + results = model(image) + pred_probs = results[0].probs.data.cpu().numpy().flatten() + class_id = int(pred_probs.argmax()) + confidence = float(pred_probs[class_id]) + class_name = CLASS_NAMES.get(class_id, f"未知类别({class_id})") + + # class1/class2 使用加权得分 + if class_id in [1, 2]: + final_class, score, p1, p2 = weighted_small_large(pred_probs, threshold=threshold) + else: + final_class = class_name + score = confidence + p1 = float(pred_probs[1]) + p2 = float(pred_probs[2]) + + return final_class, score, p1, p2 + +# --------------------------- +# 批量推理主函数 +# --------------------------- +def batch_classify_images(model_path, input_folder, output_root, roi_file, target_size=640, threshold=0.5): + # 加载模型 + model = YOLO(model_path) + + # 确保输出根目录存在 + output_root = Path(output_root) + output_root.mkdir(parents=True, exist_ok=True) + + # 为所有类别创建目录 + class_dirs = {} + for name in CLASS_NAMES.values(): + d = output_root / name + d.mkdir(exist_ok=True) + class_dirs[name] = d + + rois = load_global_rois(roi_file) + if not rois: + print("❌ 没有有效 ROI,退出") + return + + # 遍历图片 + for img_path in Path(input_folder).glob("*.*"): + if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tif']: + continue + try: + img = cv2.imread(str(img_path)) + if img is None: + continue + + crops = crop_and_resize(img, rois, target_size) + + for roi_resized, roi_idx in crops: + final_class, score, p1, p2 = classify_image_weighted(roi_resized, model, threshold=threshold) + + # 文件名中保存 ROI、类别、加权分数、class1/class2 置信度 + suffix = f"_roi{roi_idx}_{final_class}_score{score:.2f}_p1{p1:.2f}_p2{p2:.2f}" + dst_path = class_dirs[final_class] / f"{img_path.stem}{suffix}{img_path.suffix}" + cv2.imwrite(dst_path, roi_resized) + print(f"{img_path.name}{suffix} -> {final_class} (score={score:.2f}, p1={p1:.2f}, p2={p2:.2f})") + + except Exception as e: + print(f"⚠️ 处理失败 {img_path.name}: {e}") + +# --------------------------- +# 使用示例 +# --------------------------- +if __name__ == "__main__": + model_path = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/cls_resize/exp_cls2/weights/best.pt" + input_folder = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/f6" + output_root = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/classified" + roi_file = "./roi_coordinates/1_rois.txt" + target_size = 640 + threshold = 0.4 # 可调节的比例系数 + + batch_classify_images(model_path, input_folder, output_root, roi_file, target_size, threshold)