CoTracker 環(huán)境配置&與ORB 特征點提取結合實現(xiàn)視頻特征點追蹤

Step1：配置 CoTracker 環(huán)境

首先下載 conda，然后安裝虛擬環(huán)境。

	conda craete -n cotracker python=3.8
	conda activate cotracker

然后根據(jù)官方提示從 Github 上面下載源碼。
參考官方的提示，這個項目支持在 CPU 和 GPU 上運行，因此在配置環(huán)境時建議同時安裝支持 CUDA 的 PyTorch 和 TorchVision。
官方鏈接的終端命令貼出來了，需要可自行粘帖。

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

因為官方有已經訓練好的權重文件，我們只需要下載下來就可以在 Demo 里面直接調用。命令也在此處。

	mkdir -p checkpoints
	cd checkpoints
	wget https://huggingface.co/facebook/cotracker/resolve/main/cotracker2.pth
	cd ..

當然，這個 CoTracker 在配置環(huán)境過程中肯定會有一些庫的版本不對，因此需要重新卸載再安裝一些庫的版本。
以下是我的 cotracker 虛擬環(huán)境里面需要的庫版本（只摘出來 Setup.py 文件里安裝的，以及通過命令行安裝的庫）。大家可自行對照。

	matplotlib                    3.7.3
	flow-vis                      0.1

	opencv-python                 4.8.1.78

	torch                         2.1.1
	torchaudio                    2.1.1
	torchsummary                  1.5.1
	torchvision                   0.16.1
	tqdm                          4.66.1
	tensorboard#(沒找到，不過并不影響 CoTracker 的使用)

Step2：運行官方的例程

官方有一份 demo.py 文件可以直接調用一些接口，方便進行視頻的處理，但是為了更好的了解里面的一些借口的參數(shù)。建議可以參考項目里面的 demo.ipynb 文件，按照里面的步驟，自己重新寫一個 demo 文件。

Step3：結合 ORB 特征點提取

為了下一步進行視頻幀追蹤預演，提前編寫了一個針對連續(xù)圖像讀取并追蹤的代碼（注意：代碼里面輸入的不是一個視頻，而是將一連串連續(xù)的圖片轉換成張量的數(shù)據(jù)格式傳入了 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 F


def 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 = 0
    for image_file in image_files:
        n += 1
        print(n)
        image_path = os.path.join(image_folder, image_file)
        image = cv2.imread(image_path)  # 使用OpenCV讀取圖片
        height, width, _ = image.shape
        left_half = image[:, :width//2, :]
        image = cv2.cvtColor(left_half, cv2.COLOR_BGR2RGB)  # 將圖片從BGR顏色空間轉換為RGB
        image_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.shape
    print(shape[0], shape[1], shape[2], shape[3], shape[4])
    return first_path, video_tensor


# 特征點檢測的參數(shù)
max_corners = 30
quality_level = 0.1
min_distance = 200


def orb_track_points(first_image_path):
    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)
    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()
    # 創(chuàng)建了一個包含四個子圖的2x2圖像網格，用于可視化查詢點的位置，將查詢點的幀號提取出來，并轉換為整數(shù)類型的列表 frame_numbers。幀號將用于在每個子圖上顯示對應的幀數(shù)
    frame_numbers = queries[:, 0].int().tolist()
    # plt.subplots()函數(shù)創(chuàng)建了一個圖像網格, 并將返回的“軸”對象存儲在變量axs中
    fig, axs = plt.subplots(1, 1)
    # 通過調用axs.set_title()設置子圖的標題為"Frame {}
    axs.set_title("Frame {}".format(0))
    # 通過enumerate()函數(shù)同時迭代查詢點(query)和對應的幀號(frame_number)
    for i, (query, frame_number) in enumerate(zip(queries, frame_numbers)):
        # 使用plot()函數(shù)在該子圖上繪制一個紅色的點，其坐標為(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   偏航角

# 調用函數(shù)將圖片轉換為張量
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("數(shù)據(jù)計算完畢")
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("數(shù)據(jù)文件關閉")

結果展示：

Step4：針對相機進行實時追蹤，但失敗

還是之前說的，因為 CoTracker 的神經網絡本身在訓練模型的時候就是以視頻作為輸入數(shù)據(jù)進行輸入的，因此針對連續(xù)圖片可以做到追蹤，但時如果只是單個圖片，那么追蹤將無法進行。
下面可能就有小伙伴會想，通過縮小傳入視頻的幀率再輸入。例如將 3 ～ 4 幀的圖片作為一個短視頻輸入進去，然后計算出來結果后，將結果保存并用于下一個短視頻的追蹤，如此往復，實現(xiàn)相機實時追蹤效果。
這個方向我也嘗試過，但時 CoTracker 本身在進行視頻的特征點計算的時候，就極其消耗算力。而且這個消耗的時間隨著傳入的視頻時間以及要追蹤的特征點數(shù)量線性增加。
我的設備是 RTX4060 和 i7-12650。性能還算可以。但是在傳入一個連續(xù) 5 幀的視頻，并追蹤 10 個點的時候，依舊要花費 0.3 ～ 0.4 秒時間計算。出現(xiàn)的結構就是，視頻一卡一卡的，實時跟蹤效果很差。
（為什么傳入 5 幀？ 因為 5 幀已經是網絡輸入要求的最低幀數(shù)了，再小就沒有結果輸出了。）

代碼依舊貼在下面，其實就是在上面視頻的基礎上進行的改進：

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 time


def mkdir():
    if not os.path.exists(saved_videos):
        os.makedirs(saved_videos)


def initialize(first_image):
    n = 5
    i = 0
    images_pytorch = []
    image = cv2.cvtColor(first_image, cv2.COLOR_BGR2RGB)            # 將第一張圖片從BGR顏色空間轉換為RGB
    image_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顏色空間轉換為RGB
        image_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


# 特征點檢測的參數(shù)
max_corners = 5
quality_level = 0.1
min_distance = 100


def 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 queries


def 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顏色空間轉換為RGB
    image_tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()  # 轉換為PyTorch張量
    # 再將圖片存入 pre_images 里面進行后續(xù)的跟蹤計算
    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_tensor


if __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("模型創(chuàng)建完畢")
    # 將視頻數(shù)據(jù)和模型數(shù)據(jù)轉換
    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("數(shù)據(jù)計算完畢")
    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)
        # 將視頻數(shù)據(jù)和模型數(shù)據(jù)轉換
        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("數(shù)據(jù)計算完畢")

5.內部代碼的修改

原本代碼里面為了顯示跟蹤的連續(xù)性，在可視化部分 ，將追蹤點在不同時間段的軌跡連成了一條線。
我的項目里面之前為了結果的點的軌跡可以清楚一些，因此修改了原本可視化里面連線的部分，該成了畫點。如下所示，里面注釋掉的部分為曾經畫線的代碼，下面新增的為畫點的代碼

    def _draw_pred_tracks(
        self,
        rgb: np.ndarray,  # H x W x 3
        tracks: np.ndarray,  # T x 2
        vector_colors: np.ndarray,
        alpha: float = 0.5,
    ):
        radius = 2  # 半徑
        thickness = 2  # 線條寬度
        T, N, _ = tracks.shape
        for s in range(T - 1):
            vector_color = vector_colors[s]
            original = rgb.copy()
            alpha = (s / T) ** 2
            for 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

當然不排除可能是本人技術太菜無法實現(xiàn) CoTracker 的相機實時性追蹤。如果后面有小伙伴實現(xiàn)了，歡迎在評論區(qū)里面分享。文章來源地址http://www.zghlxwxcb.cn/news/detail-779538.html

到了這里，關于CoTracker 環(huán)境配置&與ORB 特征點提取結合實現(xiàn)視頻特征點追蹤的文章就介紹完了。如果您還想了解更多內容，請在右上角搜索TOY模板網以前的文章或繼續(xù)瀏覽下面的相關文章，希望大家以后多多支持TOY模板網！