[2D Vision] 2D Point Tracking: co-tracker 사용법

💡 본 문서는 '[2D Vision] 2D Point Tracking: co-tracker 사용법'에 대해 정리해놓은 글입니다.
간단하게 사용하기 좋은 Point Tracking 모듈인 co-tracker를 사용하는 방법에 대해 정리하였으니 참고하시기 바랍니다.

1. co-tracker 사용법:  Quick Start

github 에서 단순하게 언급한 co-tracker 사용방법입니다. 여기 보이듯이 torch.hub.load를 통해 checkpoint를 불러올 경우 따로 co-tracker를 git clone 해서 사용할 필요없이 cotracker 함수를 통해 불러와서 사용할 수 있습니다. 

import torch
# Download the video
url = 'https://github.com/facebookresearch/co-tracker/raw/refs/heads/main/assets/apple.mp4'

import imageio.v3 as iio
frames = iio.imread(url, plugin="FFMPEG")  # plugin="pyav"

device = 'cuda'
grid_size = 10
video = torch.tensor(frames).permute(0, 3, 1, 2)[None].float().to(device)  # B T C H W

# Run Offline CoTracker:
cotracker = torch.hub.load("facebookresearch/co-tracker", "cotracker3_offline").to(device)
pred_tracks, pred_visibility = cotracker(video, grid_size=grid_size) # B T N 2,  B T N 1

추가로 결과를 시각화 하기 위해서는 co-tracker/cotracker/utils/visualizer.py의 파일을 복사해서 사용하고 싶은 곳에 visualizer.py로 붙여넣기 한 후에 결과를 시각화해서 확인할 수 있습니다.

from visualizer import Visualizer
    visualizer = Visualizer(
        save_dir=output_cotracker,
        fps=10,
        mode="rainbow",         # rainbow or optical_flow
        tracks_leave_trace=15,  # tracks leave trace
    )
    res_video = visualizer.visualize(
        video=video*255,
        tracks=pred_tracks,
        visibility=pred_visibility if 'visibility_tensor' in locals() else None,
        filename="pred_tracks_visualization",
        save_video=True
    )

위의 모든 과정을 시각화해놓은 파일은 co-tracker/demo.py 위치에 있으니 참고하시기 바랍니다.

+ Utils

추가적으로 현재의 코드 내에서는 mp4 파일을 읽어다가 사용하는 코드만 존재합니다. 따라서 visualize하는 위치인 co-tracker/cotracker/utils/visualizer.py에 아래의 함수를 추가하면 sequence image를 읽어 이를 바탕으로 tracker할 수 있습니다.

 

def read_image_sequence_from_path(images_path):
    # Supported image file extensions
    img_extensions = ["*.jpg", "*.jpeg", "*.png", "*.bmp"]
    img_files = []
    
    # Collect image files with supported extensions
    for ext in img_extensions:
        img_files.extend(glob.glob(os.path.join(images_path, ext)))

    # Return None if no images are found
    if not img_files:
        print(f"No images found in path: {images_path}")
        return None

    # Sort files alphabetically to ensure proper sequence order
    img_files = sorted(img_files)

    frames = []
    for img_path in img_files:
        # Open each image and convert to RGB format
        img = Image.open(img_path).convert("RGB")
        frames.append(np.array(img))

    # Stack all images into a single numpy array (T, H, W, C)
    return np.stack(frames)

이 함수를 사용하기 위해서는 아래와 같이 demo.py 파일을 수정하시면 됩니다.

    args = parser.parse_args()
    # load the input video frame by frame
    if args.video_path.endswith(".mp4"):
        video = read_video_from_path(args.video_path)
    else:
        video = read_image_sequence_from_path(args.video_path)

2. co-tracker 사용법:  Module 변경 후 활용

TODO...

참고

[Github] co-tracker: https://github.com/facebookresearch/co-tracker