yolov8训练模型、测试视频

yolov8先训练生成best.pt文件,用这个生成的模型进行视频的测试

因为本来用的代码生成的测试视频打不开,格式应该是损坏了,或者部分帧没有正常保存吧。

修改了一下代码,现状可以正常打开生成的视频了。

1、训练代码train.py

python 复制代码
import os

# os.environ["CUDA_VISIBLE_DEVICES"] = "3"  # 同样是选择第3块GPU

from ultralytics import YOLO

# Load a model
# model = YOLO("yolov8n.yaml")  # build a new model from YAML
# model = YOLO("yolov8n.pt")  # load a pretrained model (recommended for training)

# ffs = os.listdir("cfg1116/new_cfg")
# for ff in ffs:
model = YOLO(f"cfg1116/yolov8n.yaml")  # build from YAML and transfer weights
# Train the model
# results = model.train(data=r"/mnt/disk3/sunjiahui/CV-code/v8_all/data.yaml", epochs=5, imgsz=1280, workers=0, batch=2, device=[2])
results = model.train(
    data=r"/mnt/disk3/sunjiahui/CV-code/v8_all/data.yaml",
    epochs=500,
    imgsz=1280,
    workers=0,
    batch=2,
    device=[0],
    hsv_h=0.015,  # HSV色调变化
    hsv_s=0.7,    # HSV饱和度变化
    hsv_v=0.4,    # HSV亮度变化
    degrees=0.0,  # 旋转角度
    translate=0.1,  # 平移比例
    scale=0.5,    # 缩放比例
    shear=0.0,    # 剪切变换
    perspective=0.0,  # 透视变换
    flipud=0.0,   # 上下翻转概率
    fliplr=0.5,   # 左右翻转概率
    mosaic=1.0,   # Mosaic增强的概率
    mixup=0.0     # MixUp增强的概率
)
model.val(imgsz=[1280,1280])

2、测试代码:视频

python 复制代码
from ultralytics import YOLO
import cv2
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "2"  # 同样是选择第3块GPU

def process_video():
    # 初始化模型
    model = YOLO("runs/detect/train2/weights/best.pt")
    
    # 输入输出路径
    input_path = "/mnt/disk3/sunjiahui/CV-code/v8_all/XIONG_AN/shipin.mp4"
    output_path = "/mnt/disk3/sunjiahui/CV-code/v8_all/XIONG_AN/output_video15.mp4"
    
    # 尝试不同编解码器组合
    codec_options = ['mp4v', 'avc1', 'X264', 'MJPG']
    success = False
    
    for codec in codec_options:
        try:
            cap = cv2.VideoCapture(input_path)
            fps = int(cap.get(cv2.CAP_PROP_FPS)) or 30  # 处理fps为0的情况
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            
            fourcc = cv2.VideoWriter_fourcc(*codec)
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
            print(f"尝试使用编解码器 {codec}...")
            
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                
                results = model.predict(frame, conf=0.15)
                annotated_frame = results[0].plot()
                
                # 确保帧格式正确
                if annotated_frame.shape[:2] != (height, width):
                    annotated_frame = cv2.resize(annotated_frame, (width, height))
                
                out.write(annotated_frame)
            
            success = True
            break
            
        except Exception as e:
            print(f"编解码器 {codec} 失败: {str(e)}")
            if os.path.exists(output_path):
                os.remove(output_path)
            continue
            
        finally:
            cap.release()
            out.release()
    
    if success:
        print(f"视频生成成功!保存路径:{os.path.abspath(output_path)}")
        print("如果仍无法播放,请尝试以下方案:")
        print("1. 使用 VLC 播放器(兼容性最佳)")
        print("2. 执行命令:ffmpeg -i output_video.mp4 -c:v libx264 final.mp4")
    else:
        print("所有编解码器尝试失败,改用图像序列方案...")
        save_as_image_sequence(model, input_path)

def save_as_image_sequence(model, input_path):
    """备用方案:保存为图片序列"""
    output_dir = "video_frames"
    os.makedirs(output_dir, exist_ok=True)
    
    cap = cv2.VideoCapture(input_path)
    frame_count = 0
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        results = model.predict(frame)
        annotated_frame = results[0].plot()
        
        cv2.imwrite(f"{output_dir}/frame_{frame_count:04d}.jpg", annotated_frame)
        frame_count += 1
    
    cap.release()
    print(f"图像序列已保存至 {output_dir},可用以下命令合成视频:")
    print(f"ffmpeg -framerate 30 -i {output_dir}/frame_%04d.jpg -c:v libx264 output.mp4")

if __name__ == "__main__":
    process_video()
相关推荐
LiJieNiub9 小时前
YOLO-V1 与 YOLO-V2 核心技术解析:目标检测的迭代突破
人工智能·yolo·目标检测
小小测试开发9 小时前
pytest 库用法示例:Python 测试框架的高效实践
开发语言·python·pytest
至善迎风9 小时前
把 Python 应用打包成 Mac 应用程序 — 完整指南
python·macos
OpenBayes9 小时前
教程上新|重新定义下一代 OCR:IBM 最新开源 Granite-docling-258M,实现端到端的「结构+内容」统一理解
人工智能·深度学习·机器学习·自然语言处理·ocr·图像识别·文档处理
C嘎嘎嵌入式开发10 小时前
【机器学习算法篇】K-近邻算法
算法·机器学习·近邻算法
应用市场10 小时前
无人机编队飞行原理与Python仿真实现完整指南
python·无人机·cocos2d
蓝桉~MLGT10 小时前
Python学习历程——字符串相关操作及正则表达式
python·学习·正则表达式
一晌小贪欢10 小时前
Python爬虫第5课:正则表达式与数据清洗技术
爬虫·python·正则表达式·网络爬虫·python爬虫·python3·网页爬虫
ARM+FPGA+AI工业主板定制专家11 小时前
【JETSON+FPGA+GMSL】实测分享 | 如何实现激光雷达与摄像头高精度时间同步?
人工智能·数码相机·机器学习·fpga开发·机器人·自动驾驶
Nina_71711 小时前
Google提示词白皮书总结(2)
人工智能·python