diff --git a/demo1.py b/demo1.py new file mode 100644 index 0000000..b313217 --- /dev/null +++ b/demo1.py @@ -0,0 +1,47 @@ +import os +import torch + +from base64 import b64encode +from cotracker.utils.visualizer import Visualizer, read_video_from_path + + +video = read_video_from_path('./assets/F1_shorts.mp4') +video = torch.from_numpy(video).permute(0, 3, 1, 2)[None].float() + +from cotracker.predictor import CoTrackerPredictor + +model = CoTrackerPredictor( + checkpoint=os.path.join( + './checkpoints/cotracker2.pth' + ) +) + +# pred_tracks, pred_visibility = model(video, grid_size=30) + +# vis = Visualizer(save_dir='./videos', pad_value=100) +# vis.visualize(video=video, tracks=pred_tracks, visibility=pred_visibility, filename='teaser'); + +grid_query_frame=20 + +import torch.nn.functional as F +video_interp = F.interpolate(video[0], [200, 360], mode="bilinear")[None] + +import time +start_time = time.time() +pred_tracks, pred_visibility = model(video_interp, + grid_query_frame=grid_query_frame, backward_tracking=True) +end_time = time.time() + +print("Time taken: {:.2f} seconds".format(end_time - start_time)) + +vis = Visualizer( + save_dir='./videos', + pad_value=20, + linewidth=1, + mode='optical_flow' +) +vis.visualize( + video=video, + tracks=pred_tracks, + visibility=pred_visibility, + filename='dense'); \ No newline at end of file