CoTracker 环境配置与ORB 特征点提取结合实现视频特征点追踪

CoTracker 环境配置&与ORB 特征点提取结合实现视频特征点追踪

Step1：配置 CoTracker 环境

首先下载 conda，然后安装虚拟环境。

	conda craete -n cotracker python=3.8conda activate cotracker

然后根据官方提示从 Github 上面下载源码。
参考官方的提示，这个项目支持在 CPU 和 GPU 上运行，因此在配置环境时建议同时安装支持 CUDA 的 PyTorch 和 TorchVision。
官方链接的终端命令贴出来了，需要可自行粘帖。

	git clone https://github.com/facebookresearch/co-trackercd co-trackerpip install -e .pip install matplotlib flow_vis tqdm tensorboard

因为官方有已经训练好的权重文件，我们只需要下载下来就可以在 Demo 里面直接调用。命令也在此处。

	mkdir -p checkpointscd checkpointswget https://huggingface.co/facebook/cotracker/resolve/main/cotracker2.pthcd ..

当然，这个 CoTracker 在配置环境过程中肯定会有一些库的版本不对，因此需要重新卸载再安装一些库的版本。
以下是我的 cotracker 虚拟环境里面需要的库版本（只摘出来 Setup.py 文件里安装的，以及通过命令行安装的库）。大家可自行对照。

	matplotlib                    3.7.3flow-vis                      0.1opencv-python                 4.8.1.78torch                         2.1.1torchaudio                    2.1.1torchsummary                  1.5.1torchvision                   0.16.1tqdm                          4.66.1tensorboard#(没找到，不过并不影响 CoTracker 的使用)

Step2：运行官方的例程

官方有一份 demo.py 文件可以直接调用一些接口，方便进行视频的处理，但是为了更好的了解里面的一些借口的参数。建议可以参考项目里面的 demo.ipynb 文件，按照里面的步骤，自己重新写一个 demo 文件。

Step3：结合 ORB 特征点提取

为了下一步进行视频帧追踪预演，提前编写了一个针对连续图像读取并追踪的代码（注意：代码里面输入的不是一个视频，而是将一连串连续的图片转换成张量的数据格式传入了 GPU，所以虽然不是视频，但是效果差不多）。如下所示：

import os
import cv2
import torch
import argparse
import numpy as np
from base64 import b64encode
from PIL import Image
import matplotlib.pyplot as plt
from cotracker.utils.visualizer import Visualizer, read_video_from_path
from cotracker.predictor import CoTrackerPredictor
import torch.nn.functional as Fdef convert_images_to_tensor(image_folder):image_files = sorted(os.listdir(image_folder))  # 获取图片文件列表并排序first_path = os.path.join(image_folder, image_files[0])print(first_path)images = []n = 0for image_file in image_files:n += 1print(n)image_path = os.path.join(image_folder, image_file)image = cv2.imread(image_path)  # 使用OpenCV读取图片height, width, _ = image.shapeleft_half = image[:, :width//2, :]image = cv2.cvtColor(left_half, cv2.COLOR_BGR2RGB)  # 将图片从BGR颜色空间转换为RGBimage_tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()  # 转换为PyTorch张量images.append(image_tensor)video_tensor = torch.stack(images)video_tensor = video_tensor.permute(1, 0, 2, 3, 4)  # 转换成视频张量的形式shape = video_tensor.shapeprint(shape[0], shape[1], shape[2], shape[3], shape[4])return first_path, video_tensor# 特征点检测的参数
max_corners = 30
quality_level = 0.1
min_distance = 200def orb_track_points(first_image_path):raw_image = cv2.imread(first_image_path)height, width, _ = raw_image.shaperaw_left_image = raw_image[:, :width // 2, :]                # 只取左边部分corners = cv2.goodFeaturesToTrack(cv2.cvtColor(raw_left_image, cv2.COLOR_BGR2GRAY), max_corners, quality_level, min_distance)corners = np.int0(corners)queries = []# 將图像上检测到的特征点，添加到追踪里面for corner in corners:x, y = corner.ravel()# cv2.circle(raw_left_image, (x, y), 2, vector_color[i].tolist(), 2)coordinate = [0., float(x), float(y)]queries.append(coordinate)queries = torch.tensor(queries)print(queries)# 并将图像上选取的点变成张量输入if torch.cuda.is_available():queries = queries.cuda()# 创建了一个包含四个子图的2x2图像网格，用于可视化查询点的位置，将查询点的帧号提取出来，并转换为整数类型的列表 frame_numbers。帧号将用于在每个子图上显示对应的帧数frame_numbers = queries[:, 0].int().tolist()# plt.subplots()函数创建了一个图像网格, 并将返回的“轴”对象存储在变量axs中fig, axs = plt.subplots(1, 1)# 通过调用axs.set_title()设置子图的标题为"Frame {}axs.set_title("Frame {}".format(0))# 通过enumerate()函数同时迭代查询点(query)和对应的帧号(frame_number)for i, (query, frame_number) in enumerate(zip(queries, frame_numbers)):# 使用plot()函数在该子图上绘制一个红色的点，其坐标为(query[1].item(), query[2].item())axs.plot(query[1].item(), query[2].item(), 'ro')# 设置子图的x和y轴范围axs.set_xlim(0, video.shape[4])axs.set_ylim(0, video.shape[3])# 翻转y轴，以与视频的坐标系一致axs.invert_yaxis()# 调整子图之间的布局plt.tight_layout()plt.savefig('./saved_videos/image_grid.png')return queries# 指定图片文件夹路径
# images_folder = "./assets/1212/snapSave_p/Cam_2"      # Pitch 俯仰角
images_folder = "./assets/1212/snapSave_r/Cam_2"      # Roll  翻滚角
# images_folder = "./assets/1212/snapSave_y/Cam_2"      # Taw   偏航角# 调用函数将图片转换为张量
first_im_path, video = convert_images_to_tensor(images_folder)
image_queries = orb_track_points(first_im_path)model = CoTrackerPredictor(checkpoint=os.path.join('./checkpoints/cotracker_stride_4_wind_8.pth')
)if torch.cuda.is_available():model = model.cuda()video = video.cuda()# 前向
pred_tracks, pred_visibility = model(video, queries=image_queries[None])
print("数据计算完毕")
vis = Visualizer(save_dir='./saved_videos', linewidth=2, mode='cool', tracks_leave_trace=-1)# tracks_leave_trace = -1 可以显示出跟踪点的轨迹
vis.visualize(video=video, tracks=pred_tracks, visibility=pred_visibility, filename='orb_track')
print("视频存储完成")
# 原文里面有考虑对相机运动的补偿消除一些影响，但是代码里面这一部分设定为 False，即没有考虑相机运动的影响
# 因此 pred_tracks, pred_visibility 即跟踪真实值
track_save_data = './saved_videos/track_data'
if not os.path.exists(track_save_data):os.makedirs(track_save_data)for i in range(max_corners):format_i = "{:02d}".format(i)with open(track_save_data + '/save_data_' + str(format_i), 'w') as data_txt:for pred_track in pred_tracks[0]:point_track = str(pred_track[i][0].item()) + ' ' + \str(pred_track[i][1].item()) + '\n'data_txt.write(point_track)data_txt.close()print("数据文件关闭")

结果展示：

Step4：针对相机进行实时追踪，但失败

还是之前说的，因为 CoTracker 的神经网络本身在训练模型的时候就是以视频作为输入数据进行输入的，因此针对连续图片可以做到追踪，但时如果只是单个图片，那么追踪将无法进行。
下面可能就有小伙伴会想，通过缩小传入视频的帧率再输入。例如将 3 ～ 4 帧的图片作为一个短视频输入进去，然后计算出来结果后，将结果保存并用于下一个短视频的追踪，如此往复，实现相机实时追踪效果。
这个方向我也尝试过，但时 CoTracker 本身在进行视频的特征点计算的时候，就极其消耗算力。而且这个消耗的时间随着传入的视频时间以及要追踪的特征点数量线性增加。
我的设备是 RTX4060 和 i7-12650。性能还算可以。但是在传入一个连续 5 帧的视频，并追踪 10 个点的时候，依旧要花费 0.3 ～ 0.4 秒时间计算。出现的结构就是，视频一卡一卡的，实时跟踪效果很差。
（为什么传入 5 帧？ 因为 5 帧已经是网络输入要求的最低帧数了，再小就没有结果输出了。）

代码依旧贴在下面，其实就是在上面视频的基础上进行的改进：

import os
import cv2
import torch
import argparse
import numpy as np
from base64 import b64encode
from PIL import Image
import matplotlib.pyplot as plt
from cotracker.utils.visualizer import Visualizer, read_video_from_path
from cotracker.predictor import CoTrackerPredictor
import torch.nn.functional as F
import timedef mkdir():if not os.path.exists(saved_videos):os.makedirs(saved_videos)def initialize(first_image):n = 5i = 0images_pytorch = []image = cv2.cvtColor(first_image, cv2.COLOR_BGR2RGB)            # 将第一张图片从BGR颜色空间转换为RGBimage_tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()  # 转换为PyTorch张量images_pytorch.append(image_tensor)while i < 5:ret, current_image = cap.read()image = cv2.cvtColor(current_image, cv2.COLOR_BGR2RGB)  # 将第一张图片从BGR颜色空间转换为RGBimage_tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()  # 转换为PyTorch张量images_pytorch.append(image_tensor)i += 1# 將图片张量转换成网络输入视频张量的形式video_tensor = torch.stack(images_pytorch)video_tensor = video_tensor.permute(1, 0, 2, 3, 4)  # 转换成视频张量的形式print("video tensor------------------------------------------------------")print(video_tensor)return images_pytorch, video_tensor# 特征点检测的参数
max_corners = 5
quality_level = 0.1
min_distance = 100def orb_track_points(first_image_):# # 仿生眼相机图像前处理部分# raw_image = cv2.imread(first_image_path)# height, width, _ = raw_image.shape# raw_left_image = raw_image[:, :width // 2, :]                # 只取左边部分# corners = cv2.goodFeaturesToTrack(cv2.cvtColor(raw_left_image, cv2.COLOR_BGR2GRAY), max_corners, quality_level, min_distance)# corners = np.int0(corners)# # 电脑相机图像处理部分corners = cv2.goodFeaturesToTrack(cv2.cvtColor(first_image_, cv2.COLOR_BGR2GRAY), max_corners, quality_level, min_distance)corners = np.int0(corners)queries = []# 將图像上检测到的特征点，添加到追踪里面for corner in corners:x, y = corner.ravel()coordinate = [0., float(x), float(y)]queries.append(coordinate)queries = torch.tensor(queries)# 并将图像上选取的点变成张量输入if torch.cuda.is_available():queries = queries.cuda()return queriesdef convert_images_to_tensor(current_image, pre_images_pytorch):# # 將当前图像转换成 pytorch 张量，仿生眼相机图像预处理# height, width, _ = current_image.shape# left_half = current_image[:, :width // 2, :]# image = cv2.cvtColor(left_half, cv2.COLOR_BGR2RGB)  # 将图片从BGR颜色空间转换为RGB# 將当前图像转换成 pytorch 张量，电脑相机图像预处理image = cv2.cvtColor(current_image, cv2.COLOR_BGR2RGB)  # 将图片从BGR颜色空间转换为RGBimage_tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()  # 转换为PyTorch张量# 再將图片存入 pre_images 里面进行后续的跟踪计算pre_images_pytorch.append(image_tensor)update_images_pytorch = pre_images_pytorch[1:]# print("update_images_pytorch: %d", update_images_pytorch)# 將图片张量转换成网络输入视频张量的形式video_tensor = torch.stack(update_images_pytorch)video_tensor = video_tensor.permute(1, 0, 2, 3, 4)  # 转换成视频张量的形式return update_images_pytorch, video_tensorif __name__ == '__main__':saved_videos = "./assets/saved_videos/"mkdir()# 开启相机获取图像cap = cv2.VideoCapture(0)if not cap.isOpened():print("无法打开视频文件")exit()ret, first_frame = cap.read()if not ret:print("无法获取图像")exit()first_queries = orb_track_points(first_frame)first_images_pytorch, first_video = initialize(first_frame)print(first_queries)        # 分别是 0, x， y# 加载模型文件model = CoTrackerPredictor(checkpoint=os.path.join('./checkpoints/cotracker_stride_4_wind_8.pth'))print("模型创建完毕")# 將视频数据和模型数据转换if torch.cuda.is_available():model = model.cuda()first_video = first_video.cuda()# 前向first_tracks, first_visibility = model(first_video, queries=first_queries[None])        # 此处的 None 是用来增加维度的print("数据计算完毕")vis = Visualizer(save_dir=saved_videos, linewidth=2, mode='cool', tracks_leave_trace=-1)  # t_l_t:-1显示跟踪轨迹print('----------------------------------------------------------------------pre')print(first_tracks[0])vis.visualize(video=first_video, tracks=first_tracks, visibility=first_visibility, filename='orb_track')print("视频存储完成")images_pytorch = first_images_pytorch# 跟踪部分while True:ret, current_frame = cap.read()cv2.imshow("current", current_frame)cv2.waitKey(20)images_pytorch, current_video = convert_images_to_tensor(current_frame, images_pytorch)# 將视频数据和模型数据转换if torch.cuda.is_available():model = model.cuda()current_video = current_video.cuda()# 前向current_tracks, current_visibility = model(current_video, queries=first_queries[None])  # 此处的 None 是用来增加维度的print("----------------------------------------------------------")print(current_tracks[0][0])print("数据计算完毕")

5.内部代码的修改

原本代码里面为了显示跟踪的连续性，在可视化部分 ，将追踪点在不同时间段的轨迹连成了一条线。
我的项目里面之前为了结果的点的轨迹可以清楚一些，因此修改了原本可视化里面连线的部分，该成了画点。如下所示，里面注释掉的部分为曾经画线的代码，下面新增的为画点的代码

    def _draw_pred_tracks(self,rgb: np.ndarray,  # H x W x 3tracks: np.ndarray,  # T x 2vector_colors: np.ndarray,alpha: float = 0.5,):radius = 2  # 半径thickness = 2  # 线条宽度T, N, _ = tracks.shapefor s in range(T - 1):vector_color = vector_colors[s]original = rgb.copy()alpha = (s / T) ** 2for i in range(N):coord_x = (int(tracks[s, i, 0]), int(tracks[s, i, 1]))if coord_x[0] != 0 and coord_x[1] != 0:cv2.circle(rgb, coord_x, radius, vector_color[i].tolist(), thickness)   # 直接画出之前轨迹的点if self.tracks_leave_trace > 0:rgb = cv2.addWeighted(rgb, alpha, original, 1 - alpha, 0)#   遍历之前追踪的点集，然后连接相邻两点，画一条直线，构成轨迹图
#         for s in range(T - 1):
#             vector_color = vector_colors[s]
#             original = rgb.copy()
#             alpha = (s / T) ** 2
#             for i in range(N):
#                 coord_y = (int(tracks[s, i, 0]), int(tracks[s, i, 1]))
#                 coord_x = (int(tracks[s + 1, i, 0]), int(tracks[s + 1, i, 1]))
#                 if coord_y[0] != 0 and coord_y[1] != 0:
#                     cv2.line(
#                         rgb,
#                         coord_y,
#                         coord_x,
#                         vector_color[i].tolist(),
#                         self.linewidth,
#                         cv2.LINE_AA,
#                     )
#             if self.tracks_leave_trace > 0:
#                 rgb = cv2.addWeighted(rgb, alpha, original, 1 - alpha, 0)return rgb

当然不排除可能是本人技术太菜无法实现 CoTracker 的相机实时性追踪。如果后面有小伙伴实现了，欢迎在评论区里面分享。