自学内容网 自学内容网

从视频帧生成点云数据、使用PointNet++模型提取特征,并将特征保存下来的完整实现。

文件地址

https://github.com/yanx27/Pointnet_Pointnet2_pytorch?spm=5176.28103460.0.0.21a95d27ollfze

Pointnet_Pointnet2_pytorch\log\classification\pointnet2_ssg_wo_normals文件夹改名为Pointnet_Pointnet2_pytorch\log\classification\pointnet2_cls_ssg

"E:\Pointnet_Pointnet2_pytorch\provider.py" 在provider.py文件头添加

def pc_normalize(pc):
    l = pc.shape[0]
    centroid = np.mean(pc, axis=0)
    pc = pc - centroid
    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
    pc = pc / m
    return pc

详细介绍了如何从视频帧中生成点云数据并使用PointNet++模型提取特征,最后将特征保存下来。
从视频帧中生成点云数据并提取特征
1. 引言

在计算机视觉领域,点云数据是一种重要的三维数据形式,广泛应用于自动驾驶、机器人导航、物体识别等场景。本文将详细介绍如何从视频帧中生成点云数据,并使用PointNet++模型提取特征,最后将特征保存下来以供后续分析或使用。
2. 环境准备

在开始之前,确保你的环境中安装了以下依赖项:

    Python 3.6+
    PyTorch 1.7+
    Open3D
    OpenCV
    NumPy

你可以使用以下命令安装这些依赖项:

pip install torch torchvision
pip install open3d opencv-python numpy

3. 代码实现

import os
import sys
# 获取当前脚本所在的目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 将 models 目录添加到 sys.path
sys.path.insert(0, os.path.join(current_dir, 'models'))
import numpy as np
import torch
import cv2
import open3d as o3d
from models.pointnet2_cls_ssg import get_model
from provider import pc_normalize
import time
import hashlib


# 打印 sys.path 以确认路径是否正确
print(sys.path)

# 设置设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"设备设置为: {device}")

# 加载预训练模型
#"E:\Pointnet_Pointnet2_pytorch\log\classification\pointnet2_cls_ssg\checkpoints\best_model.pth"
model_path = os.path.join('log', 'classification', 'pointnet2_cls_ssg', 'checkpoints', 'best_model.pth')  # 替换为实际路径
print(f"加载预训练模型: {model_path}")
model = get_model(num_class=40, normal_channel=False).to(device)

# 只加载模型参数
checkpoint = torch.load(model_path, map_location=device)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
print("模型加载完成")

# 从视频帧生成点云数据
def frame_to_point_cloud(frame, depth_frame=None):
    if depth_frame is None:
        # 如果没有深度图,使用随机生成的点云数据作为示例
        points = np.random.rand(1024, 3)  # 生成 1024 个点,每个点有 3 个坐标
    else:
        # 使用深度图生成点云
        h, w = depth_frame.shape
        y, x = np.indices((h, w))
        z = depth_frame
        points = np.stack([x, y, z], axis=-1).reshape(-1, 3)
        points = points[~np.isnan(points).any(axis=1)]  # 去除无效点
        if points.shape[0] > 1024:
            points = points[np.random.choice(points.shape[0], 1024, replace=False)]
        elif points.shape[0] < 1024:
            points = np.pad(points, ((0, 1024 - points.shape[0]), (0, 0)), mode='constant')
    print(f"生成点云数据: {points.shape}")
    return points

# 提取特征
def extract_features(model, point_cloud):
    point_cloud = pc_normalize(point_cloud)  # 归一化点云
    point_cloud = torch.from_numpy(point_cloud).float().unsqueeze(0).transpose(2, 1).to(device)
    with torch.no_grad():
        pred, trans_feat = model(point_cloud)  # 只接收两个返回值
    print(f"提取特征完成: {pred.shape}")
    return pred.cpu().numpy()

# 处理单个视频文件
def process_video(video_path, output_folder):
    cap = cv2.VideoCapture(video_path)
    frame_count = 0
    video_name = os.path.basename(video_path).split('.')[0]
    print(f"开始处理视频: {video_path}")
    all_features = []
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # 生成点云数据
        point_cloud = frame_to_point_cloud(frame)
        if point_cloud.shape[0] < 1024:  # 确保至少有 1024 个点
            print(f"帧 {frame_count} 点云数据不足,跳过")
            continue  # 如果点不够,跳过此帧

        # 提取特征
        features = extract_features(model, point_cloud)
        print(f"处理帧 {frame_count} 特征: {features}")

        # 保存特征到 all_features 列表中
        all_features.append(features)

        frame_count += 1

    cap.release()
    print(f"视频处理完成: {video_path}")

    # 生成唯一的文件名
    output_file = os.path.join(output_folder, f'{video_name}_features.npy')

    # 将所有特征保存到一个文件中
    np.save(output_file, np.vstack(all_features))
    print(f"特征已保存到: {output_file}")

# 检查视频文件是否已处理
def is_video_processed(video_path, processed_videos):
    video_hash = hashlib.md5(video_path.encode()).hexdigest()
    return video_hash in processed_videos

# 获取已处理的视频文件列表
def get_processed_videos(output_file):
    if not os.path.exists(output_file):
        return set()
    
    processed_videos = set()
    with open(output_file, 'r') as f:
        for line in f:
            processed_videos.add(line.strip())
    return processed_videos

# 记录已处理的视频文件
def record_processed_video(video_path, output_file):
    video_hash = hashlib.md5(video_path.encode()).hexdigest()
    with open(output_file, 'a') as f:
        f.write(video_hash + '\n')

# 处理视频文件夹
def process_video_folder(folder_path, output_folder):
    processed_videos_file = os.path.join(output_folder, 'processed_videos.txt')
    processed_videos = get_processed_videos(processed_videos_file)
    print(f"开始处理视频文件夹: {folder_path}")
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.endswith('.mp4') or file.endswith('.avi'):
                video_path = os.path.join(root, file)
                if is_video_processed(video_path, processed_videos):
                    print(f"视频已处理,跳过: {video_path}")
                    continue
                process_video(video_path, output_folder)
                record_processed_video(video_path, processed_videos_file)

    print("所有视频处理完成")

# 主程序
if __name__ == "__main__":
    # 视频文件夹路径
    input_folder = r'E:\Pointnet_Pointnet2_pytorch\data\voide'
    output_folder = r'E:\Pointnet_Pointnet2_pytorch\data\voide_features'  # 特征保存路径

    # 确保输出文件夹存在
    os.makedirs(output_folder, exist_ok=True)

    # 处理视频文件夹
    process_video_folder(input_folder, output_folder)


3.1 导入必要的库

首先,我们需要导入一些必要的库,包括文件操作、数值计算、深度学习框架、图像处理和点云处理相关的库。
3.2 设置设备

检查是否有可用的GPU,并设置设备。如果存在GPU,将使用GPU进行计算;否则,使用CPU。
3.3 加载预训练模型

我们使用PointNet++模型来提取点云特征。首先,加载预训练模型。这通常涉及以下几个步骤:

    指定模型路径:提供预训练模型的路径。
    加载模型:使用 get_model 函数创建模型实例,并将其移动到指定的设备(CPU或GPU)。
    加载模型参数:从预训练模型文件中加载模型参数,并设置模型为评估模式。

3.4 从视频帧生成点云数据

定义一个函数 frame_to_point_cloud,该函数从视频帧中生成点云数据。如果没有深度图,可以生成随机点云数据作为示例。具体步骤如下:

    生成随机点云:如果没有深度图,生成1024个随机点,每个点有3个坐标。
    使用深度图生成点云:如果有深度图,从深度图中提取点云数据。具体做法是将深度图的每个像素位置(x, y)和对应的深度值z组合成一个三维点(x, y, z)。然后,去除无效点,并确保点云数据的形状为 (1024, 3)。

3.5 提取特征

定义一个函数 extract_features,该函数使用预训练模型提取点云数据的特征。具体步骤如下:

    归一化点云:对点云数据进行归一化处理,使其适合输入到模型中。
    转换为张量:将点云数据转换为PyTorch张量,并移动到指定的设备。
    提取特征:使用预训练模型提取特征,并返回特征向量。

3.6 处理单个视频文件

定义一个函数 process_video,该函数处理单个视频文件,逐帧生成点云数据并提取特征。具体步骤如下:

    打开视频文件:使用OpenCV的 cv2.VideoCapture 打开视频文件。
    读取帧:逐帧读取视频。
    生成点云数据:调用 frame_to_point_cloud 函数生成点云数据。
    提取特征:调用 extract_features 函数提取特征。
    保存特征:将提取的特征保存为 .npy 文件。

3.7 处理视频文件夹

定义一个函数 process_video_folder,该函数处理指定文件夹中的所有视频文件。具体步骤如下:

    遍历文件夹:使用 os.walk 遍历指定文件夹中的所有视频文件。
    处理每个视频:调用 process_video 函数处理每个视频文件。

3.8 主程序

在主程序中,指定输入视频文件夹和输出特征文件夹的路径,并调用 process_video_folder 函数处理所有视频文件。
4. 总结

本文详细介绍了如何从视频帧中生成点云数据,并使用PointNet++模型提取特征,最后将特征保存下来。通过这些步骤,你可以将视频数据转换为点云数据,并提取有用的特征,为后续的分析和应用提供支持。


原文地址:https://blog.csdn.net/sinat_38424109/article/details/143860753

免责声明:本站文章内容转载自网络资源,如本站内容侵犯了原著者的合法权益,可联系本站删除。更多内容请关注自学内容网(zxcms.com)!