diff --git a/algorithm/berthing/weights/best.pt b/algorithm/berthing/weights/best.pt index c113b0f..7e501c7 100644 Binary files a/algorithm/berthing/weights/best.pt and b/algorithm/berthing/weights/best.pt differ diff --git a/algorithm/berthing/weights/last.pt b/algorithm/berthing/weights/last.pt index 6380c5f..5e887f8 100644 Binary files a/algorithm/berthing/weights/last.pt and b/algorithm/berthing/weights/last.pt differ diff --git a/algorithm/board_ship/weights/best.pt b/algorithm/board_ship/weights/best.pt index 5fdb6b6..4e4f0d0 100644 Binary files a/algorithm/board_ship/weights/best.pt and b/algorithm/board_ship/weights/best.pt differ diff --git a/algorithm/board_ship/weights/last.pt b/algorithm/board_ship/weights/last.pt index 2becbe7..54f2543 100644 Binary files a/algorithm/board_ship/weights/last.pt and b/algorithm/board_ship/weights/last.pt differ diff --git a/algorithm/check_seg.py b/algorithm/check_seg.py new file mode 100644 index 0000000..5e5c154 --- /dev/null +++ b/algorithm/check_seg.py @@ -0,0 +1,169 @@ +import cv2 +import supervision as sv +from ultralytics import YOLO +import os +import time +import numpy as np +from tqdm import tqdm +from PIL import Image, ImageDraw, ImageFont + +def put_chinese_text(img, text, position, font_size=20, color=(0, 0, 255)): + # 将OpenCV图像转换为PIL图像 + img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) + draw = ImageDraw.Draw(img_pil) + + # 加载支持中文的字体 + try: + # 尝试加载系统中的微软雅黑字体 + font = ImageFont.truetype("msyh.ttc", font_size) + except: + # 如果找不到指定字体，使用默认字体 + font = ImageFont.load_default() + + # 绘制文本 + draw.text(position, text, font=font, fill=color) + + # 将PIL图像转换回OpenCV图像 + return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR) + +def process_video_with_segmentation( + input_video_path: str, + output_video_path: str, + model_name: str = "yolov8n-seg.pt", # 使用分割模型 + confidence_threshold: float = 0.5, + classes: list = None, + show_live: bool = True, + save_annotated: bool = False +): + """ + 参数: + input_video_path: 输入视频文件路径 + output_video_path: 输出视频文件路径 + model_name: YOLO分割模型名称或路径 (默认: yolov8n-seg.pt) + confidence_threshold: 检测置信度阈值 (0-1) + classes: 要检测的类别ID列表 (None表示所有类别) + show_live: 是否实时显示处理过程 + save_annotated: 是否保存标注后的视频 + """ + # 1. 初始化分割模型 + print(f"加载分割模型: {model_name}") + model = YOLO(model_name) + + # 2. 初始化视频读取器 + print(f"打开视频文件: {input_video_path}") + video_info = sv.VideoInfo.from_video_path(input_video_path) + cap = cv2.VideoCapture(input_video_path) + + if not cap.isOpened(): + print(f"无法打开视频文件: {input_video_path}") + return + + # 3. 初始化视频写入器 (如果需要保存结果) + if save_annotated: + output_dir = os.path.dirname(output_video_path) + if output_dir and not os.path.exists(output_dir): + os.makedirs(output_dir) + print(f"创建输出目录: {output_dir}") + + writer = cv2.VideoWriter( + output_video_path, + cv2.VideoWriter_fourcc(*'mp4v'), + video_info.fps, + (video_info.width, video_info.height) + ) + print(f"准备保存结果到: {output_video_path}") + + # 4. 初始化Supervision工具 - 使用MaskAnnotator替代BoxAnnotator + mask_annotator = sv.MaskAnnotator( + color=sv.Color(r=0, g=255, b=0), # 绿色掩码 + opacity=0.5 # 50%透明度 + ) + + # 5. 处理进度跟踪 + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + print(f"视频总帧数: {total_frames}") + pbar = tqdm(total=total_frames, desc="处理视频帧") + + + + # 7. 处理每一帧 + frame_count = 0 + processing_times = [] + + while cap.isOpened(): + start_time = time.time() + + # 读取帧 + ret, frame = cap.read() + if not ret: + break + + frame_count += 1 + + # 使用YOLO进行图像分割 + results = model( + frame, + conf=confidence_threshold, + classes=classes, + verbose=False + )[0] + + # 转换为Supervision检测对象 + detections = sv.Detections.from_ultralytics(results) + + # 应用分割掩码 + segmented_frame = mask_annotator.annotate( + scene=frame.copy(), + detections=detections + ) + + + # 计算处理时间 + end_time = time.time() + processing_time = end_time - start_time + processing_times.append(processing_time) + + # 实时显示处理结果 + if show_live: + cv2.imshow("Segmentation", segmented_frame) + if cv2.waitKey(1) == 27: # ESC键退出 + break + + # 保存处理后的帧 + if save_annotated: + writer.write(segmented_frame) + + # 更新进度条 + pbar.update(1) + + # 8. 清理资源 + cap.release() + if save_annotated: + writer.release() + if show_live: + cv2.destroyAllWindows() + pbar.close() + + # 打印性能统计 + if processing_times: + avg_time = sum(processing_times) / len(processing_times) + fps = 1 / avg_time if avg_time > 0 else 0 + print(f"\n处理完成! 平均处理时间: {avg_time:.4f}秒/帧, 约 {fps:.2f} FPS") + +if __name__ == "__main__": + # 示例用法 + input_video = r"D:\yolo8\新建文件夹\靠泊视频.mp4" # 替换为您的视频路径 + output_video = r"output/segmented_video.mp4" # 输出视频路径 + + # 选择要检测的类别 (0=人, 2=车等) + selected_classes = [0] # 只检测人 + + process_video_with_segmentation( + input_video_path=input_video, + output_video_path=output_video, + model_name=r"D:\yolo8\model\kaobo\weights\best.pt", # 使用分割模型 + confidence_threshold=0.3, # 置信度阈值 + classes=selected_classes, # 指定检测类别 + show_live=True, # 实时显示处理 + save_annotated=True # 保存结果视频 + ) \ No newline at end of file diff --git a/algorithm/leave_ship/weights/best.pt b/algorithm/leave_ship/weights/best.pt index 2b49426..92dd9db 100644 Binary files a/algorithm/leave_ship/weights/best.pt and b/algorithm/leave_ship/weights/best.pt differ diff --git a/algorithm/leave_ship/weights/last.pt b/algorithm/leave_ship/weights/last.pt index 973c8cd..c9101a5 100644 Binary files a/algorithm/leave_ship/weights/last.pt and b/algorithm/leave_ship/weights/last.pt differ diff --git a/algorithm/unberthing/weights/best.pt b/algorithm/unberthing/weights/best.pt index 5acea75..9de1082 100644 Binary files a/algorithm/unberthing/weights/best.pt and b/algorithm/unberthing/weights/best.pt differ diff --git a/algorithm/unberthing/weights/last.pt b/algorithm/unberthing/weights/last.pt index 8a42838..5a30746 100644 Binary files a/algorithm/unberthing/weights/last.pt and b/algorithm/unberthing/weights/last.pt differ